util/async: cancellation race can block RunLoop.Append and batchRecvLoop indefinitely

[D3Hunter]

Summary

There is a cancellation race between RunLoop.Exec(ctx) and RunLoop.Append() that can leave Append blocked forever while sending on the unbuffered ready channel.

When Append is called by async.Callback.Schedule from batchCommandsClient.batchRecvLoop, the receive loop cannot finish dispatching or retiring that response. Later responses on the same batch-command stream are then head-of-line blocked even though the gRPC connection and server are healthy.

This is reproducible on current master at 61ecd7c114166c08e747d252fb9ae592e54bb599.

Root cause

RunLoop.Append does the following:

1. Acquire RunLoop.lock.
2. Append the callback.
3. If the state is StateWaiting, change it to StateIdle and set notify=true.
4. Release the lock.
5. Send on the unbuffered ready channel.

At the same time, RunLoop.Exec waits for either ready or ctx.Done().

The problematic interleaving is:

1. Exec(ctx) enters StateWaiting.
2. The context is canceled while Append is starting.
3. Append observes StateWaiting, changes the state to StateIdle, and decides to notify Exec.
4. Exec selects ctx.Done(), resets the state, and returns.
5. Append executes l.ready <- struct{}{} after the only receiver has exited.
6. The unbuffered send blocks permanently.

Callback.Schedule calls Append from inside sync.Once.Do, so the sync.Once operation also remains in progress.

For an asynchronous batch request, batchCommandsEntry.response calls Callback.Schedule. If that blocks, batchRecvLoop cannot delete/retire the request or continue processing later responses on that stream.

One observed blocked stack had this shape:

batchCommandsClient.batchRecvLoop
  -> batchCommandsEntry.response
  -> callback.Schedule
  -> sync.Once.Do
  -> RunLoop.Append
  -> chan send

Other goroutines attempting to complete the same callback can block on the sync.Once mutex. Request counters then show a growing difference between tracked and retired requests while the gRPC connections remain READY.

When EnableAsyncBatchGet is enabled, the multi-region batch-get path can reach this code through:

KVSnapshot.BatchGetWithTier
  -> batchGetKeysByRegions
  -> asyncBatchGetByRegions
  -> SendRequestAsync
  -> batchRecvLoop
  -> Callback.Schedule
  -> RunLoop.Append

Minimal reproduction

The following test uses the package lock only to force the otherwise timing-dependent interleaving. It detects the bug when Append remains blocked after the canceled Exec has returned.

Reproduction test for util/async/runloop_test.go

func TestCanceledExecCanStrandAppend(t *testing.T) {
	oldProcs := runtime.GOMAXPROCS(1)
	defer runtime.GOMAXPROCS(oldProcs)

	for attempt := 0; attempt < 100; attempt++ {
		loop := NewRunLoop()
		ctx, cancel := context.WithCancel(context.Background())
		execDone := make(chan error, 1)
		go func() {
			_, err := loop.Exec(ctx)
			execDone <- err
		}()

		deadline := time.Now().Add(time.Second)
		for loop.State() != StateWaiting {
			if time.Now().After(deadline) {
				t.Fatal("run loop did not enter waiting state")
			}
			runtime.Gosched()
		}

		// Queue Append on the mutex before waking Exec through cancellation.
		loop.lock.Lock()
		appendDone := make(chan struct{})
		go func() {
			loop.Append(func() {})
			close(appendDone)
		}()
		runtime.Gosched()
		cancel()
		runtime.Gosched()
		loop.lock.Unlock()

		select {
		case err := <-execDone:
			if err == nil {
				t.Fatal("expected canceled Exec to return an error")
			}
		case <-time.After(time.Second):
			t.Fatal("Exec did not return after cancellation")
		}

		select {
		case <-appendDone:
			// Mutex wake-up ordering may select Exec first; retry.
		case <-time.After(20 * time.Millisecond):
			t.Logf("reproduced stranded RunLoop.Append on attempt %d", attempt+1)
			return
		}
	}

	t.Fatal("did not reproduce the cancellation/notification interleaving")
}

Required imports:

import (
	"context"
	"runtime"
	"testing"
	"time"
)

Run:

go test ./util/async -run '^TestCanceledExecCanStrandAppend$' -count=20 -v

On darwin/arm64 with Go 1.25.10, all 20 runs reproduced the stranded Append on the first attempt:

reproduced stranded RunLoop.Append on attempt 1
PASS

This is a positive reproducer: it passes when the blocking bug is observed. A permanent regression test should invert the final assertion and require Append to complete.

Expected behavior

Cancellation of the RunLoop executor must not cause a concurrent or late callback submission to block indefinitely. A late response should be safely executed, rejected, or discarded without stopping the batch receive loop.

Requests following a canceled asynchronous request on the same batch-command stream should continue to receive and retire their responses.

Actual behavior

RunLoop.Append can block forever on the unbuffered wake-up send. When this happens in batchRecvLoop, one callback causes stream-wide head-of-line blocking:

• the response being dispatched is not retired;
• later responses on that stream are not processed;
• tracked-but-not-retired requests accumulate;
• unrelated client requests sharing the stream can hang or time out.

Suggested regression coverage

In addition to the util/async interleaving test, an internal/client regression test could:

