test: stabilize asynchronous cleanup waits (#26256)

## What type of PR is this?

- [x] BUG
- [x] Test and CI

## Which issue(s) this PR fixes:

Fixes #26254
Fixes #26255

## What this PR does / why we need it:

Two unrelated asynchronous subsystems were made flaky by test-only
timing assumptions under the parallel Ubuntu/x86 UT workload.

### Root cause

- TAE checkpoint tests wrapped real flush/checkpoint I/O in 5, 10, or 20
second deadlines. A single storage `writer.Sync` can legitimately exceed
20 seconds under parallel CI, so `TestSnapshotGC` could cancel
`ForceFlush` before the asynchronous work completed.
- The user-level-lock test service latched `blockUnlock=true` when
`Unlock` entered. Clearing the flag did not release calls already in
flight; all four detached-cleanup workers remained occupied until their
one-second contexts expired. A saturated 1,024-entry test queue could
therefore miss its three-second recovery assertion.

### Changes

- Add a shared one-minute checkpoint/flush timeout for storage-backed
TAE tests and apply it to the short checkpoint helpers and nearby direct
call sites.
- Make the user-level-lock test double observe recovery while an unlock
is in flight.
- Add a deterministic regression proving an already-blocked mock unlock
resumes promptly.

Production checkpoint and user-lock behavior is unchanged.

## Validation

-
`TestUserLevelLockCleanupTestServiceUnblocksInFlightUnlock|TestDetachedUserLevelLockCleanupQueueIsBoundedAndDeduped`:
`-race -count=20`
- `TestSnapshotGC|TestAppendAndGC2|TestCkpLeak|TestReplay2`: `-count=5`
- Same TAE set: `-race -count=1`
- Full `pkg/sql/plan/function`, `pkg/vm/engine/tae/db/test`, and
`pkg/vm/engine/tae/db/testutil` tests
- `go build` for affected production/test-support packages
- `go vet` for affected test packages

Author: XuPeng-SH

pkg/sql/plan/function/func_unary_test.go

@@ -7441,7 +7441,8 @@ func (s *userLevelLockTestService) Unlock(ctx context.Context, txnID []byte, com
 	s.unlockMu.Lock()
 	s.unlockedTxnIDs = append(s.unlockedTxnIDs, append([]byte(nil), txnID...))
 	s.unlockMu.Unlock()
-	if s.blockUnlock.Load() {
+waitForUnlock:
+	for s.blockUnlock.Load() {
 		if s.unlockStarted != nil {
 			select {
 			case s.unlockStarted <- struct{}{}:
@@ -7451,12 +7452,23 @@ func (s *userLevelLockTestService) Unlock(ctx context.Context, txnID []byte, com
 		if s.unlockResume != nil {
 			select {
 			case <-s.unlockResume:
+				break waitForUnlock
 			case <-ctx.Done():
 				return ctx.Err()
 			}
 		} else {
-			<-ctx.Done()
-			return ctx.Err()
+			timer := time.NewTimer(time.Millisecond)
+			select {
+			case <-timer.C:
+			case <-ctx.Done():
+				if !timer.Stop() {
+					select {
+					case <-timer.C:
+					default:
+					}
+				}
+				return ctx.Err()
+			}
 		}
 	}
 	if s.unlockErr != nil {
@@ -7532,6 +7544,35 @@ func runUserLevelLockTest(t *testing.T, fn func([]lockservice.LockService)) {
 	resetUserLevelLocksForTest(t)
 }
 
+func TestUserLevelLockCleanupTestServiceUnblocksInFlightUnlock(t *testing.T) {
+	service := &userLevelLockTestService{
+		id:            "user-level-lock-unblock",
+		state:         &userLevelLockTestState{locks: make(map[string]string)},
+		unlockStarted: make(chan struct{}, 1),
+	}
+	service.blockUnlock.Store(true)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	done := make(chan error, 1)
+	go func() {
+		done <- service.Unlock(ctx, []byte("txn"), timestamp.Timestamp{})
+	}()
+
+	select {
+	case <-service.unlockStarted:
+	case <-time.After(time.Second):
+		t.Fatal("unlock did not reach the blocked state")
+	}
+	service.blockUnlock.Store(false)
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(500 * time.Millisecond):
+		t.Fatal("in-flight unlock did not observe the test service recovery")
+	}
+}
+
 func requireUserLevelLockTxnUnlocked(t *testing.T, service *userLevelLockTestService, txnID []byte) {
 	t.Helper()
 	requireUserLevelLockTxnUnlockedFunc(t, service, func(unlocked []byte) bool {

pkg/vm/engine/tae/db/test/db_test.go

@@ -7609,7 +7609,7 @@ func TestAppendAndGC2(t *testing.T) {
 		assert.Nil(t, err)
 	}
 	wg.Wait()
-	ckpCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	ckpCtx, cancel := context.WithTimeout(ctx, testutil.TestCheckpointTimeout)
 	err = db.ForceCheckpoint(ckpCtx, db.TxnMgr.Now())
 	cancel()
 	require.NoError(t, err)
@@ -7836,7 +7836,7 @@ func TestSnapshotGC(t *testing.T) {
 	snapWG.Wait()
 	wg.Wait()
 	t.Log(tae.Catalog.SimplePPString(common.PPL1))
-	ckpCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	ckpCtx, cancel := context.WithTimeout(ctx, testutil.TestCheckpointTimeout)
 	err = db.ForceCheckpoint(ckpCtx, db.TxnMgr.Now())
 	cancel()
 	require.NoError(t, err)
@@ -8847,7 +8847,7 @@ func TestCkpLeak(t *testing.T) {
 		assert.Nil(t, err)
 	}
 	wg.Wait()
-	ckpCtx, cancel := context.WithTimeout(ctx, time.Minute)
+	ckpCtx, cancel := context.WithTimeout(ctx, testutil.TestCheckpointTimeout)
 	defer cancel()
 	require.NoError(t, db.ForceCheckpoint(ckpCtx, db.TxnMgr.Now()))
 	testutil.WaitAllCheckpointsFinished(t, db)

pkg/vm/engine/tae/db/test/replay_test.go

@@ -463,7 +463,7 @@ func TestReplay2(t *testing.T) {
 	assert.NotNil(t, err)
 	assert.Nil(t, txn.Commit(context.Background()))
 
-	ctx2, cancel := context.WithTimeout(context.Background(), time.Second*10)
+	ctx2, cancel := context.WithTimeout(context.Background(), testutil.TestCheckpointTimeout)
 	defer cancel()
 	err = tae2.ForceFlush(ctx2, tae2.TxnMgr.Now())
 	assert.NoError(t, err)

pkg/vm/engine/tae/db/testutil/engine.go

@@ -41,9 +41,11 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-const (
-	DefaultTestDB = "db"
-)
+const DefaultTestDB = "db"
+
+// TestCheckpointTimeout leaves room for storage syncs under parallel CI while
+// still bounding a genuinely stuck checkpoint or flush.
+const TestCheckpointTimeout = time.Minute
 
 type CtxOldVersion struct{}
 
@@ -160,14 +162,14 @@ func (e *TestEngine) CheckRowsByScan(exp int, applyDelete bool) {
 	assert.NoError(e.T, txn.Commit(context.Background()))
 }
 func (e *TestEngine) ForceCheckpoint() {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	ctx, cancel := context.WithTimeout(context.Background(), TestCheckpointTimeout)
 	defer cancel()
 	err := e.DB.ForceCheckpoint(ctx, e.TxnMgr.Now())
 	assert.NoError(e.T, err)
 }
 
 func (e *TestEngine) ForceLongCheckpoint() {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second*20)
+	ctx, cancel := context.WithTimeout(context.Background(), TestCheckpointTimeout)
 	defer cancel()
 	err := e.DB.ForceCheckpoint(ctx, e.TxnMgr.Now())
 	assert.NoError(e.T, err)