From 5b747715fc3eca343cfe6731568226c6fb33979d Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Tue, 19 May 2020 20:31:38 -0700
Subject: [PATCH] invoices: return error from SubscribeNotifications on
 shutdown

---
 invoices/invoiceregistry.go      |  7 +++++--
 invoices/invoiceregistry_test.go | 15 ++++++++++-----
 rpcserver.go                     |  5 ++++-
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/invoices/invoiceregistry.go b/invoices/invoiceregistry.go
index 662a8a82..b0e06a16 100644
--- a/invoices/invoiceregistry.go
+++ b/invoices/invoiceregistry.go
@@ -1182,7 +1182,9 @@ func (i *invoiceSubscriptionKit) notify(event *invoiceEvent) error {
 // added. The invoiceIndex parameter is a streaming "checkpoint". We'll start
 // by first sending out all new events with an invoice index _greater_ than
 // this value. Afterwards, we'll send out real-time notifications.
-func (i *InvoiceRegistry) SubscribeNotifications(addIndex, settleIndex uint64) *InvoiceSubscription {
+func (i *InvoiceRegistry) SubscribeNotifications(
+	addIndex, settleIndex uint64) (*InvoiceSubscription, error) {
+
 	client := &InvoiceSubscription{
 		NewInvoices:     make(chan *channeldb.Invoice),
 		SettledInvoices: make(chan *channeldb.Invoice),
@@ -1254,9 +1256,10 @@ func (i *InvoiceRegistry) SubscribeNotifications(addIndex, settleIndex uint64) *
 	select {
 	case i.newSubscriptions <- client:
 	case <-i.quit:
+		return nil, ErrShuttingDown
 	}
 
-	return client
+	return client, nil
 }
 
 // SubscribeSingleInvoice returns an SingleInvoiceSubscription which allows the
diff --git a/invoices/invoiceregistry_test.go b/invoices/invoiceregistry_test.go
index 319c30cf..aeb52125 100644
--- a/invoices/invoiceregistry_test.go
+++ b/invoices/invoiceregistry_test.go
@@ -9,6 +9,7 @@ import (
 	"github.com/lightningnetwork/lnd/lntypes"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/record"
+	"github.com/stretchr/testify/assert"
 )
 
 // TestSettleInvoice tests settling of an invoice and related notifications.
@@ -16,7 +17,8 @@ func TestSettleInvoice(t *testing.T) {
 	ctx := newTestContext(t)
 	defer ctx.cleanup()
 
-	allSubscriptions := ctx.registry.SubscribeNotifications(0, 0)
+	allSubscriptions, err := ctx.registry.SubscribeNotifications(0, 0)
+	assert.Nil(t, err)
 	defer allSubscriptions.Cancel()
 
 	// Subscribe to the not yet existing invoice.
@@ -221,11 +223,12 @@ func TestCancelInvoice(t *testing.T) {
 	ctx := newTestContext(t)
 	defer ctx.cleanup()
 
-	allSubscriptions := ctx.registry.SubscribeNotifications(0, 0)
+	allSubscriptions, err := ctx.registry.SubscribeNotifications(0, 0)
+	assert.Nil(t, err)
 	defer allSubscriptions.Cancel()
 
 	// Try to cancel the not yet existing invoice. This should fail.
-	err := ctx.registry.CancelInvoice(testInvoicePaymentHash)
+	err = ctx.registry.CancelInvoice(testInvoicePaymentHash)
 	if err != channeldb.ErrInvoiceNotFound {
 		t.Fatalf("expected ErrInvoiceNotFound, but got %v", err)
 	}
@@ -352,7 +355,8 @@ func TestSettleHoldInvoice(t *testing.T) {
 	}
 	defer registry.Stop()
 
-	allSubscriptions := registry.SubscribeNotifications(0, 0)
+	allSubscriptions, err := registry.SubscribeNotifications(0, 0)
+	assert.Nil(t, err)
 	defer allSubscriptions.Cancel()
 
 	// Subscribe to the not yet existing invoice.
@@ -651,7 +655,8 @@ func testKeySend(t *testing.T, keySendEnabled bool) {
 
 	ctx.registry.cfg.AcceptKeySend = keySendEnabled
 
-	allSubscriptions := ctx.registry.SubscribeNotifications(0, 0)
+	allSubscriptions, err := ctx.registry.SubscribeNotifications(0, 0)
+	assert.Nil(t, err)
 	defer allSubscriptions.Cancel()
 
 	hodlChan := make(chan interface{}, 1)
diff --git a/rpcserver.go b/rpcserver.go
index 9399d10f..13d9e01b 100644
--- a/rpcserver.go
+++ b/rpcserver.go
@@ -4535,9 +4535,12 @@ func (r *rpcServer) ListInvoices(ctx context.Context,
 func (r *rpcServer) SubscribeInvoices(req *lnrpc.InvoiceSubscription,
 	updateStream lnrpc.Lightning_SubscribeInvoicesServer) error {
 
-	invoiceClient := r.server.invoices.SubscribeNotifications(
+	invoiceClient, err := r.server.invoices.SubscribeNotifications(
 		req.AddIndex, req.SettleIndex,
 	)
+	if err != nil {
+		return err
+	}
 	defer invoiceClient.Cancel()
 
 	for {