1. Start a mock batch-command server.
2. Issue an asynchronous request using a RunLoop callback.
3. Cancel its context concurrently with the server response.
4. Send an unrelated synchronous request through the same batch stream.
5. Verify that the second request completes and both request entries are retired.

[D3Hunter]

/label bug
/label severity/major

[ti-chi-bot]

@D3Hunter: The label(s) /label bug, /label severity/major cannot be applied. These labels are supported: ``. Is this label configured under labels -> additional_labels or `labels -> restricted_labels` in `plugin.yaml`?

Details

In response to this:

/label bug
/label severity/major

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository.

[D3Hunter]

Suggested labels: bug, severity/major.

[D3Hunter]

Supplemental: how a canceled SQL reaches this race from TiDB

I traced the SQL-level path in TiDB at bd29ad1ac3fd. This TiDB revision pins the affected client-go revision 61ecd7c11416 in go.mod, and performance.enable-async-batch-get is enabled by default.

SQL shape

A straightforward SQL entry point is a primary-key or unique-key IN lookup whose plan is Batch_Point_Get:

CREATE TABLE t (
    id BIGINT PRIMARY KEY,
    u  BIGINT UNIQUE,
    v  VARCHAR(100)
);

-- On a real TiKV cluster, make the requested record keys span Regions.
SPLIT TABLE t BY (1000), (2000), (3000);

EXPLAIN SELECT * FROM t WHERE id IN (1, 1001, 2001, 3001);
-- Batch_Point_Get

SELECT CONNECTION_ID();
SELECT * FROM t WHERE id IN (1, 1001, 2001, 3001);

While the last query has outstanding TiKV requests, another session can cancel it with:

KILL TIDB QUERY <connection_id>;

The same cancellation can come from MAX_EXECUTION_TIME:

SELECT /*+ MAX_EXECUTION_TIME(50) */ *
FROM t
WHERE id IN (1, 1001, 2001, 3001);

This SQL is only a path demonstrator, not a deterministic race reproducer. Reproduction needs the cancellation to overlap an async completion, so delayed TiKV responses, injected network latency, or repeated concurrent execution is normally needed.

TiDB execution and cancellation paths

The read path is:

COM_QUERY
  -> clientConn.dispatch / handleQuery / handleStmt
  -> session.ExecuteStmt
  -> ExecStmt.Exec
  -> writeResultSet / ResultSet.Next(ctx)
  -> BatchPointGetExec.Next
  -> BatchPointGetExec.initialize
  -> batchGetter.BatchGet(ctx, keys)
  -> tikvSnapshot.BatchGet
  -> client-go KVSnapshot.BatchGetWithTier
  -> batchGetKeysByRegions
  -> asyncBatchGetByRegions

BatchPointGetExec passes the statement context into both its index-key BatchGet and record-key BatchGet (TiDB source, first BatchGet, record BatchGet). For MAX_EXECUTION_TIME, initialize directly derives a context.WithTimeout around these reads.

For KILL TIDB QUERY, TiDB reaches the same ctx.Done() through:

SimpleExec.executeKillStmt
  -> Server.Kill
  -> killQuery
  -> SQLKiller.SendKillSignal(QueryInterrupted)
  -> clientConn.cancelDispatch
  -> dispatch context.CancelFunc()

Relevant sources: KILL execution, killQuery, and cancelDispatch. Therefore this is an actual cancellation of the Go context used by the async BatchGet, not only an SQLKiller flag checked between executor operations.

Conditions required to enter the affected async path

batchGetKeysByRegions first groups keys by Region and splits large groups into batches. It uses batchGetSingleRegion when there is only one batch, and calls asyncBatchGetByRegions only when len(batches) > 1 and EnableAsyncBatchGet is true (client-go source). In practice, this means:

• Point_Get such as WHERE id = 1 does not reach this path.
• A small IN lookup contained in one Region normally stays on the synchronous path.
• A Batch_Point_Get spanning multiple Regions reaches the async path.
• More than 5,120 keys in one Region can also form multiple batches.
• A unique-index Batch Point Get can have two opportunities because it first fetches index keys and then record keys.

SQL-level race timeline and impact

1. The SQL result-fetch goroutine is in asyncBatchGetByRegions -> RunLoop.Exec(statementCtx) with the run loop in StateWaiting.
2. KILL TIDB QUERY or MAX_EXECUTION_TIME closes statementCtx.Done().
3. Concurrently, a TiKV BatchGet response reaches batchRecvLoop -> batchCommandsEntry.response -> Callback.Schedule -> RunLoop.Append.
4. Append observes StateWaiting, changes it to StateIdle, and decides to notify the waiter.
5. Exec selects ctx.Done() and returns from the canceled SQL.
6. Append sends on the unbuffered ready channel after its only receiver has exited and blocks forever.
7. Because Schedule is inside sync.Once.Do, competing cancellation/response completion also blocks on that sync.Once.
8. batchRecvLoop never reaches c.batched.Delete(requestID) or c.sent.Add(-1) and cannot process later responses on that stream.

The important user-visible distinction is that the canceled SQL itself can return Query execution was interrupted promptly. The persistent symptom appears afterward: unrelated SQLs sharing the affected TiKV batch-command stream hang or time out even while the gRPC connection remains READY.