multi: move payment state handling into MPPayment

This commit moves the struct `paymentState` used in `routing` into
`channeldb` and replaces it with `MPPaymentState`. In the following
commit we'd see the benefit, that we don't need to pass variables back
and forth between the two packages. More importantly, this state is put
closer to its origin, and is strictly updated whenever a payment is read
from disk. This approach is less error-prone comparing to the previous
one, which both the `payment` and `paymentState` need to be updated at
the same time to make sure the data stay consistant in a parallel
environment.

channeldb/mp_payment.go

@@ -3,6 +3,7 @@ package channeldb
 import (
 	"bytes"
 	"errors"
+	"fmt"
 	"io"
 	"math"
 	"time"
@@ -147,6 +148,27 @@ type HTLCFailInfo struct {
 	FailureSourceIndex uint32
 }
 
+// MPPaymentState wraps a series of info needed for a given payment, which is
+// used by both MPP and AMP. This is a memory representation of the payment's
+// current state and is updated whenever the payment is read from disk.
+type MPPaymentState struct {
+	// NumAttemptsInFlight specifies the number of HTLCs the payment is
+	// waiting results for.
+	NumAttemptsInFlight int
+
+	// RemainingAmt specifies how much more money to be sent.
+	RemainingAmt lnwire.MilliSatoshi
+
+	// FeesPaid specifies the total fees paid so far that can be used to
+	// calculate remaining fee budget.
+	FeesPaid lnwire.MilliSatoshi
+
+	// Terminate indicates the payment is in its final stage and no more
+	// shards should be launched. This value is true if we have an HTLC
+	// settled or the payment has an error.
+	Terminate bool
+}
+
 // MPPayment is a wrapper around a payment's PaymentCreationInfo and
 // HTLCAttempts. All payments will have the PaymentCreationInfo set, any
 // HTLCs made in attempts to be completed will populated in the HTLCs slice.
@@ -175,6 +197,11 @@ type MPPayment struct {
 
 	// Status is the current PaymentStatus of this payment.
 	Status PaymentStatus
+
+	// State is the current state of the payment that holds a number of key
+	// insights and is used to determine what to do on each payment loop
+	// iteration.
+	State *MPPaymentState
 }
 
 // Terminated returns a bool to specify whether the payment is in a terminal
@@ -280,6 +307,55 @@ func (m *MPPayment) Registrable() error {
 	return nil
 }
 
+// setState creates and attaches a new MPPaymentState to the payment. It also
+// updates the payment's status based on its current state.
+func (m *MPPayment) setState() error {
+	// Fetch the total amount and fees that has already been sent in
+	// settled and still in-flight shards.
+	sentAmt, fees := m.SentAmt()
+
+	// Sanity check we haven't sent a value larger than the payment amount.
+	totalAmt := m.Info.Value
+	if sentAmt > totalAmt {
+		return fmt.Errorf("%w: sent=%v, total=%v", ErrSentExceedsTotal,
+			sentAmt, totalAmt)
+	}
+
+	// Get any terminal info for this payment.
+	settle, failure := m.TerminalInfo()
+
+	// If either an HTLC settled, or the payment has a payment level
+	// failure recorded, it means we should terminate the moment all shards
+	// have returned with a result.
+	terminate := settle != nil || failure != nil
+
+	// Now determine the payment's status.
+	status, err := decidePaymentStatus(m.HTLCs, m.FailureReason)
+	if err != nil {
+		return err
+	}
+
+	// Update the payment state and status.
+	m.State = &MPPaymentState{
+		NumAttemptsInFlight: len(m.InFlightHTLCs()),
+		RemainingAmt:        totalAmt - sentAmt,
+		FeesPaid:            fees,
+		Terminate:           terminate,
+	}
+	m.Status = status
+
+	return nil
+}
+
+// SetState calls the internal method setState. This is a temporary method
+// to be used by the tests in routing. Once the tests are updated to use mocks,
+// this method can be removed.
+//
+// TODO(yy): delete.
+func (m *MPPayment) SetState() error {
+	return m.setState()
+}
+
 // 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

@@ -5,6 +5,9 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/lightningnetwork/lnd/lntypes"
+	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/lightningnetwork/lnd/routing/route"
 	"github.com/stretchr/testify/require"
 )
 
@@ -119,3 +122,153 @@ func TestRegistrable(t *testing.T) {
 		})
 	}
 }
