From ca9174e166f863d7b75d0c2b50e8ef18ac8f1a1b Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Wed, 4 Apr 2018 17:43:51 -0700
Subject: [PATCH] peer: extend SendMessage to allow callers to block until msg
 is sent

---
 htlcswitch/interfaces.go |  7 +++----
 peer.go                  | 21 +++++++++++++++++----
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/htlcswitch/interfaces.go b/htlcswitch/interfaces.go
index 8086624a..ef00ff38 100644
--- a/htlcswitch/interfaces.go
+++ b/htlcswitch/interfaces.go
@@ -111,10 +111,9 @@ type ChannelLink interface {
 // Peer is an interface which represents the remote lightning node inside our
 // system.
 type Peer interface {
-	// SendMessage sends message to remote peer. The second arguments
-	// denote if the method should block until the message has been sent to
-	// the remote peer. If set, this allows the caller to more strongly
-	// synchronize.
+	// SendMessage sends message to remote peer. The second argument
+	// denotes if the method should block until the message has been sent
+	// to the remote peer.
 	SendMessage(msg lnwire.Message, sync bool) error
 
 	// WipeChannel removes the channel uniquely identified by its channel
diff --git a/peer.go b/peer.go
index c45db9df..b76c099d 100644
--- a/peer.go
+++ b/peer.go
@@ -1874,10 +1874,23 @@ func (p *peer) sendInitMsg() error {
 	return p.writeMessage(msg)
 }
 
-// SendMessage queues a message for sending to the target peer.
-func (p *peer) SendMessage(msg lnwire.Message) error {
-	p.queueMsg(msg, nil)
-	return nil
+// SendMessage sends message to remote peer. The second argument denotes if the
+// method should block until the message has been sent to the remote peer.
+func (p *peer) SendMessage(msg lnwire.Message, sync bool) error {
+	if !sync {
+		p.queueMsg(msg, nil)
+		return nil
+	}
+
+	errChan := make(chan error, 1)
+	p.queueMsg(msg, errChan)
+
+	select {
+	case err := <-errChan:
+		return err
+	case <-p.quit:
+		return fmt.Errorf("peer shutting down")
+	}
 }
 
 // PubKey returns the pubkey of the peer in compressed serialized format.