From 4413beeb6416fee66147ff80141c47bcf6bac8b4 Mon Sep 17 00:00:00 2001
From: Vincent Geddes <117534+vgeddes@users.noreply.github.com>
Date: Thu, 29 Aug 2024 18:00:09 +0200
Subject: [PATCH] Add heartbeat for gsrpc connection (#1277)

* Add heartbeat option for relayer connection

* whitespace
---
 relayer/chain/parachain/connection.go | 28 +++++++++++++++++++++++++++
 relayer/relays/beacon/main.go         |  3 ++-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/relayer/chain/parachain/connection.go b/relayer/chain/parachain/connection.go
index 1bb25f1838..38fe1a50ac 100644
--- a/relayer/chain/parachain/connection.go
+++ b/relayer/chain/parachain/connection.go
@@ -5,6 +5,7 @@ package parachain
 
 import (
 	"context"
+	"time"
 
 	"github.com/sirupsen/logrus"
 
@@ -72,6 +73,33 @@ func (co *Connection) Connect(_ context.Context) error {
 	return nil
 }
 
+func (co *Connection) ConnectWithHeartBeat(ctx context.Context, heartBeat time.Duration) error {
+	err := co.Connect(ctx)
+	if err != nil {
+		return err
+	}
+
+	ticker := time.NewTicker(heartBeat)
+
+	go func() {
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				_, err := co.API().RPC.System.Version()
+				if err != nil {
+					log.WithField("endpoint", co.endpoint).Error("Connection heartbeat failed")
+					return
+				}
+			}
+		}
+	}()
+
+	return nil
+}
+
 func (co *Connection) Close() {
 	// TODO: Fix design issue in GSRPC preventing on-demand closing of connections
 }
diff --git a/relayer/relays/beacon/main.go b/relayer/relays/beacon/main.go
index 2fc42223fc..58b46d8688 100644
--- a/relayer/relays/beacon/main.go
+++ b/relayer/relays/beacon/main.go
@@ -2,6 +2,7 @@ package beacon
 
 import (
 	"context"
+	"time"
 
 	"github.com/snowfork/snowbridge/relayer/chain/parachain"
 	"github.com/snowfork/snowbridge/relayer/crypto/sr25519"
@@ -36,7 +37,7 @@ func (r *Relay) Start(ctx context.Context, eg *errgroup.Group) error {
 
 	paraconn := parachain.NewConnection(r.config.Sink.Parachain.Endpoint, r.keypair.AsKeyringPair())
 
-	err := paraconn.Connect(ctx)
+	err := paraconn.ConnectWithHeartBeat(ctx, 60 * time.Second)
 	if err != nil {
 		return err
 	}