fix(recovery): guard claim fan-out send on context cancellation to unblock Stop

The recovery manager fanned claims out to its worker pool with an unguarded
`work <- claim` send. On shutdown the workers return from their own select on
`ctx.Done()` without draining `work`, so that send blocked forever:
`close(work)` was never reached, the recovery loop's deferred `wg.Done()` never
ran, and `Stop()`'s `wg.Wait()` hung while holding `m.mu`, wedging every later
`Start()`/`Stop()` call.

Extract the fan-out into `Manager.fanOut` and guard the send on `ctx.Done()`,
keeping `close(work)` and `workerWG.Wait()` on the cancellation path so no
worker goroutine leaks, and surface the cancellation instead of swallowing it.
Undispatched claims stay `Pending`, so the next sweep re-claims them once their
lease expires and no work is lost.

`fanOut` also returns the number of claims it dispatched. The sweep summary
counted successes as `len(records)-failures`, which was only correct while the
fan-out either completed or hung; a cancelled fan-out leaves a tail that never
reaches `errCh`, so a partial sweep reported never-attempted claims as
succeeded. Successes are now counted against the dispatched total and a short
dispatch warns on its own.

A sweep aborted by `Stop()` is an ordinary shutdown, so `recoveryLoop` logs a
cancelled sweep at debug level via `logSweepError` and keeps the warning for
genuine failures.

Signed-off-by: AkramBitar <akram@il.ibm.com>

Author: AkramBitar

docs/services/storage/recovery.md

@@ -170,4 +170,8 @@ All three terminal statuses (`Confirmed`, `Deleted`, `Orphan`) are excluded from
 
 ## Thread Safety
 
-The Manager is thread-safe and can be safely started/stopped from multiple goroutines. The Handler implementation must also be thread-safe as it will be called concurrently by multiple workers.
+The Manager is thread-safe and can be safely started/stopped from multiple goroutines. The Handler implementation must also be thread-safe as it will be called concurrently by multiple workers.
+
+## Shutdown Behaviour
+
+`Stop()` cancels the manager context and waits for the recovery loop to return. If a sweep is mid-batch, the fan-out to the worker pool aborts on cancellation: workers stop after their in-flight `Handler.Recover()` call and any claims not yet dispatched are simply left undispatched. Those rows keep their `Pending` status, so they become eligible again once their lease (`leaseDuration`) expires and are picked up by the next sweep — on this or another replica. The aborted sweep reports a `recovery fan-out cancelled` error. Because this is an ordinary shutdown rather than a failure, the loop logs it at debug level and only warns for genuine sweep errors. The sweep summary counts successes against the claims actually dispatched, so a partial sweep logs `claimed=N, dispatched=M, ...` at warn level instead of crediting the undispatched tail as succeeded.

token/services/storage/services/recovery/manager.go

@@ -162,7 +162,7 @@ func (m *Manager) recoveryLoop() {
 	}
 
 	if err := m.runSweep(m.ctx); err != nil {
-		m.logger.Warnf("initial transaction recovery sweep failed: %v", err)
+		m.logSweepError("initial transaction recovery sweep", err)
 	}
 
 	for {
@@ -173,12 +173,28 @@ func (m *Manager) recoveryLoop() {
 			return
 		case <-ticker.C:
 			if err := m.runSweep(m.ctx); err != nil {
-				m.logger.Warnf("transaction recovery sweep failed: %v", err)
+				m.logSweepError("transaction recovery sweep", err)
 			}
 		}
 	}
 }
 
