fix: prevent nested local shuffle deadlocks (#26055)

## What type of PR is this?

- [x] BUG
- [ ] Improvement
- [ ] Documentation
- [ ] Feature
- [ ] Test and CI
- [ ] Code Refactoring

## Which issue(s) this PR fixes

Fixes #25922

## What this PR does / why we need it

The unified Shuffle implementation added bounded backpressure to shared
local Shuffle pools. Per-bucket credits fixed same-level hot-bucket
starvation, but a finite fixed-bucket pool can still deadlock when more
than one local Shuffle is duplicated into the same pull-based worker
tree.

The observed 16-worker wait-for cycle was:

1. an outer Shuffle writer filled its target bucket and waited for that
bucket's worker;
2. the target worker was blocked in an inner Shuffle waiting for every
inner writer to finish;
3. the missing inner writer was the child of the blocked outer writer,
so it could only finish if the outer writer resumed pulling.

This change prevents the cyclic topology instead of increasing or
removing the memory bound:

- keep the first fixed-bucket local Shuffle fast path;
- detect a fixed-bucket Shuffle already present in a packed operator
tree;
- materialize a second join/group Shuffle through the existing producer
scopes, bounded dispatch channels, and consumer scopes;
- preserve local Shuffle reuse when no second probe Shuffle is
introduced;
- cancel an active request before connection cleanup waits for
migration/reset or transaction rollback, so cleanup is independent of a
stuck data path.

The fallback remains bounded: its existing dispatch channel capacity is
32, and this PR does not change Shuffle ready credits or introduce an
unbounded queue. It also does not introduce a new hard query timeout.

## Regression coverage

- 16-way nested local shuffle join, on both probe and build chains;
- 16-way nested local shuffle group;
- fixed Shuffle below a non-root operator to cover recursive detection;
- single-layer local group and reusable local join fast paths remain
local;
- connection cleanup cancels before a blocked lifecycle operation and
before rollback.

The nested-join regression failed on `origin/main` by producing one
packed local scope and passes with this change by producing 16
independent bucket scopes.

## Validation

- `go test -race -count=1 -timeout=600s ./pkg/sql/compile
./pkg/frontend`
- focused new concurrency/topology tests under `-race -count=20`
- `go test -v -count=1 -timeout=600s ./pkg/sql/compile ./pkg/frontend`
- `go test -count=1 ./pkg/sql/colexec/shuffle ./pkg/sql/colexec/dispatch
./pkg/sql/colexec/merge ./pkg/sql/colexec/group`
- `go build ./pkg/sql/compile ./pkg/frontend`
- `go vet ./pkg/sql/compile ./pkg/frontend`
- `git diff --check`

---------

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: XuPeng-SH

pkg/cnservice/server_query.go

@@ -531,7 +531,7 @@ func (s *service) handleMigrateConnFrom(
 	}
 	rm := s.mo.GetRoutineManager()
 	resp.MigrateConnFromResponse = &query.MigrateConnFromResponse{}
-	if err := rm.MigrateConnectionFrom(req.MigrateConnFromRequest, resp.MigrateConnFromResponse); err != nil {
+	if err := rm.MigrateConnectionFromWithContext(ctx, req.MigrateConnFromRequest, resp.MigrateConnFromResponse); err != nil {
 		logutil.Errorf("failed to migrate conn from: %v", err)
 		return err
 	}
@@ -592,7 +592,7 @@ func (s *service) handleResetSession(
 	}
 	rm := s.mo.GetRoutineManager()
 	resp.ResetSessionResponse = &query.ResetSessionResponse{}
-	if err := rm.ResetSession(req.ResetSessionRequest, resp.ResetSessionResponse); err != nil {
+	if err := rm.ResetSessionWithContext(ctx, req.ResetSessionRequest, resp.ResetSessionResponse); err != nil {
 		logutil.Errorf("failed to reset session: %v", err)
 		return err
 	}

pkg/frontend/migrate.go

@@ -14,17 +14,25 @@
 
 package frontend
 
-import "sync"
+import (
+	"context"
+	"sync"
+)
 
 // migrateController is created in Routine and used to serialize lifecycle
 // operations with routine cleanup.
 type migrateController struct {
 	sync.Mutex
 	migrateOnce sync.Once
+	// migrateErr is the stable result returned to duplicate migration calls.
+	migrateErr error
 	// closed indicates if the session has been closed.
 	closed bool
 	// inProgress indicates if a lifecycle operation is in progress.
 	inProgress bool
+	// operationCancel cancels the active lifecycle operation. It is published
+	// together with inProgress while holding the controller lock.
+	operationCancel context.CancelFunc
 	// cond coordinates lifecycle operations and routine cleanup.
 	cond *sync.Cond
 	// the id of goroutine that executes the migration
@@ -41,50 +49,96 @@ func newMigrateController() *migrateController {
 // operation is in progress, it waits for that operation to finish and marks
 // the routine closed.
 func (mc *migrateController) waitAndClose() {
+	mc.startClose()
 	mc.Lock()
 	defer mc.Unlock()
-	mc.closed = true
 	for mc.inProgress {
 		mc.cond.Wait()
 	}
 }
 
+// startClose seals lifecycle admission and cancels the active operation. It
+// does not wait, so connection-liveness control paths can remain independent
+// of the operation they are stopping.
+func (mc *migrateController) startClose() {
+	mc.Lock()
+	mc.closed = true
+	cancel := mc.operationCancel
+	mc.cond.Broadcast()
+	mc.Unlock()
+	if cancel != nil {
+		cancel()
+	}
+}
+
 // beginOperation starts a lifecycle operation unless the routine is closed or
 // another operation is already running.
 func (mc *migrateController) beginOperation() bool {
+	_, ok := mc.beginOperationWithContext(context.Background())
+	return ok
+}
+
+// beginOperationWithContext starts a cancelable lifecycle operation. A waiting
+// caller is released by operation completion, routine close, or its own
+// context cancellation.
+func (mc *migrateController) beginOperationWithContext(ctx context.Context) (context.Context, bool) {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	stopWakeup := context.AfterFunc(ctx, func() {
+		mc.Lock()
+		mc.cond.Broadcast()
+		mc.Unlock()
+	})
+	defer stopWakeup()
+
 	mc.Lock()
 	defer mc.Unlock()
-	for mc.inProgress && !mc.closed {
+	for mc.inProgress && !mc.closed && ctx.Err() == nil {
 		mc.cond.Wait()
 	}
-	return mc.startOperationLocked()
+	return mc.startOperationLocked(ctx)
 }
 
 // tryBeginOperation starts a lifecycle operation only if it can proceed
 // immediately.
 func (mc *migrateController) tryBeginOperation() bool {
+	_, ok := mc.tryBeginOperationWithContext(context.Background())
+	return ok
+}
+
+func (mc *migrateController) tryBeginOperationWithContext(ctx context.Context) (context.Context, bool) {
+	if ctx == nil {
+		ctx = context.Background()
+	}
 	mc.Lock()
 	defer mc.Unlock()
-	return mc.startOperationLocked()
+	return mc.startOperationLocked(ctx)
 }
 
-func (mc *migrateController) startOperationLocked() bool {
-	if mc.closed {
-		return false
+func (mc *migrateController) startOperationLocked(ctx context.Context) (context.Context, bool) {
+	if mc.closed || ctx.Err() != nil {
+		return nil, false
 	}
 	if mc.inProgress {
-		return false
+		return nil, false
 	}
+	operationCtx, cancel := context.WithCancel(ctx)
 	mc.goroutineID = GetRoutineId()
 	mc.inProgress = true
-	return true
+	mc.operationCancel = cancel
+	return operationCtx, true
 }
 
 // endOperation completes a lifecycle operation and wakes a routine waiting
 // to close.
 func (mc *migrateController) endOperation() {
 	mc.Lock()
 	defer mc.Unlock()
+	if mc.operationCancel != nil {
+		mc.operationCancel()
+		mc.operationCancel = nil
+	}
 	mc.inProgress = false
 	mc.goroutineID = 0
 	mc.cond.Broadcast()

pkg/frontend/migrate_test.go

@@ -15,6 +15,7 @@
 package frontend
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -107,6 +108,68 @@ func TestLifecycleControllerCloseRejectsQueuedOperation(t *testing.T) {
 	}
 }
 
+func TestLifecycleControllerCloseCancelsActiveOperation(t *testing.T) {
+	mc := newMigrateController()
+	operationCtx, ok := mc.beginOperationWithContext(context.Background())
+	assert.True(t, ok)
+
+	operationDone := make(chan struct{})
+	go func() {
+		<-operationCtx.Done()
+		mc.endOperation()
+		close(operationDone)
+	}()
+
+	closeDone := make(chan struct{})
+	go func() {
+		mc.waitAndClose()
+		close(closeDone)
+	}()
+
+	select {
+	case <-operationCtx.Done():
+	case <-time.After(time.Second):
+		t.Fatal("close did not cancel the active lifecycle operation")
+	}
+	select {
+	case <-operationDone:
+	case <-time.After(time.Second):
+		t.Fatal("active lifecycle operation did not finish after cancellation")
+	}
+	select {
+	case <-closeDone:
+	case <-time.After(time.Second):
+		t.Fatal("close did not finish after the active lifecycle operation exited")
+	}
+	assert.False(t, mc.beginOperation())
+}
+
+func TestLifecycleControllerWaitingCallerHonorsContext(t *testing.T) {
+	mc := newMigrateController()
+	assert.True(t, mc.beginOperation())
+
+	ctx, cancel := context.WithCancel(context.Background())
+	waiterReady := make(chan struct{})
+	result := make(chan bool)
+	go func() {
+		close(waiterReady)
+		_, ok := mc.beginOperationWithContext(ctx)
+		result <- ok
+	}()
+	<-waiterReady
+	cancel()
+
+	select {
+	case ok := <-result:
+		assert.False(t, ok)
+	case <-time.After(time.Second):
+		t.Fatal("waiting lifecycle operation ignored caller cancellation")
+	}
+
+	mc.endOperation()
+	mc.waitAndClose()
+}
+
 func TestCloseWaitMigrate(t *testing.T) {
 	mc := newMigrateController()
 	assert.NotNil(t, mc)