channeldb: add method Registrable to decide adding HTLCs

This commit adds a new method, `Registrable`, to help decide whether
adding new HTLCs to a given payment is allowed. A new unit test,
`TestRegistrable` is also added to test it.

channeldb/mp_payment.go

@@ -235,6 +235,44 @@ func (m *MPPayment) GetAttempt(id uint64) (*HTLCAttempt, error) {
 	return nil, errors.New("htlc attempt not found on payment")
 }
 
+// Registrable returns an error to specify whether adding more HTLCs to the
+// payment with its current status is allowed. A payment can accept new HTLC
+// registrations when it's newly created, or none of its HTLCs is in a terminal
+// state.
+func (m *MPPayment) Registrable() error {
+	// Get the terminal info.
+	settle, reason := m.TerminalInfo()
+	settled := settle != nil
+	failed := reason != nil
+
+	// If updating the payment is not allowed, we can't register new HTLCs.
+	// Otherwise, the status must be either `StatusInitiated` or
+	// `StatusInFlight`.
+	if err := m.Status.updatable(); err != nil {
+		return err
+	}
+
+	// Exit early if this is not inflight.
+	if m.Status != StatusInFlight {
+		return nil
+	}
+
+	// There are still inflight HTLCs and we need to check whether there
+	// are settled HTLCs or the payment is failed. If we already have
+	// settled HTLCs, we won't allow adding more HTLCs.
+	if settled {
+		return ErrPaymentPendingSettled
+	}
+
+	// If the payment is already failed, we won't allow adding more HTLCs.
+	if failed {
+		return ErrPaymentPendingFailed
+	}
+
+	// Otherwise we can add more HTLCs.
+	return nil
+}
+
 // serializeHTLCSettleInfo serializes the details of a settled htlc.
 func serializeHTLCSettleInfo(w io.Writer, s *HTLCSettleInfo) error {
 	if _, err := w.Write(s.Preimage[:]); err != nil {

channeldb/mp_payment_test.go

@@ -2,6 +2,7 @@ package channeldb
 
 import (
 	"bytes"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -25,3 +26,96 @@ func TestLazySessionKeyDeserialize(t *testing.T) {
 	sessionKey := attempt.SessionKey()
 	require.Equal(t, priv, sessionKey)
 }
+
+// TestRegistrable checks the method `Registrable` behaves as expected for ALL
+// possible payment statuses.
+func TestRegistrable(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		status         PaymentStatus
+		registryErr    error
+		hasSettledHTLC bool
+		paymentFailed  bool
+	}{
+		{
+			status:      StatusInitiated,
+			registryErr: nil,
+		},
+		{
+			// Test inflight status with no settled HTLC and no
+			// failed payment.
+			status:      StatusInFlight,
+			registryErr: nil,
+		},
+		{
+			// Test inflight status with settled HTLC but no failed
+			// payment.
+			status:         StatusInFlight,
+			registryErr:    ErrPaymentPendingSettled,
+			hasSettledHTLC: true,
+		},
+		{
+			// Test inflight status with no settled HTLC but failed
+			// payment.
+			status:        StatusInFlight,
+			registryErr:   ErrPaymentPendingFailed,
+			paymentFailed: true,
+		},
+		{
+			// Test error state with settled HTLC and failed
+			// payment.
+			status:         0,
+			registryErr:    ErrUnknownPaymentStatus,
+			hasSettledHTLC: true,
+			paymentFailed:  true,
+		},
+		{
+			status:      StatusSucceeded,
+			registryErr: ErrPaymentAlreadySucceeded,
+		},
+		{
+			status:      StatusFailed,
+			registryErr: ErrPaymentAlreadyFailed,
+		},
+		{
+			status:      0,
+			registryErr: ErrUnknownPaymentStatus,
+		},
+	}
+
+	// Create test objects.
+	reason := FailureReasonError
+	htlcSettled := HTLCAttempt{
+		Settle: &HTLCSettleInfo{},
+	}
+
+	for i, tc := range testCases {
+		i, tc := i, tc
+
+		p := &MPPayment{
+			Status: tc.status,
+		}
+
+		// Add the settled htlc to the payment if needed.
+		htlcs := make([]HTLCAttempt, 0)
+		if tc.hasSettledHTLC {
+			htlcs = append(htlcs, htlcSettled)
+		}
+		p.HTLCs = htlcs
+
+		// Add the failure reason if needed.
+		if tc.paymentFailed {
+			p.FailureReason = &reason
+		}
+
+		name := fmt.Sprintf("test_%d_%s", i, p.Status.String())
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			err := p.Registrable()
+			require.ErrorIs(t, err, tc.registryErr,
+				"registrable under state %v", tc.status)
+		})
+	}
+}

channeldb/payment_control.go

