From a1abb11dc5441e327bf6dd2884a11843574cbc9b Mon Sep 17 00:00:00 2001 From: Conner Fromknecht Date: Mon, 30 Jul 2018 20:51:08 -0700 Subject: [PATCH] peer: add ready chan arg to WaitForDisconnect This commit adds additional synchronization logic to WaitForDisconnect, such that it can be spawned before Start has been executed by the server. Without modification, the current version will return immediately since no goroutines will have been spawned. To solve this, we modify WaitForDisconnect to block until: 1) the peer is disconnected, 2) the peer is successfully started, before watching the waitgroup. In the first case, the waitgroup will block until all (if any) spawned goroutines have exited. Otherwise, if the Start is successful, we can switch to watching the waitgroup, knowing that waitgroup counter is positive. --- peer.go | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/peer.go b/peer.go index 74ff4c74..1040674b 100644 --- a/peer.go +++ b/peer.go @@ -567,8 +567,18 @@ func (p *peer) addLink(chanPoint *wire.OutPoint, // WaitForDisconnect waits until the peer has disconnected. A peer may be // disconnected if the local or remote side terminating the connection, or an -// irrecoverable protocol error has been encountered. -func (p *peer) WaitForDisconnect() { +// irrecoverable protocol error has been encountered. This method will only +// begin watching the peer's waitgroup after the ready channel or the peer's +// quit channel are signaled. The ready channel should only be signaled if a +// call to Start returns no error. Otherwise, if the peer fails to start, +// calling Disconnect will signal the quit channel and the method will not +// block, since no goroutines were spawned. +func (p *peer) WaitForDisconnect(ready chan struct{}) { + select { + case <-ready: + case <-p.quit: + } + p.wg.Wait() }