feat(kernel): honor ctx deadlines on connect and mid-fetch

KERNEL_REV

@@ -1 +1 @@
-4d301455f7e70de9cb747e0f1143b8a3df8c5b48
+02c0a4346d413c5466a7066c081ceeb3a7eb205e

internal/backend/kernel/backend.go

@@ -97,17 +97,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
 	}
@@ -219,9 +213,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

internal/backend/kernel/cancel.go

@@ -0,0 +1,118 @@
+//go:build cgo && databricks_kernel
+
+package kernel
+
+/*
+#include <stdlib.h>
+#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)
+}

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
+}

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 {