From 6f1238ebc1ea1d771b0cb6542038de3d1fb29a6f Mon Sep 17 00:00:00 2001
From: corverroos <corver@obol.tech>
Date: Tue, 7 Jun 2022 07:22:27 +0200
Subject: [PATCH] dkg: fix wait peer issue (#671)

Fixes `waitPeer` issue where ping is closed as soon as connected.

category: bug
ticket: #586
---
 dkg/dkg.go | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/dkg/dkg.go b/dkg/dkg.go
index db14c1ac3..a6b38b5ec 100644
--- a/dkg/dkg.go
+++ b/dkg/dkg.go
@@ -50,6 +50,9 @@ type Config struct {
 
 // Run executes a dkg ceremony and writes secret share keystore and cluster lock files as output to disk.
 func Run(ctx context.Context, conf Config) (err error) {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	ctx = log.WithTopic(ctx, "dkg")
 	defer func() {
 		if err != nil {
@@ -128,7 +131,7 @@ func Run(ctx context.Context, conf Config) (err error) {
 
 		log.Info(ctx, "Connecting to peers...")
 
-		ctx, cancel, err := waitPeers(ctx, tcpNode, peers)
+		ctx, cancel, err = waitPeers(ctx, tcpNode, peers)
 		if err != nil {
 			return err
 		}
@@ -541,12 +544,20 @@ func waitPeers(ctx context.Context, tcpNode host.Host, peers []p2p.Peer) (contex
 				tuples <- tuple{Peer: pID, RTT: rtt}
 
 				// Wait for disconnect and cancel the context.
+				var err error
 				for result := range results {
 					if result.Error != nil {
-						log.Error(ctx, "Peer connection lost", result.Error, z.Str("peer", p2p.PeerName(pID)))
-						cancel()
+						err = result.Error
+						break
 					}
 				}
+
+				if ctx.Err() == nil {
+					log.Error(ctx, "Peer connection lost", err, z.Str("peer", p2p.PeerName(pID)))
+					cancel()
+				}
+
+				return
 			}
 		}(p.ID)
 	}
@@ -571,9 +582,6 @@ func waitPeers(ctx context.Context, tcpNode host.Host, peers []p2p.Peer) (contex
 
 // waitConnect blocks until a libp2p connection (ping) is established returning the ping result chan, with the peer or the context is cancelled.
 func waitConnect(ctx context.Context, tcpNode host.Host, p peer.ID) (<-chan ping.Result, time.Duration, bool) {
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
-
 	resp := ping.Ping(ctx, tcpNode, p)
 	for result := range resp {
 		if result.Error == nil {