From dc238101183fae60ef5ea1799727ea529160349d Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Fri, 2 Nov 2018 11:33:01 +0100
Subject: [PATCH] rpcServer: wait for all goroutines to finish before exiting

Returning from the sendPayments method closes the underlying stream, so
we would risk payments still being in flight when returning.

We add the running goroutines to a waitgroup, such that we can wait for
them all to exit before exiting and closing the stream.

In case an error is encountered in the process, we will return directly,
which will close the `reqQuit` channel and prompt the gorpoutines to
shut down.
---
 rpcserver.go | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/rpcserver.go b/rpcserver.go
index 7dd1ef5e7..35eec9521 100644
--- a/rpcserver.go
+++ b/rpcserver.go
@@ -3680,18 +3680,23 @@ func (r *rpcServer) sendPayment(stream *paymentStream) error {
 		htlcSema <- struct{}{}
 	}
 
+	// We keep track of the running goroutines and set up a quit signal we
+	// can use to request them to exit if the method returns because of an
+	// encountered error.
+	var wg sync.WaitGroup
+	reqQuit := make(chan struct{})
+	defer close(reqQuit)
+
 	// Launch a new goroutine to handle reading new payment requests from
 	// the client. This way we can handle errors independently of blocking
 	// and waiting for the next payment request to come through.
-	reqQuit := make(chan struct{})
-	defer func() {
-		close(reqQuit)
-	}()
-
 	// TODO(joostjager): Callers expect result to come in in the same order
 	// as the request were sent, but this is far from guarantueed in the
 	// code below.
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
+
 		for {
 			select {
 			case <-reqQuit:
@@ -3778,11 +3783,18 @@ sendLoop:
 			// We launch a new goroutine to execute the current
 			// payment so we can continue to serve requests while
 			// this payment is being dispatched.
+			wg.Add(1)
 			go func() {
+				defer wg.Done()
+
 				// Attempt to grab a free semaphore slot, using
 				// a defer to eventually release the slot
 				// regardless of payment success.
-				<-htlcSema
+				select {
+				case <-htlcSema:
+				case <-reqQuit:
+					return
+				}
 				defer func() {
 					htlcSema <- struct{}{}
 				}()
@@ -3853,6 +3865,10 @@ sendLoop:
 			}()
 		}
 	}
+
+	// Wait for all goroutines to finish before closing the stream.
+	wg.Wait()
+	return nil
 }
 
 // SendPaymentSync is the synchronous non-streaming version of SendPayment.