Extend querythrottler with lifetime-scoped admission control - #909
Draft
bgwines wants to merge 6 commits into
Draft
Extend querythrottler with lifetime-scoped admission control#909bgwines wants to merge 6 commits into
querythrottler with lifetime-scoped admission control#909bgwines wants to merge 6 commits into
Conversation
Add a generic `app_execution_context_id` field to `ExecuteOptions` (field 22) so a downstream admission-control strategy can attribute a query to the application-level request or async job that issued it, and add an `ADMISSION_CONTROL` value to the `ThrottlingStrategy` enum so a strategy can be selected via config. Regenerated `query.pb.go`, `query_vtproto.pb.go`, `querythrottler.pb.go`, and the vtadmin JS/TS bindings via `make proto`. AI disclosure: Claude Code assisted with development. Every line of code was either written by or carefully reviewed by me :) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Brett Wines <bwines@slack-corp.com>
Add a `PoolType` enum (`PoolTypeOltpRead`, `PoolTypeOlapRead`, `PoolTypeTx`) to identify which connection pool an admission request targets. The querythrottler framework threads this through `Deps.PoolCapacities` and the seam forwards it to an admission-control strategy. AI disclosure: Claude Code assisted with development. Every line of code was either written by or carefully reviewed by me :) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Brett Wines <bwines@slack-corp.com>
Extend the querythrottler framework so a strategy can act as a two-phase admission controller in addition to the existing predicate-style throttling: - Add the `AdmissionController` interface: `Admit(ctx, attrs, pool)` returns a release func so a strategy can hold a slot for a query's lifetime and free it on completion. - Add `QueryThrottler.AcquireAdmission`, which routes to the registered strategy when it implements `AdmissionController` and is a no-op (returning a no-op release) otherwise, so the seam is inert until a strategy is registered. - Add `SetPoolCapacities` and thread `Env` + `PoolCapacities` through `Deps` so a strategy can size itself against the live pool capacities. - Add `AppExecutionContextID` and `SchemaQualifiers` to `QueryAttributes`. - Add `registry.Unregister` to support test isolation. `admission_test.go` covers the no-op path, routing to a registered controller, rejection propagation, and strategy construction via the registry. AI disclosure: Claude Code assisted with development. Every line of code was either written by or carefully reviewed by me :) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Brett Wines <bwines@slack-corp.com>
Add an `APP_EXECUTION_CONTEXT_ID` query directive that flows from a query comment through `QueryHints.AppExecutionContextID` and vtgate's `VCursorImpl.SetAppExecutionContextID` into `ExecuteOptions`, so a vttablet admission-control strategy can attribute a query to the application-level request or async job that issued it. Regenerated the affected `cached_size.go` files for the new `QueryHints` field (which `engine.Plan` embeds). AI disclosure: Claude Code assisted with development. Every line of code was either written by or carefully reviewed by me :) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Brett Wines <bwines@slack-corp.com>
…aths Add the integration seam that calls the querythrottler's admission-control extension points from the connection-acquisition paths. With no admission-control strategy registered, `AcquireAdmission` is a no-op, so this is behaviorally inert on its own. - `query_executor.go`: reshape `getConn` to return a release func, build `QueryAttributes` (including schema qualifiers and app-execution-context id), acquire admission before taking a pooled conn, and split query timings by result error code. - `tx_pool.go`: acquire admission for `PoolTypeTx` in `Begin`, storing the release on the transaction so it is freed when the transaction completes. - `stateful_connection.go`: release any held admission slot in `ReleaseString` as an idempotent teardown funnel, so slots are freed even on paths that bypass `txComplete` (kill, shutdown, taint, renew failure). - `tx/api.go`: add the `AdmissionRelease` field to transaction `Properties`. - `tabletserver.go`: `wireAdmissionControl` registers the throttler as the tx pool's admitter and publishes live pool capacities via `SetPoolCapacities`. - `plan.go` / `stats.go`: cache per-plan schema qualifiers and add the by-error-code timings the seam records. Regenerated `planbuilder/cached_size.go`. AI disclosure: Claude Code assisted with development. Every line of code was either written by or carefully reviewed by me :) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Brett Wines <bwines@slack-corp.com>
bgwines
changed the title
querythrottlerquerythrottler with lifetime-scoped admission control
Pare the admission-control extension comments down to docstrings that match each file's local convention (an enum/field/function doc only where its siblings are documented), simplify `SetAppExecutionContextID` to a direct passthrough, and regenerate `query.pb.go` so its `app_execution_context_id` docstring matches the trimmed proto source. AI disclosure: Claude Code assisted with development. Every line of code was either written by or carefully reviewed by me :) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Brett Wines <bwines@slack-corp.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Background / Why?
The
querythrottlerframework today supports only predicate-style strategies: a strategy inspects a query, returns an allow/throttle decision, and its involvement ends there. That shape can't express admission control that must hold a resource for the lifetime of a query or transaction — e.g. a concurrency limiter that grants a connection-pool slot on acquire and releases it on completion. This PR adds the generic extension points such a strategy needs, without shipping any concrete strategy.The core addition is a two-phase
AdmissionControllerinterface (Admit(...) (release func(err error), err error)) thatQueryThrottler.AcquireAdmission(...)invokes from the execution path. When noAdmissionControlleris registered it returns a no-op release, so existing deployments see no behavioral change. The call is wired into both the query path (getConnnow returns a release the callers defer) and the transaction path (TxPool.Beginacquires forPoolTypeTxand stores the release ontx.Properties), and release is funneled throughStatefulConnection.ReleaseStringso every teardown path — commit, rollback, the transaction killer, shutdown, taint, renew-failure — frees the slot exactly once. Controllers also receiveDeps.PoolCapacitiesand a newAppExecutionContextID/SchemaQualifiersonQueryAttributesto reason about pool pressure and caller identity.Testing
new unit tests covering the no-op-when-unregistered path, routing to a registered controller, rejection propagation, and pool-capacity threading through strategy construction.
AI disclosure: Claude Code assisted with development. Every line of code was either written by or carefully reviewed by me :)