fix: address review on the reclaim re-validation and the tests

Re-validation now matches the collector's delete phase: transaction and last
access, not transaction alone. Comparing the txID only left an ABA gap, since
the entry can be unlocked and re-locked under the same txID while the shard
lock is released, and reclaiming on that stale verdict would drop a fresh
entry. The last access is read before the refresh that would clobber it.

The failure log no longer pairs the current holder with the previously
observed holder's status. When the entry changed hands it says so and names
the transaction the status actually belongs to.

TestReclaimDoesNotHoldShardLockDuringStatusLookup stops the collector before
arming the hook. A long sleep timeout was not enough: scan runs a full pass
before its first sleep, so it could consume the one-shot and park there
itself, leaving the reclaim's lookup to pass through and the test passing
without ever holding a lookup open inside Lock, including on the unfixed code.

Both test bounds drop under stopTimeout so a failing run lets Stop join the
scan goroutine rather than leaving it parked past the end of the test, and the
const block moves above the doc comment it had split from its function. The
note about the hook blocking under the shard lock is no longer true after this
PR and now describes the collector's lookup phase.

Signed-off-by: atharrva01 <atharvaborade568@gmail.com>

Author: atharrva01

token/services/selector/simple/inmemory/locker.go

@@ -179,15 +179,16 @@ func (d *locker) lockInShard(ctx context.Context, s *shard, owner string, id *to
 	// The holder observed here is re-validated under the lock below, since it may
 	// change while the lock is not held.
 	var (
-		observedTxID   string
-		observedStatus int
-		statusResolved bool
+		observedTxID       string
+		observedLastAccess time.Time
+		observedStatus     int
+		statusResolved     bool
 	)
 	if reclaim {
 		s.mu.RLock()
 		e, held := s.locked[k]
 		if held {
-			observedTxID = e.TxID
+			observedTxID, observedLastAccess = e.TxID, e.LastAccess
 		}
 		s.mu.RUnlock()
 
@@ -209,30 +210,41 @@ func (d *locker) lockInShard(ctx context.Context, s *shard, owner string, id *to
 	}
 	e, ok := s.locked[k]
 	if ok {
+		// Read before the refresh below clobbers it: the re-validation compares against
+		// the value observed during the status lookup.
+		prevAccess := e.LastAccess
 		e.LastAccess = time.Now()
 
 		if reclaim {
-			// Second chance. Only act on the status resolved above if the entry is
-			// still the one it was resolved for: a concurrent Lock(reclaim=true) may
-			// have handed the token to another transaction in the meantime, and
-			// reclaiming on a stale status would take a token away from its new
-			// holder. When it no longer matches, report the token as locked and let
-			// the caller retry.
+			// Second chance. Only act on the status resolved above if the entry is still
+			// exactly the one it was resolved for, matching the collector's delete phase:
+			// same transaction and same last access. Comparing the transaction alone is
+			// not enough, since the entry may have been unlocked and re-locked under the
+			// same txID while the shard lock was released, and reclaiming on that stale
+			// verdict would drop a fresh entry. When it no longer matches, report the
+			// token as locked and let the caller retry.
 			logger.DebugfContext(ctx, "[%s] already locked by [%s], try to reclaim...", id, e)
-			status := observedStatus
-			reclaimed := statusResolved && e.TxID == observedTxID && status == ttxdb.Deleted
+			unchanged := statusResolved && e.TxID == observedTxID && prevAccess.Equal(observedLastAccess)
+			reclaimed := unchanged && observedStatus == ttxdb.Deleted
 			if reclaimed {
 				delete(s.locked, k)
 			}
 			if !reclaimed {
-				logger.DebugfContext(ctx, "[%s] already locked by [%s], reclaim failed, tx status [%s]", id, e, ttxdb.TxStatusMessage[status])
+				// Only report the status when it belongs to the holder still in place;
+				// otherwise it describes observedTxID, not e, and pairing the two sends
+				// a reader after the wrong transaction.
+				if unchanged {
+					logger.DebugfContext(ctx, "[%s] already locked by [%s], reclaim failed, tx status [%s]", id, e, ttxdb.TxStatusMessage[observedStatus])
+				} else {
+					logger.DebugfContext(ctx, "[%s] already locked by [%s], reclaim failed, entry changed since the status of [%s] was read", id, e, observedTxID)
+				}
 				if logger.IsEnabledFor(zapcore.DebugLevel) {
 					return e.TxID, errors.Errorf("already locked by [%s]", e)
 				}
 
 				return e.TxID, AlreadyLockedError
 			}
-			logger.DebugfContext(ctx, "[%s] already locked by [%s], reclaimed successful, tx status [%s]", id, e, ttxdb.TxStatusMessage[status])
+			logger.DebugfContext(ctx, "[%s] reclaimed from [%s], tx status [%s]", id, observedTxID, ttxdb.TxStatusMessage[observedStatus])
 		} else {
 			logger.DebugfContext(ctx, "[%s] already locked by [%s], no reclaim", id, e)
 			if logger.IsEnabledFor(zapcore.DebugLevel) {

token/services/selector/simple/inmemory/locker_test.go

@@ -89,20 +89,22 @@ func (m *mockTXStatusProvider) GetStatus(_ context.Context, txID string) (ttxdb.
 	return m.status(txID), "", nil
 }
 
+// Bounds for the coordination in TestScannerDoesNotDeleteReclaimed. They are far longer than the
+// 20ms scan interval the test drives, so they only fire when something is genuinely wrong, but they
+// keep a missed interleaving from turning into a 10 minute package timeout. Both stay under
+// stopTimeout so that a failing run still lets t.Cleanup's Stop join the scan goroutine instead of
+// leaving it parked in the hook and running its delete phase during later tests. See #2156.
+const (
+	scannerObserveTimeout = 3 * time.Second
+	hookReleaseTimeout    = 3 * time.Second
+)
+
 // TestScannerDoesNotDeleteReclaimed verifies the TOCTOU protection in the
 // scanner: when the scanner has observed a token as removable (its tx is
 // Deleted) and a concurrent Lock(reclaim=true) re-locks that token for a new
 // transaction before the scanner deletes, the scanner must NOT delete the
 // new entry. The test drives the real scan loop and blocks the scanner's
 // status lookup to open the race window deterministically.
-// Bounds for the coordination in TestScannerDoesNotDeleteReclaimed. Both are far longer than the
-// 20ms scan interval the test drives, so they only fire when something is genuinely wrong, but they
-// keep a missed interleaving from turning into a 10 minute package timeout. See #2156.
-const (
-	scannerObserveTimeout = 30 * time.Second
-	hookReleaseTimeout    = 30 * time.Second
-)
-
 func TestScannerDoesNotDeleteReclaimed(t *testing.T) {
 	mock := newMockTXStatusProvider()
 	tokenID := &token.ID{TxId: "tok1", Index: 0}
@@ -133,9 +135,9 @@ func TestScannerDoesNotDeleteReclaimed(t *testing.T) {
 			// (the reclaim's) must pass through or they would deadlock
 			close(entered)
 			// Bounded: if the test fails before reaching close(release), this
-			// goroutine must not sit here forever. Note reclaim() calls
-			// GetStatus while holding the shard lock, so a blocked hook can
-			// hold that lock too.
+			// goroutine must not sit here forever. This blocks the collector's
+			// lookup phase, which deliberately runs without the shard lock, so
+			// parking here does not wedge the shard.
 			select {
 			case <-release:
 			case <-time.After(hookReleaseTimeout):
@@ -207,18 +209,27 @@ func TestReclaimDoesNotHoldShardLockDuringStatusLookup(t *testing.T) {
 	txA := "tx-A"
 
 	mock.setStatus(txA, ttxdb.Pending)
-	// A long sleep timeout keeps the collector out of the way: this test is about
-	// the locking path, not the scanner.
 	d := NewLocker(mock, time.Hour, time.Hour).(*locker)
 	t.Cleanup(func() { _ = d.Stop() })
 	_, err := d.Lock(context.Background(), "w1", tokenID, txA, false)
 	require.NoError(t, err)
 
+	// Stop the collector before arming the hook. A long sleep timeout is not enough
+	// to keep it away: scan runs a full pass before its first sleep, so it can reach
+	// GetStatus while the hook is armed, consume the one-shot below and park there
+	// itself. The reclaim's own lookup would then pass straight through and the test
+	// would pass without ever holding a lookup open inside Lock, which it does even
+	// on the unfixed code. This test is about the locking path, so take the collector
+	// out of the picture entirely.
+	require.NoError(t, d.Stop())
+
 	// Block the next status lookup for tx-A, standing in for a slow transaction store.
 	inLookup := make(chan struct{})
 	unblock := make(chan struct{})
 	var once sync.Once
 	mock.setGetStatusHook(func(txID string) {
+		// Only the reclaim looks up tx-A now that the collector is stopped, so the
+		// one-shot is belt and braces rather than load bearing.
 		if txID != txA {
 			return
 		}