refac lnd: subscribe invoices

network/backend/lnd/lnd.go

@@ -8,6 +8,8 @@ import (
 	"github.com/lightningnetwork/lnd/lnrpc"
 	"github.com/pkg/errors"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	"github.com/edouardparis/lntop/config"
 	"github.com/edouardparis/lntop/logging"
@@ -44,6 +46,7 @@ func (l Backend) Info(ctx context.Context) (*models.Info, error) {
 	if err != nil {
 		return nil, err
 	}
+	defer clt.Close()
 
 	resp, err := clt.GetInfo(ctx, &lnrpc.GetInfoRequest{})
 	if err != nil {
@@ -58,6 +61,7 @@ func (l Backend) SubscribeInvoice(ctx context.Context, channelInvoice chan *mode
 	if err != nil {
 		return err
 	}
+	defer clt.Close()
 
 	cltInvoices, err := clt.SubscribeInvoices(ctx, &lnrpc.InvoiceSubscription{})
 	if err != nil {
@@ -65,12 +69,22 @@ func (l Backend) SubscribeInvoice(ctx context.Context, channelInvoice chan *mode
 	}
 
 	for {
-		invoice, err := cltInvoices.Recv()
-		if err != nil {
-			return err
+		select {
+		case <-ctx.Done():
+			break
+		default:
+			invoice, err := cltInvoices.Recv()
+			if err != nil {
+				st, ok := status.FromError(err)
+				if ok && st.Code() == codes.Canceled {
+					l.logger.Debug("stopping subscribe invoice: context canceled")
+					return nil
+				}
+				return err
+			}
+
+			channelInvoice <- lookupInvoiceProtoToInvoice(invoice)
 		}
-
-		channelInvoice <- lookupInvoiceProtoToInvoice(invoice)
 	}
 }
 

pubsub/pubsub.go

@@ -30,7 +30,7 @@ func newPubSub(logger logging.Logger, network *network.Network) *pubSub {
 }
 
 func (p *pubSub) invoices(ctx context.Context, sub chan *events.Event) {
-	p.wg.Add(2)
+	p.wg.Add(3)
 	invoices := make(chan *models.Invoice)
 	ctx, cancel := context.WithCancel(ctx)
 
@@ -51,11 +51,15 @@ func (p *pubSub) invoices(ctx context.Context, sub chan *events.Event) {
 		if err != nil {
 			p.logger.Error("SubscribeInvoice returned an error", logging.Error(err))
 		}
-		close(invoices)
+		p.wg.Done()
 	}()
 
-	<-p.stop
-	cancel()
+	go func() {
+		<-p.stop
+		cancel()
+		close(invoices)
+		p.wg.Done()
+	}()
 }
 
 func (p *pubSub) wait() {