From aa70d5d02ab7129cb93e7b44dffbbbce4efb62e9 Mon Sep 17 00:00:00 2001
From: carla <kirkcohenc@gmail.com>
Date: Fri, 13 Mar 2020 09:30:16 +0200
Subject: [PATCH] routerrpc: add quit channel to signal shutdown

---
 lnrpc/routerrpc/router_server.go | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/lnrpc/routerrpc/router_server.go b/lnrpc/routerrpc/router_server.go
index ed02a6fc..8bbdb387 100644
--- a/lnrpc/routerrpc/router_server.go
+++ b/lnrpc/routerrpc/router_server.go
@@ -9,6 +9,7 @@ import (
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"sync/atomic"
 
 	"github.com/btcsuite/btcutil"
 	"github.com/lightningnetwork/lnd/channeldb"
@@ -32,6 +33,8 @@ const (
 )
 
 var (
+	errServerShuttingDown = errors.New("routerrpc server shutting down")
+
 	// macaroonOps are the set of capabilities that our minted macaroon (if
 	// it doesn't already exist) will have.
 	macaroonOps = []bakery.Op{
@@ -90,7 +93,12 @@ var (
 // Server is a stand alone sub RPC server which exposes functionality that
 // allows clients to route arbitrary payment through the Lightning Network.
 type Server struct {
+	started  int32 // To be used atomically.
+	shutdown int32 // To be used atomically.
+
 	cfg *Config
+
+	quit chan struct{}
 }
 
 // A compile time check to ensure that Server fully implements the RouterServer
@@ -150,7 +158,8 @@ func New(cfg *Config) (*Server, lnrpc.MacaroonPerms, error) {
 	}
 
 	routerServer := &Server{
-		cfg: cfg,
+		cfg:  cfg,
+		quit: make(chan struct{}),
 	}
 
 	return routerServer, macPermissions, nil
@@ -160,6 +169,10 @@ func New(cfg *Config) (*Server, lnrpc.MacaroonPerms, error) {
 //
 // NOTE: This is part of the lnrpc.SubServer interface.
 func (s *Server) Start() error {
+	if atomic.AddInt32(&s.started, 1) != 1 {
+		return nil
+	}
+
 	return nil
 }
 
@@ -167,6 +180,11 @@ func (s *Server) Start() error {
 //
 // NOTE: This is part of the lnrpc.SubServer interface.
 func (s *Server) Stop() error {
+	if atomic.AddInt32(&s.shutdown, 1) != 1 {
+		return nil
+	}
+
+	close(s.quit)
 	return nil
 }
 
@@ -513,6 +531,9 @@ func (s *Server) trackPayment(paymentHash lntypes.Hash,
 			return err
 		}
 
+	case <-s.quit:
+		return errServerShuttingDown
+
 	case <-stream.Context().Done():
 		log.Debugf("Payment status stream %v canceled", paymentHash)
 		return stream.Context().Err()