+
+// TestPaymentSetState checks that the method setState creates the
+// MPPaymentState as expected.
+func TestPaymentSetState(t *testing.T) {
+	t.Parallel()
+
+	// Create a test preimage and failure reason.
+	preimage := lntypes.Preimage{1}
+	failureReasonError := FailureReasonError
+
+	testCases := []struct {
+		name     string
+		payment  *MPPayment
+		totalAmt int
+
+		expectedState *MPPaymentState
+		errExpected   error
+	}{
+		{
+			// Test that when the sentAmt exceeds totalAmount, the
+			// error is returned.
+			name: "amount exceeded error",
+			// SentAmt returns 90, 10
+			// TerminalInfo returns non-nil, nil
+			// InFlightHTLCs returns 0
+			payment: &MPPayment{
+				HTLCs: []HTLCAttempt{
+					makeSettledAttempt(100, 10, preimage),
+				},
+			},
+			totalAmt:    1,
+			errExpected: ErrSentExceedsTotal,
+		},
+		{
+			// Test that when the htlc is failed, the fee is not
+			// used.
+			name: "fee excluded for failed htlc",
+			payment: &MPPayment{
+				// SentAmt returns 90, 10
+				// TerminalInfo returns nil, nil
+				// InFlightHTLCs returns 1
+				HTLCs: []HTLCAttempt{
+					makeActiveAttempt(100, 10),
+					makeFailedAttempt(100, 10),
+				},
+			},
+			totalAmt: 1000,
+			expectedState: &MPPaymentState{
+				NumAttemptsInFlight: 1,
+				RemainingAmt:        1000 - 90,
+				FeesPaid:            10,
+				Terminate:           false,
+			},
+		},
+		{
+			// Test when the payment is settled, the state should
+			// be marked as terminated.
+			name: "payment settled",
+			// SentAmt returns 90, 10
+			// TerminalInfo returns non-nil, nil
+			// InFlightHTLCs returns 0
+			payment: &MPPayment{
+				HTLCs: []HTLCAttempt{
+					makeSettledAttempt(100, 10, preimage),
+				},
+			},
+			totalAmt: 1000,
+			expectedState: &MPPaymentState{
+				NumAttemptsInFlight: 0,
+				RemainingAmt:        1000 - 90,
+				FeesPaid:            10,
+				Terminate:           true,
+			},
+		},
+		{
+			// Test when the payment is failed, the state should be
+			// marked as terminated.
+			name: "payment failed",
+			// SentAmt returns 0, 0
+			// TerminalInfo returns nil, non-nil
+			// InFlightHTLCs returns 0
+			payment: &MPPayment{
+				FailureReason: &failureReasonError,
+			},
+			totalAmt: 1000,
+			expectedState: &MPPaymentState{
+				NumAttemptsInFlight: 0,
+				RemainingAmt:        1000,
+				FeesPaid:            0,
+				Terminate:           true,
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			// Attach the payment info.
+			info := &PaymentCreationInfo{
+				Value: lnwire.MilliSatoshi(tc.totalAmt),
+			}
+			tc.payment.Info = info
+
+			// Call the method that updates the payment state.
+			err := tc.payment.setState()
+			require.ErrorIs(t, err, tc.errExpected)
+
+			require.Equal(
+				t, tc.expectedState, tc.payment.State,
+				"state not updated as expected",
+			)
+		})
+	}
+}
+
+func makeActiveAttempt(total, fee int) HTLCAttempt {
+	return HTLCAttempt{
+		HTLCAttemptInfo: makeAttemptInfo(total, total-fee),
+	}
+}
+
+func makeSettledAttempt(total, fee int,
+	preimage lntypes.Preimage) HTLCAttempt {
+
+	return HTLCAttempt{
+		HTLCAttemptInfo: makeAttemptInfo(total, total-fee),
+		Settle:          &HTLCSettleInfo{Preimage: preimage},
+	}
+}
+
+func makeFailedAttempt(total, fee int) HTLCAttempt {
+	return HTLCAttempt{
+		HTLCAttemptInfo: makeAttemptInfo(total, total-fee),
+		Failure: &HTLCFailInfo{
+			Reason: HTLCFailInternal,
+		},
+	}
+}
+
+func makeAttemptInfo(total, amtForwarded int) HTLCAttemptInfo {
+	hop := &route.Hop{AmtToForward: lnwire.MilliSatoshi(amtForwarded)}
+	return HTLCAttemptInfo{
+		Route: route.Route{
+			TotalAmount: lnwire.MilliSatoshi(total),
+			Hops:        []*route.Hop{hop},
+		},
+	}
+}

channeldb/payment_control.go

@@ -93,6 +93,10 @@ var (
 	// to a payment that already has a failure reason.
 	ErrPaymentPendingFailed = errors.New("payment has failure reason")
 
+	// ErrSentExceedsTotal is returned if the payment's current total sent
+	// amount exceed the total amount.
+	ErrSentExceedsTotal = errors.New("total sent exceeds total amount")
+
 	// errNoAttemptInfo is returned when no attempt info is stored yet.
 	errNoAttemptInfo = errors.New("unable to find attempt info for " +
 		"inflight payment")

channeldb/payments.go

@@ -298,19 +298,20 @@ func fetchPayment(bucket kvdb.RBucket) (*MPPayment, error) {
 		failureReason = &reason
 	}
 
-	// Now determine the payment's status.
-	paymentStatus, err := decidePaymentStatus(htlcs, failureReason)
-	if err != nil {
-		return nil, err
-	}
-
-	return &MPPayment{
+	// Create a new payment.
+	payment := &MPPayment{
 		SequenceNum:   sequenceNum,
 		Info:          creationInfo,
 		HTLCs:         htlcs,
 		FailureReason: failureReason,
-		Status:        paymentStatus,
-	}, nil
+	}
+
+	// Set its state and status.
+	if err := payment.setState(); err != nil {
+		return nil, err
+	}
+
+	return payment, nil
 }
 
 // fetchHtlcAttempts retrieves all htlc attempts made for the payment found in