+// logSweepError reports a sweep that returned an error.
+//
+// A sweep aborted because Stop cancelled the manager context is an ordinary
+// shutdown rather than a failure, so it is logged at debug level: the fan-out
+// surfaces the cancellation as an error to keep the undispatched claims visible
+// to the caller, and warning about it would flag a healthy node shutdown.
+func (m *Manager) logSweepError(what string, err error) {
+	if errors.Is(err, context.Canceled) {
+		m.logger.Debugf("%s stopped: %v", what, err)
+
+		return
+	}
+
+	m.logger.Warnf("%s failed: %v", what, err)
+}
+
 func (m *Manager) validateConfig() error {
 	switch {
 	case m.config.TTL <= 0:
@@ -248,16 +264,7 @@ func (m *Manager) recoverTransactions(ctx context.Context) error {
 		go m.worker(ctx, &workerWG, work, errCh)
 	}
 
-	// ClaimPendingTransactions reads directly from the requests table where
-	// tx_id is the primary key, so each claim is already unique. Fan out
-	// straight to the workers; nil entries are defensive but should never
-	// occur in practice.
-	for _, claim := range records {
-		if claim == nil {
-			continue
-		}
-		work <- claim
-	}
+	dispatched, fanOutErr := m.fanOut(ctx, records, work)
 	close(work)
 
 	workerWG.Wait()
@@ -277,15 +284,59 @@ func (m *Manager) recoverTransactions(ctx context.Context) error {
 		}
 	}
 
-	if failures > 0 {
-		m.logger.Warnf("completed recovery sweep: claimed=%d, succeeded=%d, failed=%d", len(records), len(records)-failures, failures)
+	// Successes are counted against what the fan-out actually dispatched, not
+	// against len(records): a cancelled fan-out leaves the tail undispatched,
+	// and those claims never reach errCh. Attributing them to len(records)-failures
+	// would report never-attempted work as succeeded.
+	if failures > 0 || dispatched < len(records) {
+		m.logger.Warnf("completed recovery sweep: claimed=%d, dispatched=%d, succeeded=%d, failed=%d",
+			len(records), dispatched, dispatched-failures, failures)
 	} else {
 		m.logger.Debugf("completed recovery sweep: claimed=%d, all succeeded", len(records))
 	}
 
+	if fanOutErr != nil {
+		return errors.Join(fanOutErr, firstErr)
+	}
+
 	return firstErr
 }
 
+// fanOut dispatches the claimed records to the worker pool.
+//
+// ClaimPendingTransactions reads directly from the requests table where tx_id
+// is the primary key, so each claim is already unique. Fan out straight to the
+// workers; nil entries are defensive but should never occur in practice.
+//
+// The send is guarded on ctx.Done() because the workers return from their own
+// select as soon as the context is cancelled, without draining what is left in
+// work. An unguarded send would then block forever with no receiver, so
+// close(work) and the deferred wg.Done() of the recovery loop would never run
+// and Stop's wg.Wait() would hang while holding m.mu, wedging every later
+// Start/Stop call.
+//
+// It returns the number of claims actually handed to a worker. Callers need
+// this to report the sweep outcome: claims left undispatched by a cancellation
+// never reach errCh, so they cannot be inferred from the failure count.
+func (m *Manager) fanOut(ctx context.Context, records []*ttxdb.RecoveryClaim, work chan<- *ttxdb.RecoveryClaim) (int, error) {
+	dispatched := 0
+	for i, claim := range records {
+		if claim == nil {
+			continue
+		}
+		select {
+		case work <- claim:
+			dispatched++
+		case <-ctx.Done():
+			m.logger.Debugf("recovery fan-out cancelled: %d of %d claim(s) not dispatched", len(records)-i, len(records))
+
+			return dispatched, errors.Wrapf(ctx.Err(), "recovery fan-out cancelled")
+		}
+	}
+
+	return dispatched, nil
+}
+
 func (m *Manager) worker(ctx context.Context, wg *sync.WaitGroup, work <-chan *ttxdb.RecoveryClaim, errCh chan<- error) {
 	defer wg.Done()
 

token/services/storage/services/recovery/manager_test.go

@@ -7,7 +7,10 @@ SPDX-License-Identifier: Apache-2.0
 package recovery_test
 
 import (
+	"context"
 	"errors"
+	"strconv"
+	"sync"
 	"testing"
 	"time"
 
@@ -306,6 +309,105 @@ func TestManager_NoPromotionWhenGracePeriodDisabled(t *testing.T) {
 	assert.Equal(t, 0, mockDB.SetStatusCallCount(), "grace period disabled should never call SetStatus")
 }
 
+// TestManager_StopDuringFanOutDoesNotDeadlock walks the exact failure scenario
+// reported in issue #2038, step by step:
+//
+//  1. recoveryLoop's goroutine is mid-fan-out, with claims still queued to send
+//     to work — forced here by returning 64 claims to a single worker that is
+//     parked inside Recover, so the unbuffered send has no receiver.
+//  2. Stop() cancels m.ctx — called from a goroutine so the test can observe it
+//     hanging instead of hanging with it.
+//  3. All WorkerCount workers observe ctx.Done() in their select and exit before
+//     draining the remaining items from work — the worker is released after the
+//     cancellation, so both its select arms are ready and it takes ctx.Done().
+//  4. The fan-out loop's next `work <- claim` send now has no receiver and, when
+//     unguarded, blocks forever.
+//  5. That send sits in the goroutine wg.Done() is deferred on, so the deferred
+//     call never runs.
+//  6. Stop()'s m.wg.Wait() — held under m.mu — hangs indefinitely, and every
+//     later Start()/Stop() deadlocks on m.mu; asserted by the final restart.
+//
+// Verified to fail (10s timeout at step 6) against the unguarded send and to
+// pass with the ctx.Done()-guarded fan-out.
+func TestManager_StopDuringFanOutDoesNotDeadlock(t *testing.T) {
+	logger := logging.MustGetLogger()
+	mockDB := &mock2.Storage{}
+	mockHandler := &mock2.Handler{}
+	config := recovery2.Config{
+		Enabled:        true,
+		TTL:            10 * time.Millisecond,
+		ScanInterval:   10 * time.Millisecond,
+		BatchSize:      100,
+		WorkerCount:    1,
+		LeaseDuration:  time.Second,
+		AdvisoryLockID: 1,
+		InstanceID:     "test-instance",
+	}
+
+	// Step 1: many more claims than workers, so the fan-out loop is guaranteed to
+	// still have undispatched claims when the context is cancelled. 64 also makes
+	// it statistically certain (1 - 2^-63) that the worker's select picks
+	// ctx.Done() before draining the batch on its own.
+	const claimCount = 64
+	records := make([]*ttxdb.RecoveryClaim, claimCount)
+	for i := range records {
+		records[i] = &ttxdb.RecoveryClaim{TxID: "tx" + strconv.Itoa(i)}
+	}
+
+	leadership := &mock2.Leadership{}
+	leadership.CloseReturns(nil)
+	mockDB.AcquireRecoveryLeadershipReturns(leadership, true, nil)
+	mockDB.ClaimPendingTransactionsReturns(records, nil)
+	mockDB.ReleaseRecoveryClaimReturns(nil)
+
+	var once sync.Once
+	inWorker := make(chan struct{})
+	release := make(chan struct{})
+	mockHandler.RecoverStub = func(_ context.Context, _ string) error {
+		// Hold the single worker inside Recover so the fan-out loop is parked
+		// on its send, then let it go only after Stop() cancelled the context.
+		once.Do(func() {
+			close(inWorker)
+			<-release
+		})
+
+		return nil
+	}
+
+	manager := recovery2.NewManager(logger, mockDB, mockHandler, config)
+	require.NoError(t, manager.Start())
+
+	select {
+	case <-inWorker:
+	case <-time.After(10 * time.Second):
+		t.Fatal("timed out waiting for the recovery sweep to reach a worker")
+	}
+
+	// Step 2: cancel m.ctx via Stop() while the fan-out is parked on its send.
+	stopped := make(chan error, 1)
+	go func() {
+		stopped <- manager.Stop()
+	}()
+
+	// Step 3: let Stop() cancel the context first, then unblock the worker so it
+	// returns to its select with both ctx.Done() and the fan-out send ready.
+	time.Sleep(50 * time.Millisecond)
+	close(release)
+
+	// Steps 4-5-6: with the unguarded send this never returns.
+	select {
+	case err := <-stopped:
+		require.NoError(t, err)
+	case <-time.After(10 * time.Second):
+		t.Fatal("Stop() deadlocked while the sweep was fanning out claims")
+	}
+
+	// Step 6, second half: the manager must not be wedged — m.mu was released, so
+	// a subsequent Start()/Stop() pair still works.
+	require.NoError(t, manager.Start())
+	require.NoError(t, manager.Stop())
+}
+
 func TestDefaultConfig(t *testing.T) {
 	config := recovery2.DefaultConfig()