@@ -85,6 +85,14 @@ var (
 	ErrMPPTotalAmountMismatch = errors.New("mp payment total amount " +
 		"mismatch")
 
+	// ErrPaymentPendingSettled is returned when we try to add a new
+	// attempt to a payment that has at least one of its HTLCs settled.
+	ErrPaymentPendingSettled = errors.New("payment has settled htlcs")
+
+	// ErrPaymentAlreadyFailed is returned when we try to add a new attempt
+	// to a payment that already has a failure reason.
+	ErrPaymentPendingFailed = errors.New("payment has failure reason")
+
 	// errNoAttemptInfo is returned when no attempt info is stored yet.
 	errNoAttemptInfo = errors.New("unable to find attempt info for " +
 		"inflight payment")
@@ -310,16 +318,8 @@ func (p *PaymentControl) RegisterAttempt(paymentHash lntypes.Hash,
 			return err
 		}
 
-		// We cannot register a new attempt if the payment already has
-		// reached a terminal condition. We check this before
-		// ensureInFlight because it is a more general check.
-		settle, fail := payment.TerminalInfo()
-		if settle != nil || fail != nil {
-			return ErrPaymentTerminal
-		}
-
-		// Ensure the payment is in-flight.
-		if err := ensureInFlight(payment); err != nil {
+		// Check if registering a new attempt is allowed.
+		if err := payment.Registrable(); err != nil {
 			return err
 		}
 
@@ -707,34 +707,6 @@ func fetchPaymentStatus(bucket kvdb.RBucket) (PaymentStatus, error) {
 	return payment.Status, nil
 }
 
-// ensureInFlight checks whether the payment found in the given bucket has
-// status InFlight, and returns an error otherwise. This should be used to
-// ensure we only mark in-flight payments as succeeded or failed.
-func ensureInFlight(payment *MPPayment) error {
-	paymentStatus := payment.Status
-
-	switch {
-	// Newly created payment is also inflight.
-	case paymentStatus == StatusInitiated:
-		return nil
-
-	// The payment was indeed InFlight.
-	case paymentStatus == StatusInFlight:
-		return nil
-
-	// The payment succeeded previously.
-	case paymentStatus == StatusSucceeded:
-		return ErrPaymentAlreadySucceeded
-
-	// The payment was already failed.
-	case paymentStatus == StatusFailed:
-		return ErrPaymentAlreadyFailed
-
-	default:
-		return ErrUnknownPaymentStatus
-	}
-}
-
 // FetchInFlightPayments returns all payments with status InFlight.
 func (p *PaymentControl) FetchInFlightPayments() ([]*MPPayment, error) {
 	var inFlights []*MPPayment

channeldb/payment_control_test.go

@@ -818,8 +818,10 @@ func TestPaymentControlMultiShard(t *testing.T) {
 		b = *attempt
 		b.AttemptID = 3
 		_, err = pControl.RegisterAttempt(info.PaymentIdentifier, &b)
-		if err != ErrPaymentTerminal {
-			t.Fatalf("expected ErrPaymentTerminal, got: %v", err)
+		if test.settleFirst {
+			require.ErrorIs(t, err, ErrPaymentPendingSettled)
+		} else {
+			require.ErrorIs(t, err, ErrPaymentPendingFailed)
 		}
 
 		assertPaymentStatus(t, pControl, info.PaymentIdentifier, StatusInFlight)
@@ -878,19 +880,41 @@ func TestPaymentControlMultiShard(t *testing.T) {
 			}
 		}
 
-		// If any of the two attempts settled, the payment should end
-		// up in the Succeeded state. If both failed the payment should
-		// also be Failed at this poinnt.
-		finalStatus := StatusFailed
-		if test.settleFirst || test.settleLast {
+		var (
+			finalStatus PaymentStatus
+			registerErr error
+		)
+
+		switch {
+		// If one of the attempts settled but the other failed with
+		// terminal error, we would still consider the payment is
+		// settled.
+		case test.settleFirst && !test.settleLast:
 			finalStatus = StatusSucceeded
+			registerErr = ErrPaymentAlreadySucceeded
+
+		case !test.settleFirst && test.settleLast:
+			finalStatus = StatusSucceeded
+			registerErr = ErrPaymentAlreadySucceeded
+
+		// If both failed, we end up in a failed status.
+		case !test.settleFirst && !test.settleLast:
+			finalStatus = StatusFailed
+			registerErr = ErrPaymentAlreadyFailed
+
+		// Otherwise, the payment has a succeed status.
+		case test.settleFirst && test.settleLast:
+			finalStatus = StatusSucceeded
+			registerErr = ErrPaymentAlreadySucceeded
 		}
 
-		assertPaymentStatus(t, pControl, info.PaymentIdentifier, finalStatus)
+		assertPaymentStatus(
+			t, pControl, info.PaymentIdentifier, finalStatus,
+		)
 
 		// Finally assert we cannot register more attempts.
 		_, err = pControl.RegisterAttempt(info.PaymentIdentifier, &b)
-		require.Equal(t, ErrPaymentTerminal, err)
+		require.Equal(t, registerErr, err)
 	}
 
 	for _, test := range tests {