From 31c98b129cae0bab706b11efcdffd0055ab65ef4 Mon Sep 17 00:00:00 2001
From: Phuc Ta <30641530+vinhphuctadang@users.noreply.github.com>
Date: Thu, 27 Apr 2023 20:32:35 +0700
Subject: [PATCH] fix: add cancel context to stop blockHeight fetcher (#123)

* fix: add cancel context to stop blockHeight fetcher

* fix: check non-nil cancelFn on chainClient Close()
---
 client/chain/chain.go | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/client/chain/chain.go b/client/chain/chain.go
index 3dbf1314..52b709e0 100644
--- a/client/chain/chain.go
+++ b/client/chain/chain.go
@@ -143,6 +143,9 @@ type chainClient struct {
 	msgC        chan sdk.Msg
 	syncMux     *sync.Mutex
 
+	cancelCtx context.Context
+	cancelFn  func()
+
 	accNum    uint64
 	accSeq    uint64
 	gasWanted uint64
@@ -208,6 +211,7 @@ func NewChainClient(
 		}
 	}
 
+	cancelCtx, cancelFn := context.WithCancel(context.Background())
 	// build client
 	cc := &chainClient{
 		ctx:  ctx,
@@ -224,6 +228,8 @@ func NewChainClient(
 		syncMux:   new(sync.Mutex),
 		msgC:      make(chan sdk.Msg, msgCommitBatchSizeLimit),
 		doneC:     make(chan bool, 1),
+		cancelCtx: cancelCtx,
+		cancelFn:  cancelFn,
 
 		sessionEnabled: stickySessionEnabled,
 
@@ -280,15 +286,23 @@ func (c *chainClient) syncNonce() {
 }
 
 func (c *chainClient) syncTimeoutHeight() {
+	t := time.NewTicker(defaultTimeoutHeightSyncInterval)
+	defer t.Stop()
+
 	for {
-		ctx := context.Background()
-		block, err := c.ctx.Client.Block(ctx, nil)
+		block, err := c.ctx.Client.Block(c.cancelCtx, nil)
 		if err != nil {
 			c.logger.WithError(err).Errorln("failed to get current block")
 			return
 		}
 		c.txFactory.WithTimeoutHeight(uint64(block.Block.Height) + defaultTimeoutHeight)
-		time.Sleep(defaultTimeoutHeightSyncInterval)
+
+		select {
+		case <-c.cancelCtx.Done():
+			return
+		case <-t.C:
+			continue
+		}
 	}
 }
 
@@ -442,6 +456,10 @@ func (c *chainClient) Close() {
 	if atomic.CompareAndSwapInt64(&c.closed, 0, 1) {
 		close(c.msgC)
 	}
+
+	if c.cancelFn != nil {
+		c.cancelFn()
+	}
 	<-c.doneC
 	if c.conn != nil {
 		c.conn.Close()