From db804541c77e89d1115aa17e0f1f425acbe6023d Mon Sep 17 00:00:00 2001
From: Stefan Liu <stefanliu@outlook.com>
Date: Thu, 23 Sep 2021 10:33:35 +0800
Subject: [PATCH] Sink upstream context into WaitTillHeight

---
 chains/chain.go | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/chains/chain.go b/chains/chain.go
index e932d2b..cd3fe94 100644
--- a/chains/chain.go
+++ b/chains/chain.go
@@ -18,6 +18,7 @@
 package chains
 
 import (
+	"context"
 	"fmt"
 	"strings"
 	"sync"
@@ -44,7 +45,7 @@ type SDK interface {
 
 type Nodes interface {
 	Height() uint64
-	WaitTillHeight(uint64, time.Duration) uint64
+	WaitTillHeight(context.Context, uint64, time.Duration) (uint64, bool)
 	Available() bool
 	Node() SDK
 }
@@ -77,7 +78,7 @@ func (s *ChainSDK) Height() uint64 {
 	return s.height
 }
 
-func (s *ChainSDK) WaitTillHeight(height uint64, interval time.Duration) uint64 {
+func (s *ChainSDK) WaitTillHeight(ctx context.Context, height uint64, interval time.Duration) (uint64, bool) {
 	if interval == 0 {
 		interval = s.interval
 	}
@@ -86,9 +87,13 @@ func (s *ChainSDK) WaitTillHeight(height uint64, interval time.Duration) uint64
 		if err != nil {
 			log.Error("Failed to get chain latest height err ", "chain", s.ChainID, "err", err)
 		} else if h >= height {
-			return h
+			return h, true
+		}
+		select {
+		case <-ctx.Done():
+			return h, false
+		case <-time.After(interval):
 		}
-		time.Sleep(interval)
 	}
 }