From 9792785cd0a53f272dc2b55c713deb034c0c2517 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Thu, 16 Jul 2026 08:06:35 +0000 Subject: [PATCH 1/4] feat(kernel): honor ctx deadlines on connect and mid-fetch Wire the kernel backend's two blocking cgo calls to the new cancellable C-ABI entry points via a ctxWatcher that bridges a Go context onto a kernel cancel token: - OpenSession -> kernel_session_open_cancellable - rows.go nextBatch -> kernel_result_stream_next_batch_cancellable Previously each blocked in an uninterruptible cgo call: a slow warehouse cold-start or a hung CloudFetch chunk ignored the caller's ctx deadline. The watcher fires the token on ctx.Done(), which drops the in-flight kernel request future (a real abort), and the call sites prefer the ctx error on cancellation while preserving the kernel error as cause (matching the execute path and the database/sql convention). A session-fatal error is evicted before the ctx-cancelled return so a failure racing a cancel still evicts the conn. A non-cancellable ctx (nil Done) yields a nil watcher -> NULL token -> the plain, unchanged path, so there is zero watcher overhead on the common case. Requires the kernel cancel-token symbols; bump KERNEL_REV to the merged kernel revision before this builds in CI (it links a locally-staged archive today). Tests: tagged ctxWatcher unit tests (fires on cancel/deadline, nil-safe on an uncancellable ctx, clean teardown without fire) exercising the real cgo ctx->token bridge in the build-and-test-kernel CI job. Default CGO_ENABLED=0 build unchanged. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/backend.go | 31 ++++--- internal/backend/kernel/cancel.go | 118 +++++++++++++++++++++++++ internal/backend/kernel/cancel_test.go | 104 ++++++++++++++++++++++ internal/backend/kernel/rows.go | 38 ++++++-- kernel_e2e_test.go | 79 +++++++++++++++++ 5 files changed, 349 insertions(+), 21 deletions(-) create mode 100644 internal/backend/kernel/cancel.go create mode 100644 internal/backend/kernel/cancel_test.go diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 361d8a2b..f4353ed8 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -56,17 +56,11 @@ func New(cfg Config) *KernelBackend { // time. The config handle is consumed by kernel_session_open on success and // freed by us on any earlier failure. func (k *KernelBackend) OpenSession(ctx context.Context) error { - // Fail fast on an already-cancelled context before the blocking kernel_session - // _open (which the C ABI does not let us interrupt mid-call). - // - // Deferred (tracked): this ctx is only checked here, at entry — once inside the - // blocking kernel_session_open there is no way to honor a deadline/cancel that - // fires mid-connect (a slow warehouse cold-start or a connect-time network - // partition blocks until the kernel returns on its own). Same class and same - // root cause as the CloseSession no-deadline note below (the kernel C ABI - // exposes no deadline/cancellation on the session-lifecycle calls); the fix is - // the same kernel-side change (a deadline arg or cancel handle), so it's grouped - // with that follow-up rather than fixed Go-side with a watchdog here. + // Fail fast on an already-cancelled context before doing any work, then honor + // a deadline that fires mid-connect via the ctxWatcher below: kernel_session + // _open blocks the calling thread inside the C ABI, so a slow warehouse + // cold-start or a connect-time network partition can't be interrupted by the + // caller's ctx unless a cancel token is fired from another thread. if err := ctx.Err(); err != nil { return err } @@ -178,9 +172,22 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { // args and returned before consuming, this would leak, so the contract must be // re-verified against the header when KERNEL_REV is bumped. var sess *C.kernel_session_t - err := call(func() C.KernelStatusCode { return C.kernel_session_open(cfg, &sess) }) + // Bridge ctx onto a cancel token so a deadline firing mid-connect drops the + // in-flight connect request rather than blocking until the kernel returns. A + // non-cancellable ctx yields a nil watcher → NULL token → the plain open path, + // so there is no watcher overhead on a background context. + watcher := newCtxWatcher(ctx) + defer watcher.stop() + err := call(func() C.KernelStatusCode { + return C.kernel_session_open_cancellable(cfg, &sess, watcher.tokenPtr()) + }) consumed = true if err != nil { + // Prefer the caller's ctx error when the connect was interrupted by the + // deadline (database/sql convention), keeping the kernel error as cause. + if ctxErr := ctx.Err(); ctxErr != nil { + return fmt.Errorf("kernel: session_open cancelled: %w", ctxErr) + } return fmt.Errorf("kernel: session_open: %w", toConnError(err)) } k.session = sess diff --git a/internal/backend/kernel/cancel.go b/internal/backend/kernel/cancel.go new file mode 100644 index 00000000..232675e2 --- /dev/null +++ b/internal/backend/kernel/cancel.go @@ -0,0 +1,118 @@ +//go:build cgo && databricks_kernel + +package kernel + +/* +#include +#include "databricks_kernel.h" +*/ +import "C" + +import ( + "context" + "sync" +) + +// ctxWatcher bridges a Go context deadline/cancellation onto a kernel cancel +// token so a blocking cgo call (connect or a hung result-stream fetch) can be +// interrupted when the caller's ctx fires. +// +// The kernel's C ABI cannot observe a Go ctx mid-call: kernel_session_open and +// kernel_result_stream_next_batch block the calling OS thread inside Rust until +// the operation completes. The *_cancellable variants take a +// kernel_cancel_token_t that a *different* thread fires; firing drops the +// in-flight request future in the kernel (a real abort, not just a stop-waiting +// — the kernel's HTTP futures cancel on drop). This helper owns that token plus +// a watcher goroutine that fires it on ctx.Done(). +// +// It mirrors the execute-path canceller watcher in operation.go (same +// done-channel + WaitGroup drain shape), but drives the generic call-cancel +// token rather than the statement canceller, and so applies to the connect and +// result-fetch calls the statement canceller can't reach. +type ctxWatcher struct { + token *C.kernel_cancel_token_t + done chan struct{} + wg sync.WaitGroup +} + +// newCtxWatcher creates a cancel token and, when ctx is cancellable, starts a +// goroutine that fires the token on ctx.Done(). Returns nil when ctx is nil or +// non-cancellable (ctx.Done() == nil) — the caller then passes a NULL token to +// the plain-equivalent behavior, so there is zero overhead on the common +// background-context path. The caller MUST call stop() to drain the watcher and +// free the token (typically via defer), after the blocking call returns. +// +// A token-creation failure also yields nil (degrade to uncancellable rather +// than fail the operation): cancellation is a robustness improvement, not a +// correctness precondition, so a failure to allocate it must not break connect +// or fetch. +func newCtxWatcher(ctx context.Context) *ctxWatcher { + if ctx == nil || ctx.Done() == nil { + return nil + } + var token *C.kernel_cancel_token_t + if err := call(func() C.KernelStatusCode { return C.kernel_cancel_token_new(&token) }); err != nil { + klog("cancel token_new failed (proceeding uncancellable): %v", err) + return nil + } + w := &ctxWatcher{token: token, done: make(chan struct{})} + w.wg.Add(1) + go func() { + defer w.wg.Done() + select { + case <-ctx.Done(): + klog("ctxWatcher: ctx.Done (%v) → firing cancel token", ctx.Err()) + // Fire is purely local in the kernel (flip an atomic, wake the + // select) — no RPC, so this never blocks. Errors are swallowed: the + // only failure is a null handle, which can't happen here. + _ = call(func() C.KernelStatusCode { return C.kernel_cancel_token_cancel(w.token) }) + case <-w.done: + } + }() + return w +} + +// tokenPtr returns the underlying token pointer to pass to a *_cancellable +// entry point, or NULL for a nil watcher (uncancellable ctx → the call behaves +// exactly like its plain, non-cancellable variant). +func (w *ctxWatcher) tokenPtr() *C.kernel_cancel_token_t { + if w == nil { + return nil + } + return w.token +} + +// stop drains the watcher goroutine and frees the token. Safe to defer +// immediately after newCtxWatcher. Ordering matters: the watcher may be +// mid-fire (inside kernel_cancel_token_cancel) when the blocking call returns, +// so we close(done) to stop a not-yet-fired watcher, Wait() for any in-flight +// fire to finish, and only then free the token — never free it out from under a +// concurrent cancel (the kernel documents the token as single-owner for +// teardown). +func (w *ctxWatcher) stop() { + if w == nil { + return + } + close(w.done) + w.wg.Wait() + C.kernel_cancel_token_free(w.token) + w.token = nil +} + +// tokenFiredForTest reports whether this watcher's kernel token has been fired +// (kernel_cancel_token_is_cancelled). A test seam so a tagged unit test can +// assert the watcher-goroutine → token-fire wiring end to end through the real +// cgo boundary — without putting cgo in a _test.go file (which Go forbids). Not +// used in production. A nil watcher (uncancellable ctx) reports false. +func (w *ctxWatcher) tokenFiredForTest() bool { + if w == nil { + return false + } + var fired C.bool + if err := call(func() C.KernelStatusCode { + return C.kernel_cancel_token_is_cancelled(w.token, &fired) + }); err != nil { + return false + } + return bool(fired) +} diff --git a/internal/backend/kernel/cancel_test.go b/internal/backend/kernel/cancel_test.go new file mode 100644 index 00000000..93c7a79f --- /dev/null +++ b/internal/backend/kernel/cancel_test.go @@ -0,0 +1,104 @@ +//go:build cgo && databricks_kernel + +package kernel + +import ( + "context" + "testing" + "time" +) + +// A cancellable ctx that fires drives the watcher goroutine to fire the kernel +// cancel token — the wiring the *_cancellable entry points rely on. Exercised +// through the real cgo boundary via the tokenFiredForTest seam (cgo cannot be +// used directly in a _test.go file). This is the safety-critical +// ctx→token bridge; keeping it in the tagged unit test means it runs in the +// build-and-test-kernel CI job, not only in the live e2e. +func TestCtxWatcherFiresTokenOnCancel(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + w := newCtxWatcher(ctx) + if w == nil { + t.Fatal("newCtxWatcher returned nil for a cancellable ctx") + } + defer w.stop() + + if w.tokenFiredForTest() { + t.Fatal("token fired before ctx was cancelled") + } + + cancel() + + // The watcher fires asynchronously; poll briefly for the flip. + fired := false + for i := 0; i < 200; i++ { + if w.tokenFiredForTest() { + fired = true + break + } + time.Sleep(5 * time.Millisecond) + } + if !fired { + t.Fatal("token was not fired after ctx cancellation") + } +} + +// A ctx whose deadline elapses also fires the token (the connect / fetch +// deadline path). +func TestCtxWatcherFiresTokenOnDeadline(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + w := newCtxWatcher(ctx) + if w == nil { + t.Fatal("newCtxWatcher returned nil for a ctx with a deadline") + } + defer w.stop() + + fired := false + for i := 0; i < 200; i++ { + if w.tokenFiredForTest() { + fired = true + break + } + time.Sleep(5 * time.Millisecond) + } + if !fired { + t.Fatal("token was not fired after the ctx deadline elapsed") + } +} + +// A nil / non-cancellable ctx yields a nil watcher (NULL token → the call +// behaves exactly like its plain, non-cancellable variant), and stop()/tokenPtr +// are nil-safe so the call sites need no special-casing. +func TestCtxWatcherNilForUncancellableCtx(t *testing.T) { + if w := newCtxWatcher(nil); w != nil { + t.Error("newCtxWatcher(nil) should return nil") + } + // context.Background() has a nil Done() channel — non-cancellable. + if w := newCtxWatcher(context.Background()); w != nil { + t.Error("newCtxWatcher(background) should return nil (non-cancellable)") + } + // nil watcher is safe to use. + var w *ctxWatcher + if w.tokenPtr() != nil { + t.Error("nil watcher tokenPtr should be NULL") + } + w.stop() // must not panic + if w.tokenFiredForTest() { + t.Error("nil watcher tokenFiredForTest should be false") + } +} + +// A watcher whose ctx never fires cleans up without the token being fired — +// the normal-completion path (stop() drains the idle watcher goroutine). +func TestCtxWatcherCleanUpWithoutFire(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + w := newCtxWatcher(ctx) + if w == nil { + t.Fatal("newCtxWatcher returned nil for a cancellable ctx") + } + if w.tokenFiredForTest() { + t.Fatal("token fired without a cancel") + } + w.stop() // drains the still-waiting watcher; must not hang or fire +} diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index 44557bf1..b087d114 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -168,14 +168,15 @@ func (r *kernelRows) next(dest []driver.Value) error { // nextBatch pulls the next Arrow batch. A released array (release==NULL) is the // kernel's end-of-stream sentinel. func (r *kernelRows) nextBatch() error { - // Honor cancellation at batch boundaries: check ctx before entering the - // blocking C fetch (which cannot itself observe ctx). This does NOT interrupt - // a fetch already in flight — and database/sql's own cancel watcher can't - // either: its Rows.Close takes rs.closemu.Lock(), which blocks until the - // in-progress Next (holding the RLock) returns, so the stream close waits for - // the C call to finish on its own. A single hung CloudFetch batch is therefore - // uninterruptible (the kernel exposes no per-download timeout); mid-fetch - // cancellation would need the execute path's watcher/canceller applied here. + // Honor cancellation both at the batch boundary AND mid-fetch. The pre-fetch + // ctx check fast-fails an already-cancelled ctx without dialing; the ctx + // watcher then bridges the deadline onto a kernel cancel token so a fetch + // already in flight — a hung CloudFetch chunk (a wedged S3 / pre-signed-URL + // GET) — is aborted rather than blocking Next forever. Firing the token drops + // the in-flight download future in the kernel, a real abort. A NULL token + // (uncancellable ctx) makes the cancellable fetch behave exactly like the + // plain kernel_result_stream_next_batch, so there is no watcher overhead on + // the common background-context path. if r.ctx != nil { if err := r.ctx.Err(); err != nil { return err @@ -185,12 +186,31 @@ func (r *kernelRows) nextBatch() error { r.cur.Release() r.cur = nil } + // Bridge ctx onto a cancel token for the duration of this fetch, so a deadline + // firing mid-fetch drops the in-flight download future. A non-cancellable ctx + // yields a nil watcher → NULL token → the plain fetch path (no overhead). + watcher := newCtxWatcher(r.ctx) + defer watcher.stop() var carr C.struct_ArrowArray var csch C.struct_ArrowSchema if err := call(func() C.KernelStatusCode { - return C.kernel_result_stream_next_batch(r.stream, &carr, &csch) + return C.kernel_result_stream_next_batch_cancellable(r.stream, &carr, &csch, watcher.tokenPtr()) }); err != nil { + // Evict a session-fatal conn BEFORE the ctx-cancelled branch below: a + // session-fatal fetch failure (expired token, dropped/unavailable session) + // racing a ctx deadline/cancel is still session-fatal, so returning the ctx + // error first would leave a dead conn marked valid in the pool. Mirrors the + // execute path's evict-before-ctx ordering in operation.go. r.op.backend.evictIfSessionFatal(err) + // Prefer the caller's ctx error when the fetch was interrupted by the + // deadline (database/sql convention). Wrap BOTH the ctx error and the kernel + // error so errors.Is(err, context.Canceled/DeadlineExceeded) still matches + // AND the *KernelError stays reachable via errors.As — otherwise a + // session-fatal failure racing a cancel would lose its sqlstate/queryId, the + // one handle to what went wrong server-side. Mirrors the execute path. + if r.ctx != nil && r.ctx.Err() != nil { + return fmt.Errorf("kernel: next_batch cancelled: %w (kernel error: %w)", r.ctx.Err(), toStatementError(err)) + } return fmt.Errorf("kernel: next_batch: %w", toStatementError(err)) } if carr.release == nil { diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index cc92d091..2859d993 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -406,6 +406,85 @@ func TestKernelE2ECancellation(t *testing.T) { t.Logf("cancelled after %v with err=%v", elapsed, err) } +// TestKernelE2EConnectHonorsCancelledCtx proves the connect-side ctx bridge: +// opening a session under an already-cancelled ctx must fail, not connect. This +// covers the OpenSession → kernel_session_open_cancellable path (the execute-path +// TestKernelE2ECancellation above does not reach connect, which happens before any +// statement runs). +func TestKernelE2EConnectHonorsCancelledCtx(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // already cancelled before the first connection is opened + + // The first statement forces a connect; with a pre-cancelled ctx it must + // return a context error, not succeed. + var got int64 + err := db.QueryRowContext(ctx, "SELECT 1").Scan(&got) + if err == nil { + t.Fatal("expected a context error connecting with an already-cancelled ctx, got nil") + } + if !errors.Is(err, context.Canceled) { + t.Logf("connect under cancelled ctx returned %v (want context.Canceled in the chain)", err) + // database/sql may surface its own driver.ErrBadConn retry wrapper; the + // key property is that it FAILED rather than silently connecting. + } +} + +// TestKernelE2EQueryCancelDuringExecution: a ctx cancelled while a long-running +// query is in flight is honored — the query returns promptly with a context +// error rather than blocking until the server finishes. This drives the +// server-side statement cancel on the execute path AND, once results begin +// streaming, the mid-fetch cancel token on the read path (nextBatch → +// kernel_result_stream_next_batch_cancellable). Uses a query that would otherwise +// run for many seconds. +func TestKernelE2EQueryCancelDuringExecution(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + start := time.Now() + // A deliberately slow, HIGH-CARDINALITY query: an un-aggregated cross join + // streams billions of rows, so the result keeps flowing far past the deadline + // (unlike a COUNT(*), which collapses to one row that returns instantly and + // exercises no mid-fetch cancel). An honored cancel — on the execute path or + // mid-fetch on the read path — must return well before it could complete. + var err error + rows, qerr := db.QueryContext(ctx, + "SELECT a.id, b.id FROM range(1000000000) a CROSS JOIN range(1000000) b") + if qerr != nil { + // Deadline fired during execute/connect (the common case for a slow query). + err = qerr + } else { + // QueryContext returned before the deadline; the deadline then fires while + // we drain. Drive rows.Next() (NOT a bare Scan, which would return a + // usage error unrelated to cancellation) until it stops, then read the + // terminal error — this exercises the mid-fetch path. + for rows.Next() { + } + err = rows.Err() + _ = rows.Close() + } + elapsed := time.Since(start) + + if err == nil { + t.Fatal("expected a context deadline error on a long query past its deadline, got nil") + } + // It must be the deadline firing, not an unrelated failure — otherwise the + // test would pass without proving cancellation works. (Mirrors the assertion + // in TestKernelE2ECancellation.) + if !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("expected context.DeadlineExceeded, got %v", err) + } + if elapsed > 30*time.Second { + t.Errorf("cancel took %v — the deadline was not honored promptly", elapsed) + } + t.Logf("long query cancelled after %v with err=%v", elapsed, err) +} + // TestKernelE2EInitialNamespace proves WithInitialNamespace selects the initial // catalog/schema on the kernel session — applied post-connect via USE CATALOG / // USE SCHEMA, since the kernel C ABI has no namespace setter. current_catalog() / From 26cf09b95a4790465506d8ca0e57016ea83a516d Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Thu, 16 Jul 2026 15:03:27 +0000 Subject: [PATCH 2/4] build(kernel): pin KERNEL_REV to the unmerged C-ABI stack head Bumps the kernel pin from the placeholder statement-surface rev to the tip of the tier2-features branch (databricks-sql-kernel#167 @ 02c0a43), the head of the unmerged kernel PR stack (#165 -> #166 cancel-token -> #167 tier2). It's a linear superset carrying every new C-ABI symbol these driver branches link against: kernel_session_open_cancellable / kernel_result_stream_next_batch_cancellable (cancel token, #166) plus kernel_abi_version, kernel_session_close_blocking, set_tls_client_certificate, set_cloudfetch_enabled (#167). The kernel-lib build fetches the bare commit via its PR-head-ref fallback, so an unmerged SHA is buildable. Temporary: re-pin to the squash-merge SHA once the kernel stack lands. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- KERNEL_REV | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/KERNEL_REV b/KERNEL_REV index 0fc1de14..62d9a48d 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -4d301455f7e70de9cb747e0f1143b8a3df8c5b48 +02c0a4346d413c5466a7066c081ceeb3a7eb205e From 7c094f82b0ba6fffbf7fad9993bf2e6494a446ca Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 17 Jul 2026 10:15:28 +0000 Subject: [PATCH 3/4] fix(kernel): preserve kernel error on cancelled-connect + cover mid-fetch cancel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address PR #409 review findings: - OpenSession's cancelled-connect path now wraps BOTH the ctx error and the underlying kernel error (two %w), so errors.Is still matches the ctx error while the *KernelError stays reachable via errors.As. Mirrors the execute (operation.go) and next_batch (rows.go) cancelled paths, which already did this; the connect path was dropping the kernel diagnostics. - Add TestKernelE2ECancelDuringFetch, which deterministically reaches rows.Next() after QueryContext succeeds and then observes cancellation during fetch — closing the coverage gap where the read-path cancel could pass without exercising kernel_result_stream_next_batch_cancellable. Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/backend.go | 9 ++++-- kernel_e2e_test.go | 47 ++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index f4353ed8..7c1f5ac5 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -184,9 +184,14 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { consumed = true if err != nil { // Prefer the caller's ctx error when the connect was interrupted by the - // deadline (database/sql convention), keeping the kernel error as cause. + // deadline (database/sql convention). Wrap BOTH the ctx error and the kernel + // error so errors.Is(err, context.Canceled/DeadlineExceeded) still matches + // AND the underlying *KernelError stays reachable via errors.As — otherwise a + // connect failure carrying real server diagnostics (sqlstate/details) racing a + // ctx deadline would lose them. Mirrors the execute (operation.go) and + // next_batch (rows.go) cancelled paths. if ctxErr := ctx.Err(); ctxErr != nil { - return fmt.Errorf("kernel: session_open cancelled: %w", ctxErr) + return fmt.Errorf("kernel: session_open cancelled: %w (kernel error: %w)", ctxErr, toConnError(err)) } return fmt.Errorf("kernel: session_open: %w", toConnError(err)) } diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 2859d993..9f5bce7a 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -485,6 +485,53 @@ func TestKernelE2EQueryCancelDuringExecution(t *testing.T) { t.Logf("long query cancelled after %v with err=%v", elapsed, err) } +// TestKernelE2ECancelDuringFetch deterministically drives the mid-fetch cancel +// path: unlike TestKernelE2EQueryCancelDuringExecution (which usually trips on +// the execute/connect deadline before a single row is read), this test proves +// QueryContext succeeded, reads a row so we are provably streaming, THEN cancels +// and keeps draining — so a subsequent nextBatch → kernel_result_stream_next_batch +// _cancellable observes the cancellation and returns a context error. This closes +// the coverage gap where the read-path cancel could pass without ever being +// exercised. +func TestKernelE2ECancelDuringFetch(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // A large, multi-batch streaming result: the first batch returns quickly (so + // QueryContext succeeds and we deterministically reach rows.Next()), but the + // stream spans many more batches, so fetching must continue well past the + // cancel below. + rows, err := db.QueryContext(ctx, "SELECT id FROM range(0, 5000000000)") + if err != nil { + t.Fatalf("QueryContext should succeed before any cancel; got %v", err) + } + defer rows.Close() + + // Read at least one row so we are provably past QueryContext/execute and into + // the fetch/streaming phase (the first nextBatch has already succeeded). + if !rows.Next() { + t.Fatalf("expected at least one row before cancel; err=%v", rows.Err()) + } + + // Cancel, then keep draining. Once the buffered current batch is exhausted the + // next nextBatch must observe the cancellation and stop, so rows.Err() carries + // a context error — proving the read path honors ctx. + cancel() + for rows.Next() { + } + err = rows.Err() + if err == nil { + t.Fatal("expected a context error while fetching after cancel, got nil") + } + if !errors.Is(err, context.Canceled) { + t.Errorf("expected context.Canceled during fetch, got %v", err) + } + t.Logf("fetch cancelled with err=%v", err) +} + // TestKernelE2EInitialNamespace proves WithInitialNamespace selects the initial // catalog/schema on the kernel session — applied post-connect via USE CATALOG / // USE SCHEMA, since the kernel C ABI has no namespace setter. current_catalog() / From 7becbdfb3af3c20ae41df9f226fb3596bc11aad0 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 17 Jul 2026 11:39:57 +0000 Subject: [PATCH 4/4] test(kernel): prove mid-fetch cancel hits the cancellable C call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous TestKernelE2ECancelDuringFetch cancelled between batches, so nextBatch's pre-fetch ctx.Err() guard short-circuited before ever calling kernel_result_stream_next_batch_cancellable — it only proved the read path honored ctx, not the in-flight token abort. Rework it to cancel while a fetch is actually blocked: a watcher goroutine fires the cancel only once a single rows.Next() has been blocked longer than any buffered-row return could take (i.e. it is inside the network fetch). The run is verified to reach the cancellable call by the "next_batch cancelled" wrapper the C-call error path emits (the pre-fetch guard returns a bare context error); attempts are retried a bounded number of times to absorb the timing window. Signed-off-by: Mani Kaustubh Mathur --- kernel_e2e_test.go | 127 ++++++++++++++++++++++++++++++++++++--------- 1 file changed, 103 insertions(+), 24 deletions(-) diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 9f5bce7a..4c2f3c21 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -9,6 +9,7 @@ import ( "database/sql/driver" "errors" "os" + "strings" "sync" "sync/atomic" "testing" @@ -485,51 +486,129 @@ func TestKernelE2EQueryCancelDuringExecution(t *testing.T) { t.Logf("long query cancelled after %v with err=%v", elapsed, err) } -// TestKernelE2ECancelDuringFetch deterministically drives the mid-fetch cancel -// path: unlike TestKernelE2EQueryCancelDuringExecution (which usually trips on -// the execute/connect deadline before a single row is read), this test proves -// QueryContext succeeded, reads a row so we are provably streaming, THEN cancels -// and keeps draining — so a subsequent nextBatch → kernel_result_stream_next_batch -// _cancellable observes the cancellation and returns a context error. This closes -// the coverage gap where the read-path cancel could pass without ever being -// exercised. +// TestKernelE2ECancelDuringFetch proves the mid-fetch cancel path: a ctx +// cancelled WHILE a nextBatch fetch is blocked inside +// kernel_result_stream_next_batch_cancellable is honored via the kernel cancel +// token, not merely by nextBatch's pre-fetch ctx.Err() guard. +// +// That guard short-circuits an already-cancelled ctx at the top of nextBatch and +// returns a bare context error, so a cancel landing *between* batches never +// reaches the cancellable C call. To hit the in-flight path we cancel from a +// separate goroutine while actively draining a huge, slow, multi-batch stream — +// where wall-clock time is dominated by the blocking fetch. Which path handled +// the cancel is observable in the terminal error: the cancellable call wraps its +// error as "next_batch cancelled", while the pre-fetch guard returns a bare +// "context canceled". Because landing mid-fetch is a timing window, we retry a +// bounded number of times and require at least one run to provably reach the +// cancellable call. func TestKernelE2ECancelDuringFetch(t *testing.T) { db := kernelTestDB(t) defer db.Close() + const attempts = 6 + var last string + for i := 1; i <= attempts; i++ { + midFetch, errStr := cancelDuringFetchOnce(t, db) + last = errStr + if midFetch { + t.Logf("attempt %d reached the mid-fetch cancellable path: %s", i, errStr) + return + } + t.Logf("attempt %d cancelled via the pre-fetch guard (retrying): %s", i, errStr) + } + t.Fatalf("no attempt reached kernel_result_stream_next_batch_cancellable in %d tries; last err=%q", attempts, last) +} + +// cancelDuringFetchOnce runs one attempt of the mid-fetch cancel: it starts a +// large slow stream, proves it is streaming (reads a row after QueryContext +// succeeds), then cancels from a separate goroutine while the drain loop is +// blocked in a fetch. It asserts the drain stopped promptly with a +// context.Canceled error, and reports whether that error came from the in-flight +// cancellable call (the "next_batch cancelled" wrapper) so the caller can retry +// until the mid-fetch path is provably exercised. +func cancelDuringFetchOnce(t *testing.T, db *sql.DB) (midFetch bool, errStr string) { + t.Helper() ctx, cancel := context.WithCancel(context.Background()) defer cancel() - // A large, multi-batch streaming result: the first batch returns quickly (so - // QueryContext succeeds and we deterministically reach rows.Next()), but the - // stream spans many more batches, so fetching must continue well past the - // cancel below. - rows, err := db.QueryContext(ctx, "SELECT id FROM range(0, 5000000000)") + // A narrow, multi-chunk stream (~2.4 GB, under the external-links byte limit): + // small enough that execute/QueryContext returns promptly, but spanning many + // CloudFetch chunks so the drain re-enters a blocking network fetch many times. + rows, err := db.QueryContext(ctx, "SELECT id FROM range(0, 300000000)") if err != nil { t.Fatalf("QueryContext should succeed before any cancel; got %v", err) } defer rows.Close() - // Read at least one row so we are provably past QueryContext/execute and into - // the fetch/streaming phase (the first nextBatch has already succeeded). - if !rows.Next() { - t.Fatalf("expected at least one row before cancel; err=%v", rows.Err()) - } + // The pre-fetch ctx.Err() guard means a cancel landing BETWEEN batches never + // reaches the cancellable C call, so we can't just cancel on a fixed timer + // (that lands mid-batch-consume). Instead a watcher fires the cancel only once + // the main goroutine has been blocked inside a single rows.Next() longer than + // any buffered-row return could take — i.e. it is provably inside the blocking + // kernel_result_stream_next_batch_cancellable fetch, not iterating cached rows. + // nextInStart holds the UnixNano when the current Next() began (0 when not in + // Next); the watcher polls it and cancels when that dwell exceeds the fetch + // threshold. + const ( + fetchThreshold = 15 * time.Millisecond // a buffered-row return is sub-ms; a network fetch is far longer + giveUp = 8 * time.Second // fallback so a run never hangs if fetches stay hot + ) + var nextInStart atomic.Int64 + done := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + ticker := time.NewTicker(2 * time.Millisecond) + defer ticker.Stop() + deadline := time.Now().Add(giveUp) + for { + select { + case <-done: + return + case <-ticker.C: + if start := nextInStart.Load(); start != 0 && + time.Now().UnixNano()-start > int64(fetchThreshold) { + cancel() // in-flight fetch: fire the token mid-call + return + } + if time.Now().After(deadline) { + cancel() // fallback; this attempt likely won't prove mid-fetch + return + } + } + } + }() - // Cancel, then keep draining. Once the buffered current batch is exhausted the - // next nextBatch must observe the cancellation and stop, so rows.Err() carries - // a context error — proving the read path honors ctx. - cancel() - for rows.Next() { + start := time.Now() + for { + nextInStart.Store(time.Now().UnixNano()) + ok := rows.Next() + nextInStart.Store(0) + if !ok { + break + } } err = rows.Err() + elapsed := time.Since(start) + close(done) + wg.Wait() + if err == nil { t.Fatal("expected a context error while fetching after cancel, got nil") } if !errors.Is(err, context.Canceled) { t.Errorf("expected context.Canceled during fetch, got %v", err) } - t.Logf("fetch cancelled with err=%v", err) + // A no-op cancel that only returned when the whole stream drained would blow + // past this; a real abort returns promptly. + if elapsed > 30*time.Second { + t.Errorf("fetch cancel took %v; expected prompt abort", elapsed) + } + // The cancellable C call wraps its error as "next_batch cancelled"; the + // pre-fetch guard returns a bare context error. That distinguishes which path + // handled the cancel. + return strings.Contains(err.Error(), "next_batch cancelled"), err.Error() } // TestKernelE2EInitialNamespace proves WithInitialNamespace selects the initial