diff --git a/network/backend/lnd/lnd.go b/network/backend/lnd/lnd.go index ab7a5c8..e7ad384 100644 --- a/network/backend/lnd/lnd.go +++ b/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) } } diff --git a/pubsub/pubsub.go b/pubsub/pubsub.go index 9bee295..620847e 100644 --- a/pubsub/pubsub.go +++ b/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() {