From c5876e5d21a97ca77e83642c85b3590e31400353 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Wed, 27 Sep 2017 20:22:40 -0700
Subject: [PATCH] peer: avoid exit deadlock by ensuring WaitGroups are
 decremented before disconnect

---
 peer.go | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/peer.go b/peer.go
index 3fdc08d8..73e62f55 100644
--- a/peer.go
+++ b/peer.go
@@ -584,7 +584,6 @@ func (c *chanMsgStream) AddMsg(msg lnwire.Message) {
 //
 // NOTE: This method MUST be run as a goroutine.
 func (p *peer) readHandler() {
-	defer p.wg.Done()
 
 	chanMsgStreams := make(map[lnwire.ChannelID]*chanMsgStream)
 out:
@@ -708,6 +707,8 @@ out:
 		}
 	}
 
+	p.wg.Done()
+
 	p.Disconnect(errors.New("read handler closed"))
 
 	for cid, chanStream := range chanMsgStreams {
@@ -793,11 +794,8 @@ func (p *peer) writeMessage(msg lnwire.Message) error {
 //
 // NOTE: This method MUST be run as a goroutine.
 func (p *peer) writeHandler() {
-	defer func() {
-		p.wg.Done()
-		peerLog.Tracef("writeHandler for peer %v done", p)
-	}()
-
+	var exitErr error
+out:
 	for {
 		select {
 		case outMsg := <-p.sendQueue:
@@ -823,15 +821,21 @@ func (p *peer) writeHandler() {
 			}
 
 			if err != nil {
-				p.Disconnect(errors.Errorf("unable to write message: %v",
-					err))
-				return
+				exitErr = errors.Errorf("unable to write message: %v", err)
+				break out
 			}
 
 		case <-p.quit:
-			return
+			exitErr = errors.Errorf("peer exiting")
+			break out
 		}
 	}
+
+	p.wg.Done()
+
+	p.Disconnect(exitErr)
+
+	peerLog.Tracef("writeHandler for peer %v done", p)
 }
 
 // queueHandler is responsible for accepting messages from outside subsystems