accounts: add new UpsertAccountPayment Store method

And use it in the InterceptorService's paymentUpdate method.

Author: ellemouton

accounts/interface.go

@@ -235,6 +235,14 @@ type Store interface {
 	IncreaseAccountBalance(ctx context.Context, id AccountID,
 		amount lnwire.MilliSatoshi) error
 
+	// UpsertAccountPayment updates or inserts a payment entry for the given
+	// account. Various functional options can be passed to modify the
+	// behavior of the method.
+	UpsertAccountPayment(_ context.Context, id AccountID,
+		paymentHash lntypes.Hash, fullAmount lnwire.MilliSatoshi,
+		status lnrpc.Payment_PaymentStatus,
+		options ...UpsertPaymentOption) error
+
 	// RemoveAccount finds an account by its ID and removes it from the¨
 	// store.
 	RemoveAccount(ctx context.Context, id AccountID) error
@@ -316,3 +324,41 @@ type RequestValuesStore interface {
 	// DeleteValues deletes any values stored for the given request ID.
 	DeleteValues(reqID uint64)
 }
+
+// UpsertPaymentOption is a functional option that can be passed to the
+// UpsertAccountPayment method to modify its behavior.
+type UpsertPaymentOption func(*upsertAcctPaymentOption)
+
+// upsertAcctPaymentOption is a struct that holds optional parameters for the
+// UpsertAccountPayment method.
+type upsertAcctPaymentOption struct {
+	debitAccount        bool
+	errIfAlreadyPending bool
+}
+
+// newUpsertPaymentOption creates a new upsertAcctPaymentOption with default
+// values.
+func newUpsertPaymentOption() *upsertAcctPaymentOption {
+	return &upsertAcctPaymentOption{
+		debitAccount:        false,
+		errIfAlreadyPending: false,
+	}
+}
+
+// WithDebitAccount is a functional option that can be passed to the
+// UpsertAccountPayment method to indicate that the account balance should be
+// debited by the full amount of the payment.
+func WithDebitAccount() UpsertPaymentOption {
+	return func(o *upsertAcctPaymentOption) {
+		o.debitAccount = true
+	}
+}
+
+// WithErrIfAlreadyPending is a functional option that can be passed to the
+// UpsertAccountPayment method to indicate that an error should be returned if
+// the payment is already pending or succeeded.
+func WithErrIfAlreadyPending() UpsertPaymentOption {
+	return func(o *upsertAcctPaymentOption) {
+		o.errIfAlreadyPending = true
+	}
+}

accounts/service.go

@@ -498,44 +498,22 @@ func (s *InterceptorService) AssociatePayment(ctx context.Context, id AccountID,
 	s.Lock()
 	defer s.Unlock()
 
-	account, err := s.store.Account(ctx, id)
-	if err != nil {
-		return err
-	}
-
-	// Check if this payment is associated with the account already.
-	_, ok := account.Payments[paymentHash]
-	if ok {
-		// We do not allow another payment to the same hash if the
-		// payment is already in-flight or succeeded. This mitigates a
-		// user being able to launch a second RPC-erring payment with
-		// the same hash that would remove the payment from being
-		// tracked. Note that this prevents launching multipart
-		// payments, but allows retrying a payment if it has failed.
-		if account.Payments[paymentHash].Status !=
-			lnrpc.Payment_FAILED {
-
-			return fmt.Errorf("payment with hash %s is already in "+
-				"flight or succeeded (status %v)", paymentHash,
-				account.Payments[paymentHash].Status)
-		}
-
-		// Otherwise, we fall through to correctly update the payment
-		// amount, in case we have a zero-amount invoice that is
-		// retried.
-	}
-
-	// Associate the payment with the account and store it.
-	account.Payments[paymentHash] = &PaymentEntry{
-		Status:     lnrpc.Payment_UNKNOWN,
-		FullAmount: fullAmt,
-	}
-
-	if err := s.store.UpdateAccount(ctx, account); err != nil {
-		return fmt.Errorf("error updating account: %w", err)
-	}
-
-	return nil
+	// We add the WithErrIfAlreadyPending option to ensure that if the
+	// payment is already associated with the account, then we return
+	// an error if the payment is already in-flight or succeeded. This
+	// prevents a user from being able to launch a second RPC-erring payment
+	// with the same hash that would remove the payment from being tracked.
+	//
+	// NOTE: this prevents launching multipart payments, but allows
+	// retrying a payment if it has failed.
+	//
+	// If the payment is already associated with the account but not in
+	// flight, we update the payment amount in case we have a zero-amount
+	// invoice that is retried.
+	return s.store.UpsertAccountPayment(
+		ctx, id, paymentHash, fullAmt, lnrpc.Payment_UNKNOWN,
+		WithErrIfAlreadyPending(),
+	)
 }
 
 // invoiceUpdate credits the account an invoice was registered with, in case the
@@ -832,23 +810,14 @@ func (s *InterceptorService) paymentUpdate(ctx context.Context,
 
 	// The payment went through! We now need to debit the full amount from
 	// the account.
-	account, err := s.store.Account(ctx, pendingPayment.accountID)
-	if err != nil {
-		err = s.disableAndErrorfUnsafe("error fetching account: %w",
-			err)
-
-		return terminalState, err
-	}
-
 	fullAmount := status.Value + status.Fee
 
-	// Update the account and store it in the database.
-	account.CurrentBalance -= int64(fullAmount)
-	account.Payments[hash] = &PaymentEntry{
-		Status:     lnrpc.Payment_SUCCEEDED,
-		FullAmount: fullAmount,
-	}
-	if err := s.store.UpdateAccount(ctx, account); err != nil {
+	// Update the persisted account.
+	err := s.store.UpsertAccountPayment(
+		ctx, pendingPayment.accountID, hash, fullAmount,
+		lnrpc.Payment_SUCCEEDED, WithDebitAccount(),
+	)
+	if err != nil {
 		err = s.disableAndErrorfUnsafe("error updating account: %w",
 			err)
 

accounts/store_kvdb.go

@@ -14,6 +14,7 @@ import (
 	"github.com/btcsuite/btcwallet/walletdb"
 	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/kvdb"
+	"github.com/lightningnetwork/lnd/lnrpc"
 	"github.com/lightningnetwork/lnd/lntypes"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"go.etcd.io/bbolt"
@@ -255,6 +256,52 @@ func (s *BoltStore) IncreaseAccountBalance(_ context.Context, id AccountID,
 	return s.updateAccount(id, update)
 }
 
+// UpsertAccountPayment updates or inserts a payment entry for the given
+// account. Various functional options can be passed to modify the behavior of
+// the method.
+//
+// NOTE: This is part of the Store interface.
+func (s *BoltStore) UpsertAccountPayment(_ context.Context, id AccountID,
+	paymentHash lntypes.Hash, fullAmount lnwire.MilliSatoshi,
+	status lnrpc.Payment_PaymentStatus,
+	options ...UpsertPaymentOption) error {
+
+	opts := newUpsertPaymentOption()
+	for _, o := range options {
+		o(opts)
+	}
+
+	update := func(account *OffChainBalanceAccount) error {
+		entry, ok := account.Payments[paymentHash]
+		if ok {
+			// If the errIfAlreadyPending option is set, we return
+			// an error if the payment is already in-flight or
+			// succeeded.
+			if opts.errIfAlreadyPending &&
+				entry.Status != lnrpc.Payment_FAILED {
+
+				return fmt.Errorf("payment with hash %s is "+
+					"already in flight or succeeded "+
+					"(status %v)", paymentHash,
+					account.Payments[paymentHash].Status)
+			}
+		}
+
+		account.Payments[paymentHash] = &PaymentEntry{
+			Status:     status,
+			FullAmount: fullAmount,
+		}
+
+		if opts.debitAccount {
+			account.CurrentBalance -= int64(fullAmount)
+		}
+
+		return nil
+	}
+
+	return s.updateAccount(id, update)
+}
+
 func (s *BoltStore) updateAccount(id AccountID,
 	updateFn func(*OffChainBalanceAccount) error) error {
 

accounts/store_test.go

@@ -257,6 +257,123 @@ func TestAccountUpdateMethods(t *testing.T) {
 
 		assertBalance(223)
 	})
+
+	t.Run("UpsertAccountPayment", func(t *testing.T) {
+		store := NewTestDB(t)
+
+		acct, err := store.NewAccount(ctx, 1000, time.Time{}, "foo")
+		require.NoError(t, err)
+
+		assertBalanceAndPayments := func(balance int64,
+			payments AccountPayments) {
+
+			dbAcct, err := store.Account(ctx, acct.ID)
+			require.NoError(t, err)
+			require.EqualValues(t, balance, dbAcct.CurrentBalance)
+
+			require.Len(t, dbAcct.Payments, len(payments))
+			for hash, payment := range payments {
+				dbPayment, ok := dbAcct.Payments[hash]
+				require.True(t, ok)
+				require.Equal(t, payment, dbPayment)
+			}
+		}
+
+		// The account initially has a balance of 1000 and no payments.
+		assertBalanceAndPayments(1000, nil)
+
+		// Assert that calling the method for a non-existent account
+		// errors out.
+		err = store.UpsertAccountPayment(
+			ctx, AccountID{}, lntypes.Hash{}, 0,
+			lnrpc.Payment_UNKNOWN,
+		)
+		require.ErrorIs(t, err, ErrAccNotFound)
+
+		// Add a payment to the account but don't update the balance.
+		// We do add a WithErrIfAlreadyPending option here just to show
+		// that no error is returned since the payment does not exist
+		// yet.
+		hash1 := lntypes.Hash{1, 2, 3, 4}
+		err = store.UpsertAccountPayment(
+			ctx, acct.ID, hash1, 600, lnrpc.Payment_UNKNOWN,
+			WithErrIfAlreadyPending(),
+		)
+		require.NoError(t, err)
+
+		assertBalanceAndPayments(1000, AccountPayments{
+			hash1: &PaymentEntry{
+				Status:     lnrpc.Payment_UNKNOWN,
+				FullAmount: 600,
+			},
+		})
+
+		// Add a second payment to the account and again don't update
+		// the balance.
+		hash2 := lntypes.Hash{5, 6, 7, 8}
+		err = store.UpsertAccountPayment(
+			ctx, acct.ID, hash2, 100, lnrpc.Payment_UNKNOWN,
+		)
+		require.NoError(t, err)
+
+		assertBalanceAndPayments(1000, AccountPayments{
+			hash1: &PaymentEntry{
+				Status:     lnrpc.Payment_UNKNOWN,
+				FullAmount: 600,
+			},
+			hash2: &PaymentEntry{
+				Status:     lnrpc.Payment_UNKNOWN,
+				FullAmount: 100,
+			},
+		})
+
+		// Now, update the first payment to have a new status and this
+		// time, debit the account.
+		err = store.UpsertAccountPayment(
+			ctx, acct.ID, hash1, 600, lnrpc.Payment_SUCCEEDED,
+			WithDebitAccount(),
+		)
+		require.NoError(t, err)
+
+		// The account should now have a balance of 400 and the first
+		// payment should have a status of succeeded.
+		assertBalanceAndPayments(400, AccountPayments{
+			hash1: &PaymentEntry{
+				Status:     lnrpc.Payment_SUCCEEDED,
+				FullAmount: 600,
+			},
+			hash2: &PaymentEntry{
+				Status:     lnrpc.Payment_UNKNOWN,
+				FullAmount: 100,
+			},
+		})
+
+		// Calling the same method again with the same payment hash
+		// should have no effect by default.
+		err = store.UpsertAccountPayment(
+			ctx, acct.ID, hash1, 600, lnrpc.Payment_SUCCEEDED,
+		)
+		require.NoError(t, err)
+
+		assertBalanceAndPayments(400, AccountPayments{
+			hash1: &PaymentEntry{
+				Status:     lnrpc.Payment_SUCCEEDED,
+				FullAmount: 600,
+			},
+			hash2: &PaymentEntry{
+				Status:     lnrpc.Payment_UNKNOWN,
+				FullAmount: 100,
+			},
+		})
+
+		// But, if we use the WithErrIfAlreadyPending option, we should
+		// get an error since the payment already exists.
+		err = store.UpsertAccountPayment(
+			ctx, acct.ID, hash1, 600, lnrpc.Payment_SUCCEEDED,
+			WithErrIfAlreadyPending(),
+		)
+		require.ErrorContains(t, err, "is already in flight")
+	})
 }
 
 // TestLastInvoiceIndexes makes sure the last known invoice indexes can be