multi: extend InvoiceDB methods with a context argument (#8066)

* multi: extend InvoiceDB methods with a context argument

This commit adds a context to InvoiceDB's methods. Along this refactor
we also extend InvoiceRegistry methods with contexts where it makes
sense. This change is essential to be able to provide kvdb and sqldb
implementations for InvoiceDB.

* channeldb: restrict invoice tests to only use an InvoiceDB instance

* docs: update release notes for 0.18.0

invoices/interface.go

@@ -1,6 +1,7 @@
 package invoices
 
 import (
+	"context"
 	"time"
 
 	"github.com/lightningnetwork/lnd/channeldb/models"
@@ -21,7 +22,8 @@ type InvoiceDB interface {
 	//
 	// NOTE: A side effect of this function is that it sets AddIndex on
 	// newInvoice.
-	AddInvoice(invoice *Invoice, paymentHash lntypes.Hash) (uint64, error)
+	AddInvoice(ctx context.Context, invoice *Invoice,
+		paymentHash lntypes.Hash) (uint64, error)
 
 	// InvoicesAddedSince can be used by callers to seek into the event
 	// time series of all the invoices added in the database. The specified
@@ -31,7 +33,8 @@ type InvoiceDB interface {
 	//
 	// NOTE: The index starts from 1, as a result. We enforce that
 	// specifying a value below the starting index value is a noop.
-	InvoicesAddedSince(sinceAddIndex uint64) ([]Invoice, error)
+	InvoicesAddedSince(ctx context.Context, sinceAddIndex uint64) (
+		[]Invoice, error)
 
 	// LookupInvoice attempts to look up an invoice according to its 32 byte
 	// payment hash. If an invoice which can settle the HTLC identified by
@@ -40,7 +43,7 @@ type InvoiceDB interface {
 	// Before setting the incoming HTLC, the values SHOULD be checked to
 	// ensure the payer meets the agreed upon contractual terms of the
 	// payment.
-	LookupInvoice(ref InvoiceRef) (Invoice, error)
+	LookupInvoice(ctx context.Context, ref InvoiceRef) (Invoice, error)
 
 	// ScanInvoices scans through all invoices and calls the passed scanFunc
 	// for each invoice with its respective payment hash. Additionally a
@@ -50,11 +53,12 @@ type InvoiceDB interface {
 	//
 	// TODO(positiveblue): abstract this functionality so it makes sense for
 	// other backends like sql.
-	ScanInvoices(scanFunc InvScanFunc, reset func()) error
+	ScanInvoices(ctx context.Context, scanFunc InvScanFunc,
+		reset func()) error
 
 	// QueryInvoices allows a caller to query the invoice database for
 	// invoices within the specified add index range.
-	QueryInvoices(q InvoiceQuery) (InvoiceSlice, error)
+	QueryInvoices(ctx context.Context, q InvoiceQuery) (InvoiceSlice, error)
 
 	// UpdateInvoice attempts to update an invoice corresponding to the
 	// passed payment hash. If an invoice matching the passed payment hash
@@ -67,7 +71,7 @@ type InvoiceDB interface {
 	//
 	// TODO(positiveblue): abstract this functionality so it makes sense for
 	// other backends like sql.
-	UpdateInvoice(ref InvoiceRef, setIDHint *SetID,
+	UpdateInvoice(ctx context.Context, ref InvoiceRef, setIDHint *SetID,
 		callback InvoiceUpdateCallback) (*Invoice, error)
 
 	// InvoicesSettledSince can be used by callers to catch up any settled
@@ -77,13 +81,15 @@ type InvoiceDB interface {
 	//
 	// NOTE: The index starts from 1, as a result. We enforce that
 	// specifying a value below the starting index value is a noop.
-	InvoicesSettledSince(sinceSettleIndex uint64) ([]Invoice, error)
+	InvoicesSettledSince(ctx context.Context, sinceSettleIndex uint64) (
+		[]Invoice, error)
 
 	// DeleteInvoice attempts to delete the passed invoices from the
 	// database in one transaction. The passed delete references hold all
 	// keys required to delete the invoices without also needing to
 	// deserialze them.
-	DeleteInvoice(invoicesToDelete []InvoiceDeleteRef) error
+	DeleteInvoice(ctx context.Context,
+		invoicesToDelete []InvoiceDeleteRef) error
 }
 
 // Payload abstracts access to any additional fields provided in the final hop's

invoices/invoiceregistry.go

@@ -1,6 +1,7 @@
 package invoices
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"sync"
@@ -179,7 +180,7 @@ func NewRegistry(idb InvoiceDB, expiryWatcher *InvoiceExpiryWatcher,
 // scanInvoicesOnStart will scan all invoices on start and add active invoices
 // to the invoice expiry watcher while also attempting to delete all canceled
 // invoices.
-func (i *InvoiceRegistry) scanInvoicesOnStart() error {
+func (i *InvoiceRegistry) scanInvoicesOnStart(ctx context.Context) error {
 	var (
 		pending   []invoiceExpiry
 		removable []InvoiceDeleteRef
@@ -225,7 +226,7 @@ func (i *InvoiceRegistry) scanInvoicesOnStart() error {
 		return nil
 	}
 
-	err := i.idb.ScanInvoices(scanFunc, reset)
+	err := i.idb.ScanInvoices(ctx, scanFunc, reset)
 	if err != nil {
 		return err
 	}
@@ -237,7 +238,7 @@ func (i *InvoiceRegistry) scanInvoicesOnStart() error {
 	if len(removable) > 0 {
 		log.Infof("Attempting to delete %v canceled invoices",
 			len(removable))
-		if err := i.idb.DeleteInvoice(removable); err != nil {
+		if err := i.idb.DeleteInvoice(ctx, removable); err != nil {
 			log.Warnf("Deleting canceled invoices failed: %v", err)
 		} else {
 			log.Infof("Deleted %v canceled invoices",
@@ -252,7 +253,9 @@ func (i *InvoiceRegistry) scanInvoicesOnStart() error {
 func (i *InvoiceRegistry) Start() error {
 	// Start InvoiceExpiryWatcher and prepopulate it with existing active
 	// invoices.
-	err := i.expiryWatcher.Start(i.cancelInvoiceImpl)
+	err := i.expiryWatcher.Start(func(hash lntypes.Hash, force bool) error {
+		return i.cancelInvoiceImpl(context.Background(), hash, force)
+	})
 	if err != nil {
 		return err
 	}
@@ -264,7 +267,7 @@ func (i *InvoiceRegistry) Start() error {
 
 	// Now scan all pending and removable invoices to the expiry watcher or
 	// delete them.
-	err = i.scanInvoicesOnStart()
+	err = i.scanInvoicesOnStart(context.Background())
 	if err != nil {
 		_ = i.Stop()
 		return err
@@ -484,15 +487,15 @@ func (i *InvoiceRegistry) dispatchToClients(event *invoiceEvent) {
 
 // deliverBacklogEvents will attempts to query the invoice database for any
 // notifications that the client has missed since it reconnected last.
-func (i *InvoiceRegistry) deliverBacklogEvents(
+func (i *InvoiceRegistry) deliverBacklogEvents(ctx context.Context,
 	client *InvoiceSubscription) error {
 
-	addEvents, err := i.idb.InvoicesAddedSince(client.addIndex)
+	addEvents, err := i.idb.InvoicesAddedSince(ctx, client.addIndex)
 	if err != nil {
 		return err
 	}
 
-	settleEvents, err := i.idb.InvoicesSettledSince(client.settleIndex)
+	settleEvents, err := i.idb.InvoicesSettledSince(ctx, client.settleIndex)
 	if err != nil {
 		return err
 	}
@@ -536,10 +539,10 @@ func (i *InvoiceRegistry) deliverBacklogEvents(
 // invoice subscribers will always receive the current state right after
 // subscribing. Only in case the invoice does not yet exist, nothing is sent
 // yet.
-func (i *InvoiceRegistry) deliverSingleBacklogEvents(
+func (i *InvoiceRegistry) deliverSingleBacklogEvents(ctx context.Context,
 	client *SingleInvoiceSubscription) error {
 
-	invoice, err := i.idb.LookupInvoice(client.invoiceRef)
+	invoice, err := i.idb.LookupInvoice(ctx, client.invoiceRef)
 
 	// It is possible that the invoice does not exist yet, but the client is
 	// already watching it in anticipation.
@@ -579,7 +582,7 @@ func (i *InvoiceRegistry) deliverSingleBacklogEvents(
 // addIndex of the newly created invoice which monotonically increases for each
 // new invoice added.  A side effect of this function is that it also sets
 // AddIndex on the invoice argument.
-func (i *InvoiceRegistry) AddInvoice(invoice *Invoice,
+func (i *InvoiceRegistry) AddInvoice(ctx context.Context, invoice *Invoice,
 	paymentHash lntypes.Hash) (uint64, error) {
 
 	i.Lock()
@@ -587,7 +590,7 @@ func (i *InvoiceRegistry) AddInvoice(invoice *Invoice,
 	ref := InvoiceRefByHash(paymentHash)
 	log.Debugf("Invoice%v: added with terms %v", ref, invoice.Terms)
 
-	addIndex, err := i.idb.AddInvoice(invoice, paymentHash)
+	addIndex, err := i.idb.AddInvoice(ctx, invoice, paymentHash)
 	if err != nil {
 		i.Unlock()
 		return 0, err
@@ -613,17 +616,21 @@ func (i *InvoiceRegistry) AddInvoice(invoice *Invoice,
 // then we're able to pull the funds pending within an HTLC.
 //
 // TODO(roasbeef): ignore if settled?
-func (i *InvoiceRegistry) LookupInvoice(rHash lntypes.Hash) (Invoice, error) {
+func (i *InvoiceRegistry) LookupInvoice(ctx context.Context,
+	rHash lntypes.Hash) (Invoice, error) {
+
 	// We'll check the database to see if there's an existing matching
 	// invoice.
 	ref := InvoiceRefByHash(rHash)
-	return i.idb.LookupInvoice(ref)
+	return i.idb.LookupInvoice(ctx, ref)
 }
 
 // LookupInvoiceByRef looks up an invoice by the given reference, if found
 // then we're able to pull the funds pending within an HTLC.
-func (i *InvoiceRegistry) LookupInvoiceByRef(ref InvoiceRef) (Invoice, error) {
-	return i.idb.LookupInvoice(ref)
+func (i *InvoiceRegistry) LookupInvoiceByRef(ctx context.Context,
+	ref InvoiceRef) (Invoice, error) {
+
+	return i.idb.LookupInvoice(ctx, ref)
 }
 
 // startHtlcTimer starts a new timer via the invoice registry main loop that
@@ -723,7 +730,8 @@ func (i *InvoiceRegistry) cancelSingleHtlc(invoiceRef InvoiceRef,
 	// no invoice update is performed, we can return early.
 	setID := (*SetID)(invoiceRef.SetID())
 	var updated bool
-	invoice, err := i.idb.UpdateInvoice(invoiceRef, setID,
+	invoice, err := i.idb.UpdateInvoice(
+		context.Background(), invoiceRef, setID,
 		func(invoice *Invoice) (
 			*InvoiceUpdateDesc, error) {
 
@@ -830,7 +838,7 @@ func (i *InvoiceRegistry) processKeySend(ctx invoiceUpdateCtx) error {
 
 	// Insert invoice into database. Ignore duplicates, because this
 	// may be a replay.
-	_, err = i.AddInvoice(invoice, ctx.hash)
+	_, err = i.AddInvoice(context.Background(), invoice, ctx.hash)
 	if err != nil && !errors.Is(err, ErrDuplicateInvoice) {
 		return err
 	}
@@ -888,7 +896,7 @@ func (i *InvoiceRegistry) processAMP(ctx invoiceUpdateCtx) error {
 	// Insert invoice into database. Ignore duplicates payment hashes and
 	// payment addrs, this may be a replay or a different HTLC for the AMP
 	// invoice.
-	_, err := i.AddInvoice(invoice, ctx.hash)
+	_, err := i.AddInvoice(context.Background(), invoice, ctx.hash)
 	isDuplicatedInvoice := errors.Is(err, ErrDuplicateInvoice)
 	isDuplicatedPayAddr := errors.Is(err, ErrDuplicatePayAddr)
 	switch {
@@ -1045,7 +1053,9 @@ func (i *InvoiceRegistry) notifyExitHopHtlcLocked(
 
 	invoiceRef := ctx.invoiceRef()
 	setID := (*SetID)(ctx.setID())
-	invoice, err := i.idb.UpdateInvoice(invoiceRef, setID, callback)
+	invoice, err := i.idb.UpdateInvoice(
+		context.Background(), invoiceRef, setID, callback,
+	)
 
 	var duplicateSetIDErr ErrDuplicateSetID
 	if errors.As(err, &duplicateSetIDErr) {
@@ -1228,7 +1238,9 @@ func (i *InvoiceRegistry) notifyExitHopHtlcLocked(
 }
 
 // SettleHodlInvoice sets the preimage of a hodl invoice.
-func (i *InvoiceRegistry) SettleHodlInvoice(preimage lntypes.Preimage) error {
+func (i *InvoiceRegistry) SettleHodlInvoice(ctx context.Context,
+	preimage lntypes.Preimage) error {
+
 	i.Lock()
 	defer i.Unlock()
 
@@ -1255,7 +1267,7 @@ func (i *InvoiceRegistry) SettleHodlInvoice(preimage lntypes.Preimage) error {
 
 	hash := preimage.Hash()
 	invoiceRef := InvoiceRefByHash(hash)
-	invoice, err := i.idb.UpdateInvoice(invoiceRef, nil, updateInvoice)
+	invoice, err := i.idb.UpdateInvoice(ctx, invoiceRef, nil, updateInvoice)
 	if err != nil {
 		log.Errorf("SettleHodlInvoice with preimage %v: %v",
 			preimage, err)
@@ -1290,8 +1302,10 @@ func (i *InvoiceRegistry) SettleHodlInvoice(preimage lntypes.Preimage) error {
 
 // CancelInvoice attempts to cancel the invoice corresponding to the passed
 // payment hash.
-func (i *InvoiceRegistry) CancelInvoice(payHash lntypes.Hash) error {
-	return i.cancelInvoiceImpl(payHash, true)
+func (i *InvoiceRegistry) CancelInvoice(ctx context.Context,
+	payHash lntypes.Hash) error {
+
+	return i.cancelInvoiceImpl(ctx, payHash, true)
 }
 
 // shouldCancel examines the state of an invoice and whether we want to
@@ -1312,8 +1326,8 @@ func shouldCancel(state ContractState, cancelAccepted bool) bool {
 // payment hash. Accepted invoices will only be canceled if explicitly
 // requested to do so. It notifies subscribing links and resolvers that
 // the associated htlcs were canceled if they change state.
-func (i *InvoiceRegistry) cancelInvoiceImpl(payHash lntypes.Hash,
-	cancelAccepted bool) error {
+func (i *InvoiceRegistry) cancelInvoiceImpl(ctx context.Context,
+	payHash lntypes.Hash, cancelAccepted bool) error {
 
 	i.Lock()
 	defer i.Unlock()
@@ -1338,7 +1352,7 @@ func (i *InvoiceRegistry) cancelInvoiceImpl(payHash lntypes.Hash,
 	}
 
 	invoiceRef := InvoiceRefByHash(payHash)
-	invoice, err := i.idb.UpdateInvoice(invoiceRef, nil, updateInvoice)
+	invoice, err := i.idb.UpdateInvoice(ctx, invoiceRef, nil, updateInvoice)
 
 	// Implement idempotency by returning success if the invoice was already
 	// canceled.
@@ -1391,7 +1405,7 @@ func (i *InvoiceRegistry) cancelInvoiceImpl(payHash lntypes.Hash,
 			deleteRef.PayAddr = &invoice.Terms.PaymentAddr
 		}
 
-		err = i.idb.DeleteInvoice([]InvoiceDeleteRef{deleteRef})
+		err = i.idb.DeleteInvoice(ctx, []InvoiceDeleteRef{deleteRef})
 		// If by any chance deletion failed, then log it instead of
 		// returning the error, as the invoice itself has already been
 		// canceled.
@@ -1522,7 +1536,7 @@ 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(
+func (i *InvoiceRegistry) SubscribeNotifications(ctx context.Context,
 	addIndex, settleIndex uint64) (*InvoiceSubscription, error) {
 
 	client := &InvoiceSubscription{
@@ -1612,7 +1626,7 @@ func (i *InvoiceRegistry) SubscribeNotifications(
 
 	// Query the database to see if based on the provided addIndex and
 	// settledIndex we need to deliver any backlog notifications.
-	err := i.deliverBacklogEvents(client)
+	err := i.deliverBacklogEvents(ctx, client)
 	if err != nil {
 		return nil, err
 	}
@@ -1624,7 +1638,7 @@ func (i *InvoiceRegistry) SubscribeNotifications(
 
 // SubscribeSingleInvoice returns an SingleInvoiceSubscription which allows the
 // caller to receive async notifications for a specific invoice.
-func (i *InvoiceRegistry) SubscribeSingleInvoice(
+func (i *InvoiceRegistry) SubscribeSingleInvoice(ctx context.Context,
 	hash lntypes.Hash) (*SingleInvoiceSubscription, error) {
 
 	client := &SingleInvoiceSubscription{
@@ -1686,7 +1700,7 @@ func (i *InvoiceRegistry) SubscribeSingleInvoice(
 	i.singleNotificationClients[client.id] = client
 	i.notificationClientMux.Unlock()
 
-	err := i.deliverSingleBacklogEvents(client)
+	err := i.deliverSingleBacklogEvents(ctx, client)
 	if err != nil {
 		return nil, err
 	}

invoices/invoiceregistry_test.go

@@ -1,6 +1,7 @@
 package invoices_test
 
 import (
+	"context"
 	"crypto/rand"
 	"math"
 	"testing"
@@ -23,13 +24,14 @@ func TestSettleInvoice(t *testing.T) {
 
 	ctx := newTestContext(t, nil)
 
-	allSubscriptions, err := ctx.registry.SubscribeNotifications(0, 0)
+	ctxb := context.Background()
+	allSubscriptions, err := ctx.registry.SubscribeNotifications(ctxb, 0, 0)
 	require.Nil(t, err)
 	defer allSubscriptions.Cancel()
 
 	// Subscribe to the not yet existing invoice.
 	subscription, err := ctx.registry.SubscribeSingleInvoice(
-		testInvoicePaymentHash,
+		ctxb, testInvoicePaymentHash,
 	)
 	if err != nil {
 		t.Fatal(err)
@@ -41,7 +43,7 @@ func TestSettleInvoice(t *testing.T) {
 	// Add the invoice.
 	testInvoice := newInvoice(t, false)
 	addIdx, err := ctx.registry.AddInvoice(
-		testInvoice, testInvoicePaymentHash,
+		ctxb, testInvoice, testInvoicePaymentHash,
 	)
 	if err != nil {
 		t.Fatal(err)
@@ -174,7 +176,7 @@ func TestSettleInvoice(t *testing.T) {
 
 	// Check that settled amount is equal to the sum of values of the htlcs
 	// 0 and 1.
-	inv, err := ctx.registry.LookupInvoice(testInvoicePaymentHash)
+	inv, err := ctx.registry.LookupInvoice(ctxb, testInvoicePaymentHash)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -184,7 +186,7 @@ func TestSettleInvoice(t *testing.T) {
 	}
 
 	// Try to cancel.
-	err = ctx.registry.CancelInvoice(testInvoicePaymentHash)
+	err = ctx.registry.CancelInvoice(ctxb, testInvoicePaymentHash)
 	require.ErrorIs(t, err, invpkg.ErrInvoiceAlreadySettled)
 
 	// As this is a direct settle, we expect nothing on the hodl chan.
@@ -205,31 +207,30 @@ func testCancelInvoice(t *testing.T, gc bool) {
 	cfg.GcCanceledInvoicesOnTheFly = gc
 	ctx := newTestContext(t, &cfg)
 
-	allSubscriptions, err := ctx.registry.SubscribeNotifications(0, 0)
+	ctxb := context.Background()
+	allSubscriptions, err := ctx.registry.SubscribeNotifications(ctxb, 0, 0)
 	require.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(ctxb, testInvoicePaymentHash)
 	require.ErrorIs(t, err, invpkg.ErrInvoiceNotFound)
 
 	// Subscribe to the not yet existing invoice.
 	subscription, err := ctx.registry.SubscribeSingleInvoice(
-		testInvoicePaymentHash,
+		ctxb, testInvoicePaymentHash,
 	)
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 	defer subscription.Cancel()
 
 	require.Equal(t, subscription.PayHash(), &testInvoicePaymentHash)
 
 	// Add the invoice.
 	testInvoice := newInvoice(t, false)
-	_, err = ctx.registry.AddInvoice(testInvoice, testInvoicePaymentHash)
-	if err != nil {
-		t.Fatal(err)
-	}
+	_, err = ctx.registry.AddInvoice(
+		ctxb, testInvoice, testInvoicePaymentHash,
+	)
+	require.NoError(t, err)
 
 	// We expect the open state to be sent to the single invoice subscriber.
 	select {
@@ -258,10 +259,8 @@ func testCancelInvoice(t *testing.T, gc bool) {
 	}
 
 	// Cancel invoice.
-	err = ctx.registry.CancelInvoice(testInvoicePaymentHash)
-	if err != nil {
-		t.Fatal(err)
-	}
+	err = ctx.registry.CancelInvoice(ctxb, testInvoicePaymentHash)
+	require.NoError(t, err)
 
 	// We expect the canceled state to be sent to the single invoice
 	// subscriber.
@@ -280,7 +279,7 @@ func testCancelInvoice(t *testing.T, gc bool) {
 	if gc {
 		// Check that the invoice has been deleted from the db.
 		_, err = ctx.idb.LookupInvoice(
-			invpkg.InvoiceRefByHash(testInvoicePaymentHash),
+			ctxb, invpkg.InvoiceRefByHash(testInvoicePaymentHash),
 		)
 		require.Error(t, err)
 	}
@@ -291,7 +290,7 @@ func testCancelInvoice(t *testing.T, gc bool) {
 	// Try to cancel again. Expect that we report ErrInvoiceNotFound if the
 	// invoice has been garbage collected (since the invoice has been
 	// deleted when it was canceled), and no error otherwise.
-	err = ctx.registry.CancelInvoice(testInvoicePaymentHash)
+	err = ctx.registry.CancelInvoice(ctxb, testInvoicePaymentHash)
 
 	if gc {
 		require.Error(t, err, invpkg.ErrInvoiceNotFound)
@@ -366,32 +365,27 @@ func TestSettleHoldInvoice(t *testing.T) {
 	registry := invpkg.NewRegistry(idb, expiryWatcher, &cfg)
 
 	err = registry.Start()
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 	defer registry.Stop()
 
-	allSubscriptions, err := registry.SubscribeNotifications(0, 0)
+	ctxb := context.Background()
+	allSubscriptions, err := registry.SubscribeNotifications(ctxb, 0, 0)
 	require.Nil(t, err)
 	defer allSubscriptions.Cancel()
 
 	// Subscribe to the not yet existing invoice.
 	subscription, err := registry.SubscribeSingleInvoice(
-		testInvoicePaymentHash,
+		ctxb, testInvoicePaymentHash,
 	)
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 	defer subscription.Cancel()
 
 	require.Equal(t, subscription.PayHash(), &testInvoicePaymentHash)
 
 	// Add the invoice.
 	invoice := newInvoice(t, true)
-	_, err = registry.AddInvoice(invoice, testInvoicePaymentHash)
-	if err != nil {
-		t.Fatal(err)
-	}
+	_, err = registry.AddInvoice(ctxb, invoice, testInvoicePaymentHash)
+	require.NoError(t, err)
 
 	// We expect the open state to be sent to the single invoice subscriber.
 	update := <-subscription.Updates
@@ -475,10 +469,8 @@ func TestSettleHoldInvoice(t *testing.T) {
 	}
 
 	// Settling with preimage should succeed.
-	err = registry.SettleHodlInvoice(testInvoicePreimage)
-	if err != nil {
-		t.Fatal("expected set preimage to succeed")
-	}
+	err = registry.SettleHodlInvoice(ctxb, testInvoicePreimage)
+	require.NoError(t, err, "expected set preimage to succeed")
 
 	htlcResolution, _ := (<-hodlChan).(invpkg.HtlcResolution)
 	require.NotNil(t, htlcResolution)
@@ -507,14 +499,14 @@ func TestSettleHoldInvoice(t *testing.T) {
 	}
 
 	// Idempotency.
-	err = registry.SettleHodlInvoice(testInvoicePreimage)
+	err = registry.SettleHodlInvoice(ctxb, testInvoicePreimage)
 	require.ErrorIs(t, err, invpkg.ErrInvoiceAlreadySettled)
 
 	// Try to cancel.
-	err = registry.CancelInvoice(testInvoicePaymentHash)
-	if err == nil {
-		t.Fatal("expected cancellation of a settled invoice to fail")
-	}
+	err = registry.CancelInvoice(ctxb, testInvoicePaymentHash)
+	require.Error(
+		t, err, "expected cancellation of a settled invoice to fail",
+	)
 }
 
 // TestCancelHoldInvoice tests canceling of a hold invoice and related
@@ -545,12 +537,12 @@ func TestCancelHoldInvoice(t *testing.T) {
 		require.NoError(t, registry.Stop())
 	})
 
+	ctxb := context.Background()
+
 	// Add the invoice.
 	invoice := newInvoice(t, true)
-	_, err = registry.AddInvoice(invoice, testInvoicePaymentHash)
-	if err != nil {
-		t.Fatal(err)
-	}
+	_, err = registry.AddInvoice(ctxb, invoice, testInvoicePaymentHash)
+	require.NoError(t, err)
 
 	amtPaid := lnwire.MilliSatoshi(100000)
 	hodlChan := make(chan interface{}, 1)
@@ -569,10 +561,8 @@ func TestCancelHoldInvoice(t *testing.T) {
 	}
 
 	// Cancel invoice.
-	err = registry.CancelInvoice(testInvoicePaymentHash)
-	if err != nil {
-		t.Fatal("cancel invoice failed")
-	}
+	err = registry.CancelInvoice(ctxb, testInvoicePaymentHash)
+	require.NoError(t, err, "cancel invoice failed")
 
 	htlcResolution, _ := (<-hodlChan).(invpkg.HtlcResolution)
 	require.NotNil(t, htlcResolution)
@@ -642,8 +632,10 @@ func testKeySend(t *testing.T, keySendEnabled bool) {
 	cfg.AcceptKeySend = keySendEnabled
 	ctx := newTestContext(t, &cfg)
 
-	allSubscriptions, err := ctx.registry.SubscribeNotifications(0, 0)
-	require.Nil(t, err)
+	allSubscriptions, err := ctx.registry.SubscribeNotifications(
+		context.Background(), 0, 0,
+	)
+	require.NoError(t, err)
 	defer allSubscriptions.Cancel()
 
 	hodlChan := make(chan interface{}, 1)
@@ -772,8 +764,9 @@ func testHoldKeysend(t *testing.T, timeoutKeysend bool) {
 	cfg.KeysendHoldTime = holdDuration
 	ctx := newTestContext(t, &cfg)
 
-	allSubscriptions, err := ctx.registry.SubscribeNotifications(0, 0)
-	require.Nil(t, err)
+	ctxb := context.Background()
+	allSubscriptions, err := ctx.registry.SubscribeNotifications(ctxb, 0, 0)
+	require.NoError(t, err)
 	defer allSubscriptions.Cancel()
 
 	hodlChan := make(chan interface{}, 1)
@@ -840,8 +833,8 @@ func testHoldKeysend(t *testing.T, timeoutKeysend bool) {
 	}
 
 	// Settle keysend payment manually.
-	require.Nil(t, ctx.registry.SettleHodlInvoice(
-		*newInvoice.Terms.PaymentPreimage,
+	require.NoError(t, ctx.registry.SettleHodlInvoice(
+		ctxb, *newInvoice.Terms.PaymentPreimage,
 	))
 
 	// We expect a settled notification to be sent out.
@@ -857,13 +850,14 @@ func TestMppPayment(t *testing.T) {
 	defer timeout()()
 
 	ctx := newTestContext(t, nil)
+	ctxb := context.Background()
 
 	// Add the invoice.
 	testInvoice := newInvoice(t, false)
-	_, err := ctx.registry.AddInvoice(testInvoice, testInvoicePaymentHash)
-	if err != nil {
-		t.Fatal(err)
-	}
+	_, err := ctx.registry.AddInvoice(
+		ctxb, testInvoice, testInvoicePaymentHash,
+	)
+	require.NoError(t, err)
 
 	mppPayload := &mockPayload{
 		mpp: record.NewMPP(testInvoiceAmount, [32]byte{}),
@@ -933,10 +927,8 @@ func TestMppPayment(t *testing.T) {
 
 	// Check that settled amount is equal to the sum of values of the htlcs
 	// 2 and 3.
-	inv, err := ctx.registry.LookupInvoice(testInvoicePaymentHash)
-	if err != nil {
-		t.Fatal(err)
-	}
+	inv, err := ctx.registry.LookupInvoice(ctxb, testInvoicePaymentHash)
+	require.NoError(t, err)
 	if inv.State != invpkg.ContractSettled {
 		t.Fatal("expected invoice to be settled")
 	}
@@ -952,17 +944,16 @@ func TestMppPayment(t *testing.T) {
 func TestMppPaymentWithOverpayment(t *testing.T) {
 	t.Parallel()
 
+	ctxb := context.Background()
 	f := func(overpaymentRand uint64) bool {
 		ctx := newTestContext(t, nil)
 
 		// Add the invoice.
 		testInvoice := newInvoice(t, false)
 		_, err := ctx.registry.AddInvoice(
-			testInvoice, testInvoicePaymentHash,
+			ctxb, testInvoice, testInvoicePaymentHash,
 		)
-		if err != nil {
-			t.Fatal(err)
-		}
+		require.NoError(t, err)
 
 		mppPayload := &mockPayload{
 			mpp: record.NewMPP(testInvoiceAmount, [32]byte{}),
@@ -1009,10 +1000,10 @@ func TestMppPaymentWithOverpayment(t *testing.T) {
 
 		// Check that settled amount is equal to the sum of values of
 		// the htlcs 1 and 2.
-		inv, err := ctx.registry.LookupInvoice(testInvoicePaymentHash)
-		if err != nil {
-			t.Fatal(err)
-		}
+		inv, err := ctx.registry.LookupInvoice(
+			ctxb, testInvoicePaymentHash,
+		)
+		require.NoError(t, err)
 		if inv.State != invpkg.ContractSettled {
 			t.Fatal("expected invoice to be settled")
 		}
@@ -1050,10 +1041,12 @@ func TestInvoiceExpiryWithRegistry(t *testing.T) {
 		t, testTime, 0, numExpired, numPending,
 	)
 
+	ctxb := context.Background()
+
 	var expectedCancellations []lntypes.Hash
 	expiredInvoices := existingInvoices.expiredInvoices
 	for paymentHash, expiredInvoice := range expiredInvoices {
-		_, err := idb.AddInvoice(expiredInvoice, paymentHash)
+		_, err := idb.AddInvoice(ctxb, expiredInvoice, paymentHash)
 		require.NoError(t, err)
 		expectedCancellations = append(
 			expectedCancellations, paymentHash,
@@ -1062,7 +1055,7 @@ func TestInvoiceExpiryWithRegistry(t *testing.T) {
 
 	pendingInvoices := existingInvoices.pendingInvoices
 	for paymentHash, pendingInvoice := range pendingInvoices {
-		_, err := idb.AddInvoice(pendingInvoice, paymentHash)
+		_, err := idb.AddInvoice(ctxb, pendingInvoice, paymentHash)
 		require.NoError(t, err)
 	}
 
@@ -1078,24 +1071,20 @@ func TestInvoiceExpiryWithRegistry(t *testing.T) {
 
 	var invoicesThatWillCancel []lntypes.Hash
 	for paymentHash, pendingInvoice := range newInvoices.pendingInvoices {
-		_, err := registry.AddInvoice(pendingInvoice, paymentHash)
+		_, err := registry.AddInvoice(ctxb, pendingInvoice, paymentHash)
+		require.NoError(t, err)
 		invoicesThatWillCancel = append(
 			invoicesThatWillCancel, paymentHash,
 		)
-		if err != nil {
-			t.Fatal(err)
-		}
 	}
 
 	// Check that they are really not canceled until before the clock is
 	// advanced.
 	for i := range invoicesThatWillCancel {
 		invoice, err := registry.LookupInvoice(
-			invoicesThatWillCancel[i],
+			ctxb, invoicesThatWillCancel[i],
 		)
-		if err != nil {
-			t.Fatalf("cannot find invoice: %v", err)
-		}
+		require.NoError(t, err, "cannot find invoice")
 
 		if invoice.State == invpkg.ContractCanceled {
 			t.Fatalf("expected pending invoice, got canceled")
@@ -1115,7 +1104,7 @@ func TestInvoiceExpiryWithRegistry(t *testing.T) {
 	canceled := func() bool {
 		for i := range expectedCancellations {
 			invoice, err := registry.LookupInvoice(
-				expectedCancellations[i],
+				ctxb, expectedCancellations[i],
 			)
 			require.NoError(t, err)
 
@@ -1162,6 +1151,8 @@ func TestOldInvoiceRemovalOnStart(t *testing.T) {
 		t, testTime, 0, numExpired, numPending,
 	)
 
+	ctxb := context.Background()
+
 	i := 0
 	for paymentHash, invoice := range existingInvoices.expiredInvoices {
 		// Mark half of the invoices as settled, the other half as
@@ -1172,7 +1163,7 @@ func TestOldInvoiceRemovalOnStart(t *testing.T) {
 			invoice.State = invpkg.ContractCanceled
 		}
 
-		_, err := idb.AddInvoice(invoice, paymentHash)
+		_, err := idb.AddInvoice(ctxb, invoice, paymentHash)
 		require.NoError(t, err)
 		i++
 	}
@@ -1186,7 +1177,7 @@ func TestOldInvoiceRemovalOnStart(t *testing.T) {
 		NumMaxInvoices: math.MaxUint64,
 	}
 
-	response, err := idb.QueryInvoices(query)
+	response, err := idb.QueryInvoices(ctxb, query)
 	require.NoError(t, err)
 
 	// Save all settled invoices for our expectation set.
@@ -1202,7 +1193,7 @@ func TestOldInvoiceRemovalOnStart(t *testing.T) {
 	require.NoError(t, err, "cannot start the registry")
 
 	// Perform a scan query to collect all invoices.
-	response, err = idb.QueryInvoices(query)
+	response, err = idb.QueryInvoices(ctxb, query)
 	require.NoError(t, err)
 
 	// Check that we really only kept the settled invoices after the
@@ -1248,7 +1239,10 @@ func testHeightExpiryWithRegistry(t *testing.T, numParts int, settle bool) {
 	testInvoice.HodlInvoice = true
 	testInvoice.PaymentRequest = []byte{1, 2, 3}
 
-	_, err := ctx.registry.AddInvoice(testInvoice, testInvoicePaymentHash)
+	ctxb := context.Background()
+	_, err := ctx.registry.AddInvoice(
+		ctxb, testInvoice, testInvoicePaymentHash,
+	)
 	require.NoError(t, err)
 
 	payLoad := testPayload
@@ -1275,7 +1269,9 @@ func testHeightExpiryWithRegistry(t *testing.T, numParts int, settle bool) {
 	}
 
 	require.Eventually(t, func() bool {
-		inv, err := ctx.registry.LookupInvoice(testInvoicePaymentHash)
+		inv, err := ctx.registry.LookupInvoice(
+			ctxb, testInvoicePaymentHash,
+		)
 		require.NoError(t, err)
 
 		return inv.State == invpkg.ContractAccepted
@@ -1299,7 +1295,7 @@ func testHeightExpiryWithRegistry(t *testing.T, numParts int, settle bool) {
 
 	// If we want to settle our invoice in this test, we do so now.
 	if settle {
-		err = ctx.registry.SettleHodlInvoice(testInvoicePreimage)
+		err = ctx.registry.SettleHodlInvoice(ctxb, testInvoicePreimage)
 		require.NoError(t, err)
 
 		for i := 0; i < numParts; i++ {
@@ -1333,7 +1329,7 @@ func testHeightExpiryWithRegistry(t *testing.T, numParts int, settle bool) {
 
 	// Finally, lookup the invoice and assert that we have the state we
 	// expect.
-	inv, err := ctx.registry.LookupInvoice(testInvoicePaymentHash)
+	inv, err := ctx.registry.LookupInvoice(ctxb, testInvoicePaymentHash)
 	require.NoError(t, err)
 	require.Equal(t, expectedState, inv.State, "expected "+
 		"hold invoice: %v, got: %v", expectedState, inv.State)
@@ -1352,7 +1348,10 @@ func TestMultipleSetHeightExpiry(t *testing.T) {
 	// Add a hold invoice.
 	testInvoice := newInvoice(t, true)
 
-	_, err := ctx.registry.AddInvoice(testInvoice, testInvoicePaymentHash)
+	ctxb := context.Background()
+	_, err := ctx.registry.AddInvoice(
+		ctxb, testInvoice, testInvoicePaymentHash,
+	)
 	require.NoError(t, err)
 
 	mppPayload := &mockPayload{
@@ -1409,7 +1408,7 @@ func TestMultipleSetHeightExpiry(t *testing.T) {
 
 	// Assert that we've reached an accepted state because the invoice has
 	// been paid with a complete set.
-	inv, err := ctx.registry.LookupInvoice(testInvoicePaymentHash)
+	inv, err := ctx.registry.LookupInvoice(ctxb, testInvoicePaymentHash)
 	require.NoError(t, err)
 	require.Equal(t, invpkg.ContractAccepted, inv.State, "expected "+
 		"hold invoice accepted")
@@ -1421,7 +1420,9 @@ func TestMultipleSetHeightExpiry(t *testing.T) {
 	}
 
 	require.Eventuallyf(t, func() bool {
-		inv, err := ctx.registry.LookupInvoice(testInvoicePaymentHash)
+		inv, err := ctx.registry.LookupInvoice(
+			ctxb, testInvoicePaymentHash,
+		)
 		require.NoError(t, err)
 
 		return inv.State == invpkg.ContractCanceled
@@ -1435,14 +1436,15 @@ func TestSettleInvoicePaymentAddrRequired(t *testing.T) {
 	t.Parallel()
 
 	ctx := newTestContext(t, nil)
+	ctxb := context.Background()
 
-	allSubscriptions, err := ctx.registry.SubscribeNotifications(0, 0)
-	require.Nil(t, err)
+	allSubscriptions, err := ctx.registry.SubscribeNotifications(ctxb, 0, 0)
+	require.NoError(t, err)
 	defer allSubscriptions.Cancel()
 
 	// Subscribe to the not yet existing invoice.
 	subscription, err := ctx.registry.SubscribeSingleInvoice(
-		testInvoicePaymentHash,
+		ctxb, testInvoicePaymentHash,
 	)
 	require.NoError(t, err)
 	defer subscription.Cancel()
@@ -1467,7 +1469,7 @@ func TestSettleInvoicePaymentAddrRequired(t *testing.T) {
 	// Add the invoice, which requires the MPP payload to always be
 	// included due to its set of feature bits.
 	addIdx, err := ctx.registry.AddInvoice(
-		invoice, testInvoicePaymentHash,
+		ctxb, invoice, testInvoicePaymentHash,
 	)
 	require.NoError(t, err)
 	require.Equal(t, int(addIdx), 1)
@@ -1523,14 +1525,15 @@ func TestSettleInvoicePaymentAddrRequiredOptionalGrace(t *testing.T) {
 	t.Parallel()
 
 	ctx := newTestContext(t, nil)
+	ctxb := context.Background()
 
-	allSubscriptions, err := ctx.registry.SubscribeNotifications(0, 0)
-	require.Nil(t, err)
+	allSubscriptions, err := ctx.registry.SubscribeNotifications(ctxb, 0, 0)
+	require.NoError(t, err)
 	defer allSubscriptions.Cancel()
 
 	// Subscribe to the not yet existing invoice.
 	subscription, err := ctx.registry.SubscribeSingleInvoice(
-		testInvoicePaymentHash,
+		ctxb, testInvoicePaymentHash,
 	)
 	require.NoError(t, err)
 	defer subscription.Cancel()
@@ -1556,7 +1559,7 @@ func TestSettleInvoicePaymentAddrRequiredOptionalGrace(t *testing.T) {
 	// Add the invoice, which does not require the MPP payload to always be
 	// included due to its set of feature bits.
 	addIdx, err := ctx.registry.AddInvoice(
-		invoice, testInvoicePaymentHash,
+		ctxb, invoice, testInvoicePaymentHash,
 	)
 	require.NoError(t, err)
 	require.Equal(t, int(addIdx), 1)
@@ -1725,8 +1728,9 @@ func testSpontaneousAmpPayment(
 	cfg := defaultRegistryConfig()
 	cfg.AcceptAMP = ampEnabled
 	ctx := newTestContext(t, &cfg)
+	ctxb := context.Background()
 
-	allSubscriptions, err := ctx.registry.SubscribeNotifications(0, 0)
+	allSubscriptions, err := ctx.registry.SubscribeNotifications(ctxb, 0, 0)
 	require.Nil(t, err)
 	defer allSubscriptions.Cancel()