From 1adeb41a9d67801d79302a94119b2fbf60da51fd Mon Sep 17 00:00:00 2001
From: whythat <whythat@protonmail.com>
Date: Sat, 18 Jul 2020 23:41:11 +0300
Subject: [PATCH] rpcserver+lnrpc: make Subscribe RPCs context aware

This commit makes all the Subscribe RCP's context aware so that they
stop executing when the request context is cancelled.
---
 lnrpc/invoicesrpc/invoices_server.go |  3 +++
 lntest/itest/log_error_whitelist.txt |  4 ++++
 rpcserver.go                         | 22 ++++++++++++++++++++++
 3 files changed, 29 insertions(+)

diff --git a/lnrpc/invoicesrpc/invoices_server.go b/lnrpc/invoicesrpc/invoices_server.go
index 2fa881db..0c76b1d7 100644
--- a/lnrpc/invoicesrpc/invoices_server.go
+++ b/lnrpc/invoicesrpc/invoices_server.go
@@ -253,6 +253,9 @@ func (s *Server) SubscribeSingleInvoice(req *SubscribeSingleInvoiceRequest,
 				return nil
 			}
 
+		case <-updateStream.Context().Done():
+			return updateStream.Context().Err()
+
 		case <-s.quit:
 			return nil
 		}
diff --git a/lntest/itest/log_error_whitelist.txt b/lntest/itest/log_error_whitelist.txt
index 89dd0f35..08b3419f 100644
--- a/lntest/itest/log_error_whitelist.txt
+++ b/lntest/itest/log_error_whitelist.txt
@@ -269,3 +269,7 @@
 <time> [ERR] RPCS: WS: error writing message: websocket: close sent
 <time> [ERR] RPCS: [/routerrpc.Router/XImportMissionControl]: pair: <hex> -> <hex>: invalid failure: msat: <amt> and sat: 0.0000002 BTC values not equal
 <time> [ERR] BTCN: utxo scan failed: neutrino shutting down
+<time> [ERR] RPCS: [/lnrpc.Lightning/SubscribeChannelGraph]: context canceled
+<time> [ERR] RPCS: [/lnrpc.Lightning/SubscribeInvoices]: context canceled
+<time> [ERR] RPCS: [/lnrpc.Lightning/SubscribeChannelGraph]: context deadline exceeded
+<time> [ERR] RPCS: [/invoicesrpc.Invoices/SubscribeSingleInvoice]: context canceled
diff --git a/rpcserver.go b/rpcserver.go
index dd7e42c9..acb46076 100644
--- a/rpcserver.go
+++ b/rpcserver.go
@@ -2756,6 +2756,10 @@ func (r *rpcServer) SubscribePeerEvents(req *lnrpc.PeerEventSubscription,
 			if err := eventStream.Send(event); err != nil {
 				return err
 			}
+
+		case <-eventStream.Context().Done():
+			return eventStream.Context().Err()
+
 		case <-r.quit:
 			return nil
 		}
@@ -4029,6 +4033,10 @@ func (r *rpcServer) SubscribeChannelEvents(req *lnrpc.ChannelEventSubscription,
 			if err := updateStream.Send(update); err != nil {
 				return err
 			}
+
+		case <-updateStream.Context().Done():
+			return updateStream.Context().Err()
+
 		case <-r.quit:
 			return nil
 		}
@@ -4946,6 +4954,9 @@ func (r *rpcServer) SubscribeInvoices(req *lnrpc.InvoiceSubscription,
 				return err
 			}
 
+		case <-updateStream.Context().Done():
+			return updateStream.Context().Err()
+
 		case <-r.quit:
 			return nil
 		}
@@ -5003,6 +5014,9 @@ func (r *rpcServer) SubscribeTransactions(req *lnrpc.GetTransactionsRequest,
 				return err
 			}
 
+		case <-updateStream.Context().Done():
+			return updateStream.Context().Err()
+
 		case <-r.quit:
 			return nil
 		}
@@ -5524,6 +5538,11 @@ func (r *rpcServer) SubscribeChannelGraph(req *lnrpc.GraphTopologySubscription,
 				return err
 			}
 
+		// The context was cancelled so we report a cancellation error
+		// and exit immediately.
+		case <-updateStream.Context().Done():
+			return updateStream.Context().Err()
+
 		// The server is quitting, so we'll exit immediately. Returning
 		// nil will close the clients read end of the stream.
 		case <-r.quit:
@@ -6446,6 +6465,9 @@ func (r *rpcServer) SubscribeChannelBackups(req *lnrpc.ChannelBackupSubscription
 				return err
 			}
 
+		case <-updateStream.Context().Done():
+			return updateStream.Context().Err()
+
 		case <-r.quit:
 			return nil
 		}