diff --git a/KERNEL_REV b/KERNEL_REV index 0fc1de14..62d9a48d 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -4d301455f7e70de9cb747e0f1143b8a3df8c5b48 +02c0a4346d413c5466a7066c081ceeb3a7eb205e diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 2ad46cc8..7cf9a8d7 100644 --- a/internal/backend/kernel/backend.go +++ b/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 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() /