From 8b0f6c355cc0e2bd42d7408484f72b282b7c5764 Mon Sep 17 00:00:00 2001
From: Yusef Napora <yusef@napora.org>
Date: Thu, 25 Jun 2020 13:00:55 -0400
Subject: [PATCH] read from sub.Done()

---
 lotus-soup/drand.go | 28 +++++++++++++++++-----------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/lotus-soup/drand.go b/lotus-soup/drand.go
index 0c455fb9d..507b05060 100644
--- a/lotus-soup/drand.go
+++ b/lotus-soup/drand.go
@@ -97,19 +97,20 @@ func prepareDrandNode(t *TestEnvironment) (*DrandInstance, error) {
 	var publicAddrs []string
 	var leaderAddr string
 	ch := make(chan *NodeAddr)
-	t.SyncClient.MustPublishSubscribe(ctx, addrTopic, &NodeAddr{
+	_, sub := t.SyncClient.MustPublishSubscribe(ctx, addrTopic, &NodeAddr{
 		PrivateAddr: n.PrivateAddr(),
 		PublicAddr:  n.PublicAddr(),
 		IsLeader:    isLeader,
 	}, ch)
 	for i := 0; i < nNodes; i++ {
-		msg, ok := <-ch
-		if !ok {
-			return nil, fmt.Errorf("failed to read drand node addr from sync service")
-		}
-		publicAddrs = append(publicAddrs, fmt.Sprintf("http://%s", msg.PublicAddr))
-		if msg.IsLeader {
-			leaderAddr = msg.PrivateAddr
+		select {
+		case msg := <-ch:
+			publicAddrs = append(publicAddrs, fmt.Sprintf("http://%s", msg.PublicAddr))
+			if msg.IsLeader {
+				leaderAddr = msg.PrivateAddr
+			}
+		case err := <-sub.Done():
+			return nil, fmt.Errorf("unable to read drand addrs from sync service: %w", err)
 		}
 	}
 	if leaderAddr == "" {
@@ -196,10 +197,15 @@ func prepareDrandNode(t *TestEnvironment) (*DrandInstance, error) {
 		}
 		infoCh := make(chan *peer.AddrInfo, nNodes)
 		infoTopic := sync.NewTopic("drand-gossip-addrs", &peer.AddrInfo{})
-		t.SyncClient.MustPublishSubscribe(ctx, infoTopic, relayInfo, infoCh)
+
+		_, sub := t.SyncClient.MustPublishSubscribe(ctx, infoTopic, relayInfo, infoCh)
 		for i := 0; i < nNodes; i++ {
-			ai := <-infoCh
-			relayAddrs = append(relayAddrs, *ai)
+			select {
+			case ai := <-infoCh:
+				relayAddrs = append(relayAddrs, *ai)
+			case err := <-sub.Done():
+				return nil, fmt.Errorf("unable to get drand relay addr from sync service: %w", err)
+			}
 		}
 	}