routing: always update payment in the same goroutine

This commit refactors `collectResultAsync` such that this method is now
only responsible for collecting results from the switch. The method
`decideNextStep` is expanded to process these results in the same
goroutine where we fetch the payment from db, to make sure the lifecycle
loop always have a consistent view of a given payment.

routing/payment_lifecycle.go

@@ -24,6 +24,17 @@ import (
 // the payment lifecycle is exiting .
 var ErrPaymentLifecycleExiting = errors.New("payment lifecycle exiting")
 
+// switchResult is the result sent back from the switch after processing the
+// HTLC.
+type switchResult struct {
+	// attempt is the HTLC sent to the switch.
+	attempt *channeldb.HTLCAttempt
+
+	// result is sent from the switch which contains either a preimage if
+	// ths HTLC is settled or an error if it's failed.
+	result *htlcswitch.PaymentResult
+}
+
 // paymentLifecycle holds all information about the current state of a payment
 // needed to resume if from any point.
 type paymentLifecycle struct {
@@ -39,11 +50,9 @@ type paymentLifecycle struct {
 	// to stop.
 	quit chan struct{}
 
-	// resultCollected is used to signal that the result of an attempt has
+	// resultCollected is used to send the result returned from the switch
-	// been collected. A nil error means the attempt is either successful
+	// for a given HTLC attempt.
-	// or failed with temporary error. Otherwise, we should exit the
+	resultCollected chan *switchResult
-	// lifecycle loop as a terminal error has occurred.
-	resultCollected chan error
 
 	// resultCollector is a function that is used to collect the result of
 	// an HTLC attempt, which is always mounted to `p.collectResultAsync`
@@ -66,7 +75,7 @@ func newPaymentLifecycle(r *ChannelRouter, feeLimit lnwire.MilliSatoshi,
 		shardTracker:          shardTracker,
 		currentHeight:         currentHeight,
 		quit:                  make(chan struct{}),
-		resultCollected:       make(chan error, 1),
+		resultCollected:       make(chan *switchResult, 1),
 		firstHopCustomRecords: firstHopCustomRecords,
 	}
 
@@ -137,20 +146,27 @@ func (p *paymentLifecycle) decideNextStep(
 		log.Tracef("Waiting for attempt results for payment %v",
 			p.identifier)
 
-		// Otherwise we wait for one HTLC attempt then continue
+		// Otherwise we wait for the result for one HTLC attempt then
-		// the lifecycle.
+		// continue the lifecycle.
-		//
-		// NOTE: we don't check `p.quit` since `decideNextStep` is
-		// running in the same goroutine as `resumePayment`.
 		select {
-		case err := <-p.resultCollected:
+		case r := <-p.resultCollected:
-			// If an error is returned, exit with it.
+			log.Tracef("Received attempt result for payment %v",
+				p.identifier)
+
+			// Handle the result here. If there's no error, we will
+			// return stepSkip and move to the next lifecycle
+			// iteration, which will refresh the payment and wait
+			// for the next attempt result, if any.
+			_, err := p.handleAttemptResult(r.attempt, r.result)
+
+			// We would only get a DB-related error here, which will
+			// cause us to abort the payment flow.
 			if err != nil {
 				return stepExit, err
 			}
 
-			log.Tracef("Received attempt result for payment %v",
+		case <-p.quit:
-				p.identifier)
+			return stepExit, ErrPaymentLifecycleExiting
 
 		case <-p.router.quit:
 			return stepExit, ErrRouterShuttingDown
@@ -430,50 +446,57 @@ type attemptResult struct {
 }
 
 // collectResultAsync launches a goroutine that will wait for the result of the
-// given HTLC attempt to be available then handle its result. Once received, it
+// given HTLC attempt to be available then save its result in a map. Once
-// will send a nil error to channel `resultCollected` to indicate there's a
+// received, it will send the result returned from the switch to channel
-// result.
+// `resultCollected`.
 func (p *paymentLifecycle) collectResultAsync(attempt *channeldb.HTLCAttempt) {
 	log.Debugf("Collecting result for attempt %v in payment %v",
 		attempt.AttemptID, p.identifier)
 
 	go func() {
-		// Block until the result is available.
+		result, err := p.collectResult(attempt)
-		_, err := p.collectResult(attempt)
 		if err != nil {
-			log.Errorf("Error collecting result for attempt %v "+
+			log.Errorf("Error collecting result for attempt %v in "+
-				"in payment %v: %v", attempt.AttemptID,
+				"payment %v: %v", attempt.AttemptID,
 				p.identifier, err)
+
+			return
 		}
 
 		log.Debugf("Result collected for attempt %v in payment %v",
 			attempt.AttemptID, p.identifier)
 
-		// Once the result is collected, we signal it by writing the
+		// Create a switch result and send it to the resultCollected
-		// error to `resultCollected`.
+		// chan, which gets processed when the lifecycle is waiting for
+		// a result to be received in decideNextStep.
+		r := &switchResult{
+			attempt: attempt,
+			result:  result,
+		}
+
+		// Signal that a result has been collected.
 		select {
-		// Send the signal or quit.
+		// Send the result so decideNextStep can proceed.
-		case p.resultCollected <- err:
+		case p.resultCollected <- r:
 
 		case <-p.quit:
 			log.Debugf("Lifecycle exiting while collecting "+
 				"result for payment %v", p.identifier)
 
 		case <-p.router.quit:
-			return
 		}
 	}()
 }
 
-// collectResult waits for the result for the given attempt to be available
+// collectResult waits for the result of the given HTLC attempt to be sent by
-// from the Switch, then records the attempt outcome with the control tower.
+// the switch and returns it.
-// An attemptResult is returned, indicating the final outcome of this HTLC
+func (p *paymentLifecycle) collectResult(
-// attempt.
+	attempt *channeldb.HTLCAttempt) (*htlcswitch.PaymentResult, error) {
-func (p *paymentLifecycle) collectResult(attempt *channeldb.HTLCAttempt) (
-	*attemptResult, error) {
 
 	log.Tracef("Collecting result for attempt %v", spew.Sdump(attempt))
 
+	result := &htlcswitch.PaymentResult{}
+
 	// Regenerate the circuit for this attempt.
 	circuit, err := attempt.Circuit()
 
@@ -489,8 +512,7 @@ func (p *paymentLifecycle) collectResult(attempt *channeldb.HTLCAttempt) (
 	if err != nil {
 		log.Debugf("Unable to generate circuit for attempt %v: %v",
 			attempt.AttemptID, err)
-
+		return nil, err
-		return p.failAttempt(attempt.AttemptID, err)
 	}
 
 	// Using the created circuit, initialize the error decrypter, so we can
@@ -516,22 +538,21 @@ func (p *paymentLifecycle) collectResult(attempt *channeldb.HTLCAttempt) (
 		log.Errorf("Failed getting result for attemptID %d "+
 			"from switch: %v", attempt.AttemptID, err)
 
-		return p.handleSwitchErr(attempt, err)
+		result.Error = err
+
+		return result, nil
 	}
 
 	// The switch knows about this payment, we'll wait for a result to be
 	// available.
-	var (
-		result *htlcswitch.PaymentResult
-		ok     bool
-	)
-
 	select {
-	case result, ok = <-resultChan:
+	case r, ok := <-resultChan:
 		if !ok {
 			return nil, htlcswitch.ErrSwitchExiting
 		}
 
+		result = r
+
 	case <-p.quit:
 		return nil, ErrPaymentLifecycleExiting
 
@@ -539,46 +560,7 @@ func (p *paymentLifecycle) collectResult(attempt *channeldb.HTLCAttempt) (
 		return nil, ErrRouterShuttingDown
 	}
 
-	// In case of a payment failure, fail the attempt with the control
+	return result, nil
-	// tower and return.
-	if result.Error != nil {
-		return p.handleSwitchErr(attempt, result.Error)
-	}
-
-	// We successfully got a payment result back from the switch.
-	log.Debugf("Payment %v succeeded with pid=%v",
-		p.identifier, attempt.AttemptID)
-
-	// Report success to mission control.
-	err = p.router.cfg.MissionControl.ReportPaymentSuccess(
-		attempt.AttemptID, &attempt.Route,
-	)
-	if err != nil {
-		log.Errorf("Error reporting payment success to mc: %v", err)
-	}
-
-	// In case of success we atomically store settle result to the DB move
-	// the shard to the settled state.
-	htlcAttempt, err := p.router.cfg.Control.SettleAttempt(
-		p.identifier, attempt.AttemptID,
-		&channeldb.HTLCSettleInfo{
-			Preimage:   result.Preimage,
-			SettleTime: p.router.cfg.Clock.Now(),
-		},
-	)
-	if err != nil {
-		log.Errorf("Error settling attempt %v for payment %v with "+
-			"preimage %v: %v", attempt.AttemptID, p.identifier,
-			result.Preimage, err)
-
-		// We won't mark the attempt as failed since we already have
-		// the preimage.
-		return nil, err
-	}
-
-	return &attemptResult{
-		attempt: htlcAttempt,
-	}, nil
 }
 
 // registerAttempt is responsible for creating and saving an HTLC attempt in db
@@ -1111,3 +1093,65 @@ func (p *paymentLifecycle) reloadPayment() (DBMPPayment,
 
 	return payment, ps, nil
 }
+
+// handleAttemptResult processes the result of an HTLC attempt returned from
+// the htlcswitch.
+func (p *paymentLifecycle) handleAttemptResult(attempt *channeldb.HTLCAttempt,
+	result *htlcswitch.PaymentResult) (*attemptResult, error) {
+
+	// If the result has an error, we need to further process it by failing
+	// the attempt and maybe fail the payment.
+	if result.Error != nil {
+		return p.handleSwitchErr(attempt, result.Error)
+	}
+
+	// We got an attempt settled result back from the switch.
+	log.Debugf("Payment(%v): attempt(%v) succeeded", p.identifier,
+		attempt.AttemptID)
+
+	// Report success to mission control.
+	err := p.router.cfg.MissionControl.ReportPaymentSuccess(
+		attempt.AttemptID, &attempt.Route,
+	)
+	if err != nil {
+		log.Errorf("Error reporting payment success to mc: %v", err)
+	}
+
+	// In case of success we atomically store settle result to the DB and
+	// move the shard to the settled state.
+	htlcAttempt, err := p.router.cfg.Control.SettleAttempt(
+		p.identifier, attempt.AttemptID,
+		&channeldb.HTLCSettleInfo{
+			Preimage:   result.Preimage,
+			SettleTime: p.router.cfg.Clock.Now(),
+		},
+	)
+	if err != nil {
+		log.Errorf("Error settling attempt %v for payment %v with "+
+			"preimage %v: %v", attempt.AttemptID, p.identifier,
+			result.Preimage, err)
+
+		// We won't mark the attempt as failed since we already have
+		// the preimage.
+		return nil, err
+	}
+
+	return &attemptResult{
+		attempt: htlcAttempt,
+	}, nil
+}
+
+// collectAndHandleResult waits for the result for the given attempt to be
+// available from the Switch, then records the attempt outcome with the control
+// tower. An attemptResult is returned, indicating the final outcome of this
+// HTLC attempt.
+func (p *paymentLifecycle) collectAndHandleResult(
+	attempt *channeldb.HTLCAttempt) (*attemptResult, error) {
+
+	result, err := p.collectResult(attempt)
+	if err != nil {
+		return nil, err
+	}
+
+	return p.handleAttemptResult(attempt, result)
+}

routing/payment_lifecycle_test.go

@@ -522,7 +522,8 @@ func TestRequestRouteFailPaymentError(t *testing.T) {
 	ct.AssertExpectations(t)
 }
 
-// TestDecideNextStep checks the method `decideNextStep` behaves as expected.
+// TestDecideNextStep checks the method `decideNextStep` behaves as expected
+// given the returned values from `AllowMoreAttempts` and `NeedWaitAttempts`.
 func TestDecideNextStep(t *testing.T) {
 	t.Parallel()
 
@@ -537,15 +538,8 @@ func TestDecideNextStep(t *testing.T) {
 		name              string
 		allowMoreAttempts *mockReturn
 		needWaitAttempts  *mockReturn
-
+		expectedStep      stateStep
-		// When the attemptResultChan has returned.
+		expectedErr       error
-		closeResultChan bool
-
-		// Whether the router has quit.
-		routerQuit bool
-
-		expectedStep stateStep
-		expectedErr  error
 	}{
 		{
 			name:              "allow more attempts",
@@ -554,52 +548,36 @@ func TestDecideNextStep(t *testing.T) {
 			expectedErr:       nil,
 		},
 		{
-			name:              "error on allow more attempts",
+			name:              "error checking allow more attempts",
 			allowMoreAttempts: &mockReturn{false, errDummy},
 			expectedStep:      stepExit,
 			expectedErr:       errDummy,
 		},
 		{
-			name:              "no wait and exit",
+			name:              "no need to wait attempts",
 			allowMoreAttempts: &mockReturn{false, nil},
 			needWaitAttempts:  &mockReturn{false, nil},
 			expectedStep:      stepExit,
 			expectedErr:       nil,
 		},
 		{
-			name:              "wait returns an error",
+			name:              "error checking wait attempts",
 			allowMoreAttempts: &mockReturn{false, nil},
 			needWaitAttempts:  &mockReturn{false, errDummy},
 			expectedStep:      stepExit,
 			expectedErr:       errDummy,
 		},
-
-		{
-			name:              "wait and exit on result chan",
-			allowMoreAttempts: &mockReturn{false, nil},
-			needWaitAttempts:  &mockReturn{true, nil},
-			closeResultChan:   true,
-			expectedStep:      stepSkip,
-			expectedErr:       nil,
-		},
-		{
-			name:              "wait and exit on router quit",
-			allowMoreAttempts: &mockReturn{false, nil},
-			needWaitAttempts:  &mockReturn{true, nil},
-			routerQuit:        true,
-			expectedStep:      stepExit,
-			expectedErr:       ErrRouterShuttingDown,
-		},
 	}
 
 	for _, tc := range testCases {
 		tc := tc
 
 		// Create a test paymentLifecycle.
-		p := createTestPaymentLifecycle()
+		p, _ := newTestPaymentLifecycle(t)
 
 		// Make a mock payment.
 		payment := &mockMPPayment{}
+		defer payment.AssertExpectations(t)
 
 		// Mock the method AllowMoreAttempts.
 		payment.On("AllowMoreAttempts").Return(
@@ -615,29 +593,190 @@ func TestDecideNextStep(t *testing.T) {
 			).Once()
 		}
 
-		// Send a nil error to the attemptResultChan if requested.
-		if tc.closeResultChan {
-			p.resultCollected = make(chan error, 1)
-			p.resultCollected <- nil
-		}
-
-		// Quit the router if requested.
-		if tc.routerQuit {
-			close(p.router.quit)
-		}
-
 		// Once the setup is finished, run the test cases.
 		t.Run(tc.name, func(t *testing.T) {
 			step, err := p.decideNextStep(payment)
 			require.Equal(t, tc.expectedStep, step)
 			require.ErrorIs(t, tc.expectedErr, err)
 		})
-
-		// Check the payment's methods are called as expected.
-		payment.AssertExpectations(t)
 	}
 }
 
+// TestDecideNextStepOnRouterQuit checks the method `decideNextStep` behaves as
+// expected when the router is quit.
+func TestDecideNextStepOnRouterQuit(t *testing.T) {
+	t.Parallel()
+
+	// Create a test paymentLifecycle.
+	p, _ := newTestPaymentLifecycle(t)
+
+	// Make a mock payment.
+	payment := &mockMPPayment{}
+	defer payment.AssertExpectations(t)
+
+	// Mock the method AllowMoreAttempts to return false.
+	payment.On("AllowMoreAttempts").Return(false, nil).Once()
+
+	// Mock the method NeedWaitAttempts to wait for results.
+	payment.On("NeedWaitAttempts").Return(true, nil).Once()
+
+	// Quit the router.
+	close(p.router.quit)
+
+	// Call the method under test.
+	step, err := p.decideNextStep(payment)
+
+	// We expect stepExit and an error to be returned.
+	require.Equal(t, stepExit, step)
+	require.ErrorIs(t, err, ErrRouterShuttingDown)
+}
+
+// TestDecideNextStepOnLifecycleQuit checks the method `decideNextStep` behaves
+// as expected when the lifecycle is quit.
+func TestDecideNextStepOnLifecycleQuit(t *testing.T) {
+	t.Parallel()
+
+	// Create a test paymentLifecycle.
+	p, _ := newTestPaymentLifecycle(t)
+
+	// Make a mock payment.
+	payment := &mockMPPayment{}
+	defer payment.AssertExpectations(t)
+
+	// Mock the method AllowMoreAttempts to return false.
+	payment.On("AllowMoreAttempts").Return(false, nil).Once()
+
+	// Mock the method NeedWaitAttempts to wait for results.
+	payment.On("NeedWaitAttempts").Return(true, nil).Once()
+
+	// Quit the paymentLifecycle.
+	close(p.quit)
+
+	// Call the method under test.
+	step, err := p.decideNextStep(payment)
+
+	// We expect stepExit and an error to be returned.
+	require.Equal(t, stepExit, step)
+	require.ErrorIs(t, err, ErrPaymentLifecycleExiting)
+}
+
+// TestDecideNextStepHandleAttemptResultSucceed checks the method
+// `decideNextStep` behaves as expected when successfully handled the attempt
+// result.
+func TestDecideNextStepHandleAttemptResultSucceed(t *testing.T) {
+	t.Parallel()
+
+	// Create a test paymentLifecycle.
+	p, m := newTestPaymentLifecycle(t)
+
+	// Mock the clock to return a current time.
+	m.clock.On("Now").Return(time.Now())
+
+	// Make a mock payment.
+	payment := &mockMPPayment{}
+	defer payment.AssertExpectations(t)
+
+	// Mock the method AllowMoreAttempts to return false.
+	payment.On("AllowMoreAttempts").Return(false, nil).Once()
+
+	// Mock the method NeedWaitAttempts to wait for results.
+	payment.On("NeedWaitAttempts").Return(true, nil).Once()
+
+	paymentAmt := 10_000
+	preimage := lntypes.Preimage{1}
+	attempt := makeSettledAttempt(t, paymentAmt, preimage)
+
+	// Create a result that contains a preimage.
+	result := &htlcswitch.PaymentResult{
+		Preimage: preimage,
+	}
+
+	// Create a switch result and send it to the `resultCollected`` chan.
+	r := &switchResult{
+		attempt: attempt,
+		result:  result,
+	}
+	p.resultCollected <- r
+
+	// We now mock the behavior of `handleAttemptResult` - we are not
+	// testing this method's behavior here, so we simply mock it to return
+	// no error.
+	//
+	// Since the result doesn't contain an error, `ReportPaymentSuccess`
+	// should be called.
+	m.missionControl.On("ReportPaymentSuccess", mock.Anything,
+		mock.Anything).Return(nil).Once()
+
+	// The settled htlc should be returned from `SettleAttempt`.
+	m.control.On("SettleAttempt", mock.Anything, mock.Anything,
+		mock.Anything).Return(attempt, nil).Once()
+
+	// Call the method under test.
+	step, err := p.decideNextStep(payment)
+
+	// We expect stepSkip and no error to be returned.
+	require.Equal(t, stepSkip, step)
+	require.NoError(t, err)
+}
+
+// TestDecideNextStepHandleAttemptResultFail checks the method `decideNextStep`
+// behaves as expected when it fails to handle the attempt result.
+func TestDecideNextStepHandleAttemptResultFail(t *testing.T) {
+	t.Parallel()
+
+	// Create a test paymentLifecycle.
+	p, m := newTestPaymentLifecycle(t)
+
+	// Mock the clock to return a current time.
+	m.clock.On("Now").Return(time.Now())
+
+	// Make a mock payment.
+	payment := &mockMPPayment{}
+	defer payment.AssertExpectations(t)
+
+	// Mock the method AllowMoreAttempts to return false.
+	payment.On("AllowMoreAttempts").Return(false, nil).Once()
+
+	// Mock the method NeedWaitAttempts to wait for results.
+	payment.On("NeedWaitAttempts").Return(true, nil).Once()
+
+	paymentAmt := 10_000
+	preimage := lntypes.Preimage{1}
+	attempt := makeSettledAttempt(t, paymentAmt, preimage)
+
+	// Create a result that contains a preimage.
+	result := &htlcswitch.PaymentResult{
+		Preimage: preimage,
+	}
+
+	// Create a switch result and send it to the `resultCollected`` chan.
+	r := &switchResult{
+		attempt: attempt,
+		result:  result,
+	}
+	p.resultCollected <- r
+
+	// We now mock the behavior of `handleAttemptResult` - we are not
+	// testing this method's behavior here, so we simply mock it to return
+	// an error.
+	//
+	// Since the result doesn't contain an error, `ReportPaymentSuccess`
+	// should be called.
+	m.missionControl.On("ReportPaymentSuccess",
+		mock.Anything, mock.Anything).Return(nil).Once()
+
+	// Mock SettleAttempt to return an error.
+	m.control.On("SettleAttempt", mock.Anything, mock.Anything,
+		mock.Anything).Return(attempt, errDummy).Once()
+
+	// Call the method under test.
+	step, err := p.decideNextStep(payment)
+
+	// We expect stepExit and the above error to be returned.
+	require.Equal(t, stepExit, step)
+	require.ErrorIs(t, err, errDummy)
+}
+
 // TestResumePaymentFailOnFetchPayment checks when we fail to fetch the
 // payment, the error is returned.
 //
@@ -1333,7 +1472,7 @@ func TestCollectResultExitOnErr(t *testing.T) {
 	m.clock.On("Now").Return(time.Now())
 
 	// Now call the method under test.
-	result, err := p.collectResult(attempt)
+	result, err := p.collectAndHandleResult(attempt)
 	require.ErrorIs(t, err, errDummy, "expected dummy error")
 	require.Nil(t, result, "expected nil attempt")
 }
@@ -1379,7 +1518,7 @@ func TestCollectResultExitOnResultErr(t *testing.T) {
 	m.clock.On("Now").Return(time.Now())
 
 	// Now call the method under test.
-	result, err := p.collectResult(attempt)
+	result, err := p.collectAndHandleResult(attempt)
 	require.ErrorIs(t, err, errDummy, "expected dummy error")
 	require.Nil(t, result, "expected nil attempt")
 }
@@ -1405,7 +1544,7 @@ func TestCollectResultExitOnSwitchQuit(t *testing.T) {
 	})
 
 	// Now call the method under test.
-	result, err := p.collectResult(attempt)
+	result, err := p.collectAndHandleResult(attempt)
 	require.ErrorIs(t, err, htlcswitch.ErrSwitchExiting,
 		"expected switch exit")
 	require.Nil(t, result, "expected nil attempt")
@@ -1432,7 +1571,7 @@ func TestCollectResultExitOnRouterQuit(t *testing.T) {
 	})
 
 	// Now call the method under test.
-	result, err := p.collectResult(attempt)
+	result, err := p.collectAndHandleResult(attempt)
 	require.ErrorIs(t, err, ErrRouterShuttingDown, "expected router exit")
 	require.Nil(t, result, "expected nil attempt")
 }
@@ -1458,7 +1597,7 @@ func TestCollectResultExitOnLifecycleQuit(t *testing.T) {
 	})
 
 	// Now call the method under test.
-	result, err := p.collectResult(attempt)
+	result, err := p.collectAndHandleResult(attempt)
 	require.ErrorIs(t, err, ErrPaymentLifecycleExiting,
 		"expected lifecycle exit")
 	require.Nil(t, result, "expected nil attempt")
@@ -1502,7 +1641,7 @@ func TestCollectResultExitOnSettleErr(t *testing.T) {
 	m.clock.On("Now").Return(time.Now())
 
 	// Now call the method under test.
-	result, err := p.collectResult(attempt)
+	result, err := p.collectAndHandleResult(attempt)
 	require.ErrorIs(t, err, errDummy, "expected settle error")
 	require.Nil(t, result, "expected nil attempt")
 }
@@ -1544,7 +1683,7 @@ func TestCollectResultSuccess(t *testing.T) {
 	m.clock.On("Now").Return(time.Now())
 
 	// Now call the method under test.
-	result, err := p.collectResult(attempt)
+	result, err := p.collectAndHandleResult(attempt)
 	require.NoError(t, err, "expected no error")
 	require.Equal(t, preimage, result.attempt.Settle.Preimage,
 		"preimage mismatch")
@@ -1561,19 +1700,97 @@ func TestCollectResultAsyncSuccess(t *testing.T) {
 	preimage := lntypes.Preimage{1}
 	attempt := makeSettledAttempt(t, paymentAmt, preimage)
 
+	// Create a mock result returned from the switch.
+	result := &htlcswitch.PaymentResult{
+		Preimage: preimage,
+	}
+
 	// Mock the htlcswitch to return a the result chan.
 	resultChan := make(chan *htlcswitch.PaymentResult, 1)
 	m.payer.On("GetAttemptResult",
 		attempt.AttemptID, p.identifier, mock.Anything,
 	).Return(resultChan, nil).Once().Run(func(args mock.Arguments) {
 		// Send the preimage to the result chan.
-		resultChan <- &htlcswitch.PaymentResult{
+		resultChan <- result
-			Preimage: preimage,
-		}
 	})
 
-	// Once the result is received, `ReportPaymentSuccess` should be
+	// Now call the method under test.
-	// called.
+	p.collectResultAsync(attempt)
+
+	var r *switchResult
+
+	// Assert the result is returned within 5 seconds.
+	waitErr := wait.NoError(func() error {
+		r = <-p.resultCollected
+		return nil
+	}, testTimeout)
+	require.NoError(t, waitErr, "timeout waiting for result")
+
+	// Assert the result is received as expected.
+	require.Equal(t, attempt, r.attempt)
+	require.Equal(t, result, r.result)
+}
+
+// TestHandleAttemptResultWithError checks that when the `Error` field in the
+// result is not nil, it's properly handled by `handleAttemptResult`.
+func TestHandleAttemptResultWithError(t *testing.T) {
+	t.Parallel()
+
+	// Create a test paymentLifecycle with the initial two calls mocked.
+	p, m := newTestPaymentLifecycle(t)
+
+	paymentAmt := 10_000
+	preimage := lntypes.Preimage{1}
+	attempt := makeSettledAttempt(t, paymentAmt, preimage)
+
+	// Create a result that contains an error.
+	//
+	// NOTE: The error is chosen so we can quickly exit `handleSwitchErr`
+	// since we are not testing its behavior here.
+	result := &htlcswitch.PaymentResult{
+		Error: htlcswitch.ErrPaymentIDNotFound,
+	}
+
+	// The above error will end up being handled by `handleSwitchErr`, in
+	// which we'd cancel the shard and fail the attempt.
+	//
+	// `CancelShard` should be called with the attemptID.
+	m.shardTracker.On("CancelShard", attempt.AttemptID).Return(nil).Once()
+
+	// Mock `FailAttempt` to return a dummy error.
+	m.control.On("FailAttempt",
+		p.identifier, attempt.AttemptID, mock.Anything,
+	).Return(nil, errDummy).Once()
+
+	// Mock the clock to return a current time.
+	m.clock.On("Now").Return(time.Now())
+
+	// Call the method under test and expect the dummy error to be
+	// returned.
+	attemptResult, err := p.handleAttemptResult(attempt, result)
+	require.ErrorIs(t, err, errDummy, "expected fail error")
+	require.Nil(t, attemptResult, "expected nil attempt result")
+}
+
+// TestHandleAttemptResultSuccess checks that when the result contains no error
+// but a preimage, it's handled correctly by `handleAttemptResult`.
+func TestHandleAttemptResultSuccess(t *testing.T) {
+	t.Parallel()
+
+	// Create a test paymentLifecycle with the initial two calls mocked.
+	p, m := newTestPaymentLifecycle(t)
+
+	paymentAmt := 10_000
+	preimage := lntypes.Preimage{1}
+	attempt := makeSettledAttempt(t, paymentAmt, preimage)
+
+	// Create a result that contains a preimage.
+	result := &htlcswitch.PaymentResult{
+		Preimage: preimage,
+	}
+
+	// Since the result doesn't contain an error, `ReportPaymentSuccess`
+	// should be called.
 	m.missionControl.On("ReportPaymentSuccess",
 		attempt.AttemptID, &attempt.Route,
 	).Return(nil).Once()
@@ -1586,17 +1803,9 @@ func TestCollectResultAsyncSuccess(t *testing.T) {
 	// Mock the clock to return a current time.
 	m.clock.On("Now").Return(time.Now())
 
-	// Now call the method under test.
+	// Call the method under test and expect the dummy error to be
-	p.collectResultAsync(attempt)
+	// returned.
-
+	attemptResult, err := p.handleAttemptResult(attempt, result)
-	// Assert the result is returned within 5 seconds.
-	var err error
-	waitErr := wait.NoError(func() error {
-		err = <-p.resultCollected
-		return nil
-	}, testTimeout)
-	require.NoError(t, waitErr, "timeout waiting for result")
-
-	// Assert that a nil error is received.
 	require.NoError(t, err, "expected no error")
+	require.Equal(t, attempt, attemptResult.attempt)
 }

routing/router.go

@@ -1186,7 +1186,7 @@ func (r *ChannelRouter) sendToRoute(htlcHash lntypes.Hash, rt *route.Route,
 
 	// The attempt was successfully sent, wait for the result to be
 	// available.
-	result, err = p.collectResult(attempt)
+	result, err = p.collectAndHandleResult(attempt)
 	if err != nil {
 		return nil, err
 	}