accounts: thread context through to accounts

Author: ellemouton

accounts/checkers_test.go

@@ -499,8 +499,8 @@ func TestSendPaymentCalls(t *testing.T) {
 
 func testSendPayment(t *testing.T, uri string) {
 	var (
-		parentCtx = context.Background()
-		zeroFee   = &lnrpc.FeeLimit{Limit: &lnrpc.FeeLimit_Fixed{
+		ctx     = context.Background()
+		zeroFee = &lnrpc.FeeLimit{Limit: &lnrpc.FeeLimit_Fixed{
 			Fixed: 0,
 		}}
 		requestID uint64
@@ -520,7 +520,7 @@ func testSendPayment(t *testing.T, uri string) {
 	service, err := NewService(t.TempDir(), errFunc)
 	require.NoError(t, err)
 
-	err = service.Start(lndMock, routerMock, chainParams)
+	err = service.Start(ctx, lndMock, routerMock, chainParams)
 	require.NoError(t, err)
 
 	assertBalance := func(id AccountID, expectedBalance int64) {
@@ -533,7 +533,7 @@ func testSendPayment(t *testing.T, uri string) {
 
 	// This should error because there is no account in the context.
 	err = service.checkers.checkIncomingRequest(
-		parentCtx, uri, &lnrpc.SendRequest{},
+		ctx, uri, &lnrpc.SendRequest{},
 	)
 	require.ErrorContains(t, err, "no account found in context")
 
@@ -543,7 +543,7 @@ func testSendPayment(t *testing.T, uri string) {
 	)
 	require.NoError(t, err)
 
-	ctxWithAcct := AddAccountToContext(parentCtx, acct)
+	ctxWithAcct := AddAccountToContext(ctx, acct)
 
 	// This should error because there is no request ID in the context.
 	err = service.checkers.checkIncomingRequest(
@@ -552,7 +552,7 @@ func testSendPayment(t *testing.T, uri string) {
 	require.ErrorContains(t, err, "no request ID found in context")
 
 	reqID1 := nextRequestID()
-	ctx := AddRequestIDToContext(ctxWithAcct, reqID1)
+	ctx = AddRequestIDToContext(ctxWithAcct, reqID1)
 
 	// This should error because no payment hash is provided.
 	err = service.checkers.checkIncomingRequest(
@@ -698,7 +698,7 @@ func testSendPayment(t *testing.T, uri string) {
 func TestSendPaymentV2(t *testing.T) {
 	var (
 		uri       = "/routerrpc.Router/SendPaymentV2"
-		parentCtx = context.Background()
+		ctx       = context.Background()
 		requestID uint64
 	)
 
@@ -716,7 +716,7 @@ func TestSendPaymentV2(t *testing.T) {
 	service, err := NewService(t.TempDir(), errFunc)
 	require.NoError(t, err)
 
-	err = service.Start(lndMock, routerMock, chainParams)
+	err = service.Start(ctx, lndMock, routerMock, chainParams)
 	require.NoError(t, err)
 
 	assertBalance := func(id AccountID, expectedBalance int64) {
@@ -729,7 +729,7 @@ func TestSendPaymentV2(t *testing.T) {
 
 	// This should error because there is no account in the context.
 	err = service.checkers.checkIncomingRequest(
-		parentCtx, uri, &routerrpc.SendPaymentRequest{},
+		ctx, uri, &routerrpc.SendPaymentRequest{},
 	)
 	require.ErrorContains(t, err, "no account found in context")
 
@@ -739,7 +739,7 @@ func TestSendPaymentV2(t *testing.T) {
 	)
 	require.NoError(t, err)
 
-	ctxWithAcct := AddAccountToContext(parentCtx, acct)
+	ctxWithAcct := AddAccountToContext(ctx, acct)
 
 	// This should error because there is no request ID in the context.
 	err = service.checkers.checkIncomingRequest(
@@ -748,7 +748,7 @@ func TestSendPaymentV2(t *testing.T) {
 	require.ErrorContains(t, err, "no request ID found in context")
 
 	reqID1 := nextRequestID()
-	ctx := AddRequestIDToContext(ctxWithAcct, reqID1)
+	ctx = AddRequestIDToContext(ctxWithAcct, reqID1)
 
 	// This should error because no payment hash is provided.
 	err = service.checkers.checkIncomingRequest(
@@ -885,7 +885,7 @@ func TestSendPaymentV2(t *testing.T) {
 func TestSendToRouteV2(t *testing.T) {
 	var (
 		uri       = "/routerrpc.Router/SendToRouteV2"
-		parentCtx = context.Background()
+		ctx       = context.Background()
 		requestID uint64
 	)
 
@@ -903,7 +903,7 @@ func TestSendToRouteV2(t *testing.T) {
 	service, err := NewService(t.TempDir(), errFunc)
 	require.NoError(t, err)
 
-	err = service.Start(lndMock, routerMock, chainParams)
+	err = service.Start(ctx, lndMock, routerMock, chainParams)
 	require.NoError(t, err)
 
 	assertBalance := func(id AccountID, expectedBalance int64) {
@@ -916,7 +916,7 @@ func TestSendToRouteV2(t *testing.T) {
 
 	// This should error because there is no account in the context.
 	err = service.checkers.checkIncomingRequest(
-		parentCtx, uri, &routerrpc.SendToRouteRequest{},
+		ctx, uri, &routerrpc.SendToRouteRequest{},
 	)
 	require.ErrorContains(t, err, "no account found in context")
 
@@ -926,7 +926,7 @@ func TestSendToRouteV2(t *testing.T) {
 	)
 	require.NoError(t, err)
 
-	ctxWithAcct := AddAccountToContext(parentCtx, acct)
+	ctxWithAcct := AddAccountToContext(ctx, acct)
 
 	// This should error because there is no request ID in the context.
 	err = service.checkers.checkIncomingRequest(
@@ -935,7 +935,7 @@ func TestSendToRouteV2(t *testing.T) {
 	require.ErrorContains(t, err, "no request ID found in context")
 
 	reqID1 := nextRequestID()
-	ctx := AddRequestIDToContext(ctxWithAcct, reqID1)
+	ctx = AddRequestIDToContext(ctxWithAcct, reqID1)
 
 	// This should error because no payment hash is provided.
 	err = service.checkers.checkIncomingRequest(

accounts/service.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/btcsuite/btcd/chaincfg"
 	"github.com/lightninglabs/lndclient"
+	"github.com/lightninglabs/taproot-assets/fn"
 	"github.com/lightningnetwork/lnd/channeldb"
 	invpkg "github.com/lightningnetwork/lnd/invoices"
 	"github.com/lightningnetwork/lnd/lnrpc"
@@ -55,7 +56,7 @@ type InterceptorService struct {
 	routerClient lndclient.RouterClient
 
 	mainCtx       context.Context
-	contextCancel context.CancelFunc
+	contextCancel fn.Option[context.CancelFunc]
 
 	requestMtx sync.Mutex
 	checkers   *AccountChecker
@@ -85,12 +86,8 @@ func NewService(dir string,
 		return nil, err
 	}
 
-	mainCtx, contextCancel := context.WithCancel(context.Background())
-
 	return &InterceptorService{
 		store:              accountStore,
-		mainCtx:            mainCtx,
-		contextCancel:      contextCancel,
 		invoiceToAccount:   make(map[lntypes.Hash]AccountID),
 		pendingPayments:    make(map[lntypes.Hash]*trackedPayment),
 		requestValuesStore: newRequestValuesStore(),
@@ -101,9 +98,14 @@ func NewService(dir string,
 }
 
 // Start starts the account service and its interceptor capability.
-func (s *InterceptorService) Start(lightningClient lndclient.LightningClient,
+func (s *InterceptorService) Start(ctx context.Context,
+	lightningClient lndclient.LightningClient,
 	routerClient lndclient.RouterClient, params *chaincfg.Params) error {
 
+	mainCtx, contextCancel := context.WithCancel(ctx)
+	s.mainCtx = mainCtx
+	s.contextCancel = fn.Some(contextCancel)
+
 	s.routerClient = routerClient
 	s.checkers = NewAccountChecker(s, params)
 
@@ -180,7 +182,7 @@ func (s *InterceptorService) Start(lightningClient lndclient.LightningClient,
 	s.wg.Add(1)
 	go func() {
 		defer s.wg.Done()
-		defer s.contextCancel()
+		defer contextCancel()
 
 		for {
 			select {
@@ -235,9 +237,8 @@ func (s *InterceptorService) Stop() error {
 	s.requestMtx.Lock()
 	defer s.requestMtx.Unlock()
 
-	s.contextCancel()
+	s.contextCancel.WhenSome(func(fn context.CancelFunc) { fn() })
 	close(s.quit)
-
 	s.wg.Wait()
 
 	return s.store.Close()

accounts/service_test.go

@@ -838,7 +838,10 @@ func TestAccountService(t *testing.T) {
 			}
 
 			// Any errors during startup expected?
-			err = service.Start(lndMock, routerMock, chainParams)
+			err = service.Start(
+				context.Background(), lndMock, routerMock,
+				chainParams,
+			)
 			if tc.startupErr != "" {
 				require.ErrorContains(tt, err, tc.startupErr)
 

terminal.go

@@ -1042,7 +1042,7 @@ func (g *LightningTerminal) startInternalSubServers(ctx context.Context,
 	log.Infof("Starting LiT account service")
 	if !g.cfg.Accounts.Disable {
 		err = g.accountService.Start(
-			g.lndClient.Client, g.lndClient.Router,
+			ctx, g.lndClient.Client, g.lndClient.Router,
 			g.lndClient.ChainParams,
 		)
 		if err != nil {