feat(kernel): ctx-deadline cancellation, ABI/mTLS/CloudFetch controls, and log forwarding#409
feat(kernel): ctx-deadline cancellation, ABI/mTLS/CloudFetch controls, and log forwarding#409mani-mathur-arch wants to merge 11 commits into
Conversation
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 <mani.mathur@databricks.com>
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 <mani.mathur@databricks.com>
42af196 to
26cf09b
Compare
…etch cancel 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.
…etch cancel 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 <mani.mathur@databricks.com>
662aedc to
7c094f8
Compare
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 <mani.mathur@databricks.com>
…ions
Wire the four new kernel C-ABI symbols on the driver side, extending the
existing experimental-config surface:
- checkABIVersion(): a sync.Once handshake at the top of OpenSession compares
the linked library's kernel_abi_version() against the header's
DATABRICKS_KERNEL_ABI_VERSION and refuses to connect on a mismatch, so a
differently-built prebuilt .a can't be silently misread.
- WithKernelClientCertificate(cert, key): mTLS client identity forwarded via
the paired kernel_session_config_set_tls_client_certificate. Both PEM halves
are required together — validateKernelConfig rejects an unpaired credential
loudly so a lone key can't be silently dropped — and the key is never logged.
- WithKernelCloudFetch(enabled): a tri-state *bool (nil keeps the kernel
default on; set forwards kernel_session_config_set_cloudfetch_enabled).
Distinct from the plain-bool WithCloudFetch, whose unset state can't be told
from false.
All three are kernel-only and rejected loudly on the Thrift path (the connector
fails when KernelExperimental is non-nil). The reflective classification guard is
extended so a new experimental field can't ship unforwarded/unrejected.
CloseSession stays fire-and-forget: the C ABI now also offers
kernel_session_close_blocking, but adopting it would make close a blocking
network round-trip with no deadline honored, so swapping to it is grouped with
the cancellable-close follow-up.
Requires the kernel ABI-version / mTLS / CloudFetch symbols; bump KERNEL_REV to
the merged kernel revision before this builds in CI (it links a locally-staged
archive today).
Tests: experimental-field classification guard + option wiring + Thrift
rejection + DeepCopy (untagged); TestSetKernelTLS mTLS case, TestABIVersionMatches,
and mTLS-pairing validation (tagged / config). Default CGO_ENABLED=0 build
unchanged.
Co-authored-by: Isaac
Signed-off-by: Mani Kaustubh Mathur <mani.mathur@databricks.com>
WithKernelClientCertificate marks the kernel experimental config as set, but an empty cert+key pair (e.g. from a failed PEM load) was indistinguishable from the option never being called: the old XOR validation accepted it and applyKernelTLS skipped the setter, so a caller who explicitly requested mTLS connected with no client identity. Add an explicit TLSClientCertConfigured marker (robust against nil-vs-empty), set it whenever the option is invoked, and tighten validateKernelConfig to reject any incomplete mTLS request (missing cert, missing key, or both empty). Covered by new option/validation/DeepCopy tests and the exhaustiveness guard. Signed-off-by: Mani Kaustubh Mathur <mani.mathur@databricks.com>
The ABI-version handshake was only exercised on the happy path (TestABIVersionMatches, cgo build). The mismatch branch — the runtime hazard the check exists for, a driver header linked against a differently-built prebuilt .a — had no coverage because it lived behind the cgo symbols and the one-shot abiCheckOnce cache. Extract the pure got-vs-want verdict into compareABI (untagged) and have checkABIVersion delegate to it, then add TestCompareABI asserting the mismatch returns a non-nil error naming both versions and the remediation. Being untagged, the negative test runs under the default CGO_ENABLED=0 build with no kernel lib linked. Production behavior is unchanged. Signed-off-by: Mani Kaustubh Mathur <mani.mathur@databricks.com>
Register a cgo-exported trampoline (kernelLogTrampoline) as the kernel's log sink (kernel_set_log_callback) so kernel-internal tracing events flow into the driver's logger — the unified Go+kernel stream, honouring logger.SetLogOutput — instead of only reaching stderr via kernel_init_logging. This completes the one-knob logging story: #399 unified the LEVEL across the Go binding lines and the kernel's Rust lines (PECOBLR-3650) against the stderr sink; this replaces that sink with the callback so the Rust lines land in the same writer as everything else. The reverse-call machinery: a runtime/cgo.Handle round-trips the *logSink through the C void* ctx (cgo pointer rules), and a defer recover() firewall converts any panic into a dropped line (a panic across the cgo boundary would abort the process) — the same shape a kernel->host token-provider callback would need. The driver maps its own log level (resolveKernelLogArg, unchanged) and passes it as the callback's level argument, so the kernel's Rust lines follow DATABRICKS_LOG_LEVEL, not just RUST_LOG; DBSQL_KERNEL_DEBUG still defers to RUST_LOG (NULL level). The forwarding sink logs through a SNAPSHOT of the driver logger taken at install and pinned to TraceLevel, for two reasons: (1) no double gating — the kernel already filtered events against the level we passed it, so re-gating through the live logger.Logger (at the driver level) would re-drop exactly the events the DBSQL_KERNEL_DEBUG override was meant to surface; (2) no data race — the kernel drain thread reads only its immutable snapshot rather than the mutable global logger a concurrent SetLogLevel/SetLogOutput would rewrite. A non-Success install is surfaced at Warn (visible at the default level), not the Debug-gated klog. The two C-ABI logging routes share the one process-global subscriber and are mutually exclusive; we install the callback, and a non-Success install (host already installed a subscriber) is logged, never fatal to connect. Level mapping and sink routing are pure Go (logforward.go, untagged) so they test under CGO_ENABLED=0 — including the per-level mapping assertions and the not-re-gated-by-driver-level property. The reverse-call round-trip, the recover firewall (a panicking sink must not crash the process), the wrong-type/nil-ctx guards, and delivery after a panic are exercised by tagged tests via a C invoke seam. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur <mani.mathur@databricks.com>
Remove ticket codes, review shorthand, and ops jargon so the PR reads cleanly for OSS reviewers. Signed-off-by: Mani Kaustubh Mathur <mani.mathur@databricks.com>
Rely on resolveKernelLogArg returning an empty level for the RUST_LOG override instead of normalizing it twice. Signed-off-by: Mani Kaustubh Mathur <mani.mathur@databricks.com>
The consolidated branch now carries the log-callback forwarding (log_callback.go references kernel_set_log_callback / kernel_log_callback), so KERNEL_REV must point at a kernel revision that exports those symbols. Bump from the tier2 rev 02c0a434 (pre-log-callback) to 946c263, the head of kernel PR #169 (mani/c-abi-log-callback), which defines them. The kernel-lib build resolves this unmerged SHA via its PR-head-ref fallback, so the tagged kernel-backend build links a header carrying the symbols. Temporary: re-pin to the squash-merge SHA once kernel #169 lands. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur <mani.mathur@databricks.com>
mani-mathur-arch
changed the title
| kc.TLSSkipHostnameVerify = ke.TLSSkipHostnameVerify | ||
| kc.TLSClientCertPEM = ke.TLSClientCertPEM | ||
| kc.TLSClientKeyPEM = ke.TLSClientKeyPEM | ||
| kc.CloudFetchEnabled = ke.CloudFetchEnabled |
There was a problem hiding this comment.
High — buildKernelConfig field forwarding is untested; a dropped CloudFetchEnabled line would ship green
TestBuildKernelConfig (kernel_config_test.go:411) asserts only that TLSTrustedCertsPEM and TLSSkipHostnameVerify forward. The three new forwarded lines here — TLSClientCertPEM / TLSClientKeyPEM / CloudFetchEnabled (kernel_config.go:149-151) — have no runtime-copy assertion. TestKernelExperimentalFieldsClassified only checks the disposition map, not the actual copy (its own docstring warns of exactly this).
Deleting kc.CloudFetchEnabled = ke.CloudFetchEnabled makes WithKernelCloudFetch(false) a silent no-op (CloudFetch stays on) with zero test failures — the precise silent-option-drop the classified/forwarded machinery exists to prevent. Same exposure for the two mTLS halves.
Suggested fix
Extend the "experimental TLS fields forwarded" subtest (or add one) to set TLSClientCertPEM, TLSClientKeyPEM, and CloudFetchEnabled on KernelExperimental and assert all three land in the returned kernel.Config, plus a CloudFetchEnabled == nil (unset) case.
/full-review · feedback: #code-review-squad-feedback
| // 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) { |
There was a problem hiding this comment.
High — connect-cancel e2e never reaches the kernel_session_open_cancellable path it claims to cover
This test passes an already-cancelled ctx to the first query. OpenSession returns at its entry guard if err := ctx.Err(); err != nil { return err } (backend.go:203) — pre-existing behavior — long before the new kernel_session_open_cancellable + ctxWatcher and the dual-wrap "session_open cancelled" branch (backend.go:199-217). database/sql may even return ctx.Err() before calling the driver at all. The wrong-error-type check is only t.Logf, not a failure, so the docstring's claim ("covers the OpenSession → kernel_session_open_cancellable path") is false.
The headline connect-side mid-connect cancellation (ctxWatcher firing during a blocking connect, and the branch that preserves both the ctx error and *KernelError) is unexercised — it could be deleted or broken and every test still passes. Contrast the rigorous timing dance in the mid-fetch test.
Suggested fix
Add a test that fires the deadline during a slow connect (cold-start warehouse / unreachable host with a short ctx timeout) and assert the terminal error contains "session_open cancelled" and matches errors.Is(context.DeadlineExceeded) — mirroring the fetch test's mid-call proof.
/full-review · feedback: #code-review-squad-feedback
| // without the marker, so a config assembled without the option can't fail open | ||
| // either. | ||
| if ke := cfg.KernelExperimental; ke != nil { | ||
| certLen, keyLen := len(ke.TLSClientCertPEM), len(ke.TLSClientKeyPEM) |
There was a problem hiding this comment.
Medium — WithCloudFetch(false) is silently ignored on the kernel path, violating this backend's "nothing silently ignored" contract (flagged independently by 3 reviewers: architecture, agent-compat, devil's-advocate)
validateKernelConfig rejects every other unsupported neutral option loudly (Port, Transport, Timeout, retries-disabling WithRetries — all wrap ErrNotSupportedByKernel), and the mTLS pair is validated right here. But the backend-neutral WithCloudFetch(false) (sets cfg.UseCloudFetch) is classified inert (kernel_config_test.go:359), never read by buildKernelConfig, and never rejected here — so it is silently dropped and CloudFetch stays on (the kernel default).
doc.go:230 claims "Nothing is silently ignored" and lists only WithMaxRows/WithRetries as inert; WithCloudFetch's inertness is documented nowhere. An operator disabling CloudFetch via the standard documented option (e.g. to bypass a proxy that blocks the S3 pre-signed-URL path) silently still uses it; WithCloudFetch(false) + WithKernelCloudFetch(true) is contradictory config resolved silently in favor of the kernel knob.
Suggested fix
UseCloudFetch defaults to true, so an explicit false is distinguishable from unset — reject UseCloudFetch == false on the kernel path (wrapping ErrNotSupportedByKernel, pointing at WithKernelCloudFetch(false)), OR forward it into CloudFetchEnabled when WithKernelCloudFetch is unset. At minimum, add WithCloudFetch to doc.go's inert list.
/full-review · feedback: #code-review-squad-feedback
| case kernelLevelInfo: | ||
| s.log.Info().Str("target", target).Msg(message) | ||
| case kernelLevelDebug: | ||
| s.log.Debug().Str("target", target).Msg(message) |
There was a problem hiding this comment.
Medium — kernel WARN/ERROR logs now inflate the driver's shared Warn/Error stream, unthrottled and unclassified
This PR changes the kernel's Rust log destination from stderr to the shared logger.Logger sink. At the default Warn level the kernel forwards WARN+ERROR, and forward re-emits them at driver Warn/Error with no rate-limiting, sampling, dedup, or user-fault/infra reclassification. Kernel-internal transient events (retry "retrying request" WARN lines during a 503 / S3 storm) now interleave into the driver's own Warn/Error stream.
This reintroduces exactly the operational noise the lastError isUserFault split (cgo.go:241 — "don't inflate the WARN rate on-call alerts key on") was built to suppress, through a different door. A transient outage across many concurrent connections can make log-based alerting keyed on the driver's WARN/ERROR rate fire on recoverable events.
Suggested fix
Either forward kernel events one severity lower by default (kernel ERROR→Warn, WARN→Info) and require DBSQL_KERNEL_DEBUG to preserve native severity, or add coarse rate-limiting/dedup in forward. Document that enabling kernel forwarding raises the driver's Warn/Error volume.
/full-review · feedback: #code-review-squad-feedback
| // 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) |
There was a problem hiding this comment.
Low — a fresh ctxWatcher (goroutine + channel + 2 cgo token calls) is allocated per Arrow batch on any cancellable context (devil's-advocate conf 86; performance conf 64)
nextBatch calls newCtxWatcher(r.ctx) + defer watcher.stop() on every batch. For a cancellable ctx (the idiomatic request-scoped / WithTimeout case) each batch does kernel_cancel_token_new, spawns a goroutine, and stop() does close(done) + wg.Wait() + kernel_cancel_token_free. The "no watcher overhead on the common background-context path" comment is true only for context.Background(); the common service case pays per-batch churn. The execute path (operation.go), by contrast, creates exactly one canceller for the whole statement.
A multi-GB CloudFetch result (thousands of batches) under a deadline ctx incurs thousands of goroutine spawn/teardown cycles, chan+WaitGroup syncs, and cgo token alloc/free. Dominated by per-batch network latency in the normal case (likely <2%), but strictly wasted work on the driver's known-sensitive large-result path — and if the kernel returns buffered/prefetched batches sub-millisecond, the fixed watcher tax approaches the per-batch cost.
Suggested fix
Create one ctxWatcher per kernelRows (tied to r.ctx) in newKernelRows, reuse its token across all nextBatch calls, and stop() it in Close — matching the execute path's once-per-statement amortization. A fired token correctly aborts the current and all subsequent fetches.
/full-review · feedback: #code-review-squad-feedback
| }) | ||
| consumed = true | ||
| if err != nil { | ||
| // Prefer the caller's ctx error when the connect was interrupted by the |
There was a problem hiding this comment.
Low — the cancelled-path error-wrap rationale and dual-%w line are triplicated across backend.go / rows.go / operation.go
The same ~6-line rationale comment ("Prefer the caller's ctx error… Wrap BOTH… so errors.Is … AND errors.As … Mirrors the execute path") plus the identical fmt.Errorf("kernel: X cancelled: %w (kernel error: %w)", ctxErr, conv(err)) dual-%w line now appears in three places: here (backend.go:206-214), rows.go:205-213, and the pre-existing operation.go:176-186. The dual-%w trick (making both errors.Is and errors.As work) is the load-bearing, subtle part.
Changing the wrap contract (add queryId, adapt to a Go multi-%w semantics change) is a 3-site edit, and the three prose copies drift independently.
Suggested fix
Extract cancelledErr(op string, ctxErr, kernelErr error) error holding the format + dual-%w, document the rationale once on it, and have the call sites use it.
/full-review · feedback: #code-review-squad-feedback
Code Review Squad — ReviewScore: 69/100 — HIGH RISK The C-ABI/FFI mechanics are sound — security and language reviewers independently verified cert handling, memory ownership, cancel-token lifecycle, and the log trampoline against the kernel header at the pinned rev, both clean. Risk is concentrated in two test-coverage gaps (a dropped forwarding line or a broken connect-cancel path would ship green) and a UX/contract inconsistency where the standard WithCloudFetch(false) is silently ignored on the kernel path (flagged by 3 reviewers). No confirmed production correctness bug; the HIGH RISK band is driven by the test gaps, not by shipped defects. Medium FindingsM2 — This PR adds mid-fetch cancellation:
This finding has no in-diff line to anchor an inline comment (the paragraph sits just past the last doc.go hunk, which ends at line 266).
|
1 participant
What
Hardens the SEA-via-kernel backend across three areas — context cancellation, connection-security/CloudFetch controls, and log forwarding — all behind the
cgo && databricks_kernelbuild tag. The defaultCGO_ENABLED=0Thrift build is unchanged; every kernel-only option is rejected loudly on the Thrift path.1. Honor
contextdeadlines on connect and mid-fetchBoth blocking cgo calls now route through cancellable C-ABI entry points via a
ctxWatcherthat bridges a Gocontextonto a kernel cancel token:OpenSession→kernel_session_open_cancellable(connect)rows.go nextBatch→kernel_result_stream_next_batch_cancellable(mid-fetch)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(), dropping the in-flight kernel request future (a real abort, not a detach). The call sites return the ctx error while preserving the kernel error as its cause, and evict a session-fatal connection before returning. A non-cancellable ctx yields a nil watcher → NULL token → the plain, unchanged path (zero overhead).CloseSessionstays fire-and-forget.Closes PECOBLR-3645 and the connect half of PECOBLR-3646.
2. ABI-version handshake + mTLS + CloudFetch toggle
Wires four new kernel C-ABI symbols onto the driver's experimental-config surface:
sync.Oncecheck at the top ofOpenSessioncompares the linked library'skernel_abi_version()against the header'sDATABRICKS_KERNEL_ABI_VERSIONand refuses to connect on a mismatch, so a differently-built prebuilt archive can't be silently misread.WithKernelClientCertificate(cert, key)— mTLS client identity forwarded via the pairedkernel_session_config_set_tls_client_certificate. Both PEM halves are required together;validateKernelConfigrejects an incomplete/unpaired credential loudly (via an explicit "configured" marker robust against nil-vs-empty) rather than failing open with no client identity. The private key is never logged.WithKernelCloudFetch(enabled)— a tri-state toggle (nil keeps the kernel default on; set forwardskernel_session_config_set_cloudfetch_enabled), distinct from the plain-boolWithCloudFetchwhose unset state can't be told from false.A reflective exhaustiveness guard ensures a future
KernelExperimentalfield can't ship unforwarded/unrejected.3. Forward kernel (Rust) logs into the driver logger
The kernel's internal Rust tracing events are routed into the same driver logger sink as the Go binding lines (including a custom
logger.SetLogOutput) via a reverse-call callback (kernel_set_log_callback) — instead of going to stderr throughkernel_init_logging. The one-knob log level (DATABRICKS_LOG_LEVEL/DBSQL_KERNEL_DEBUG+RUST_LOG) is mirrored into the kernel subscriber. The Go trampoline has a panic firewall (recover) and is nil-/wrong-type-safe; the forwarding sink snapshots the logger at first connect so forwarded events aren't re-gated by a later driver-level change.Build
KERNEL_REVis pinned to the kernel revision that exports the cancellable, ABI-version, mTLS, CloudFetch, and log-callback symbols this driver code links against.Tests
ctxWatcher— fires the token on cancel and on deadline, is nil-safe on an uncancellable ctx, and tears down cleanly without firing (TestCtxWatcher*).kernel_config_test.go,kernel_experimental_test.go).TestLogSinkForward*,TestLogCallback*).CGO_ENABLED=0): the ABI-version compare (match + mismatch,TestCompareABI/TestABIVersionMatches), config classification/exhaustiveness guard, and pure config-assembly all run in the default build with no kernel lib linked.kernel_e2e_test.go): connect under an already-cancelled ctx fails fast; a long streaming query past its deadline is cancelled promptly withcontext.DeadlineExceeded; query cancel mid-execution.How we know it works: the tagged suite exercises the real cgo ctx→token bridge and the real callback round-trip (not mocks), the untagged suite guarantees the default build and its config contract are unaffected, and the staging e2e confirms the cancellation and connection-security behavior end-to-end against a live warehouse.
Co-authored-by: Isaac