fix(plugins): memoise BigQuery analytics readiness across re-initialisation - #7020
fix(plugins): memoise BigQuery analytics readiness across re-initialisation#7020chelsealong wants to merge 1 commit into
Conversation
…sation BigQueryAgentAnalyticsPlugin re-ran its full table/view readiness pass on every re-initialisation instead of once per process. Each pass issues one CREATE OR REPLACE VIEW statement per _EVENT_VIEW_DEFS entry, awaited from before_run_callback, so the DDL sat on the request path. A host that builds a short-lived Runner per request over one shared plugin closes the plugin after every request (Runner.close() -> PluginManager.close() -> plugin.close()), which clears _started, so the next request re-ran all the view DDL. In production this exhausted the per-table BigQuery quota and added ~25s to median latency (2.3.0 -> 2.8.0 regression). Remember that readiness succeeded in a flag (_schema_ready) that, like _schema, survives close()/shutdown(), and gate the readiness pass on it. A *failed* attempt raises before the flag is set and is still retried on the next setup (the 2.8.0 intent); a *successful* one is not repeated (restoring the 2.3.0 cost profile). Also log one WARNING when a generation mismatch aborts an otherwise-successful setup, so a plugin churning through full setups is no longer silent. Closes google#7017
|
Reporter of #7017 here. Thanks for picking this up so quickly — the fix matches what we landed as a local mitigation, and putting One gap worth closing before merge:
The plugin has a custom
Suggested one-liner, alongside the existing resets: state["_started"] = False
+ state["_schema_ready"] = False
state["_startup_error"] = NoneCost is one readiness pass per process, which is the 2.3.0 profile this PR is restoring anyway — it does not weaken the fix, since the whole problem was re-verifying many times within one process. A test, if useful — it slots onto the existing def test_getstate_resets_pid(self):
"""Pickle state should have _init_pid = 0 to force re-init."""
plugin = self._make_plugin()
plugin._schema_ready = True
state = plugin.__getstate__()
assert state["_init_pid"] == 0
assert state["_started"] is False
# Readiness is a claim about the destination dataset, so a pickle
# restored elsewhere must verify it again rather than assume it.
assert state["_schema_ready"] is FalseTwo smaller notes, neither blocking:
The middleware side of the trigger is filed at ag-ui-protocol/ag-ui#2642, so the per-request teardown that drives the re-init gets fixed there too. |
lupuletic
left a comment
There was a problem hiding this comment.
Verified this branch (4196a06) against the reported production failure. The fix works and the approach is right.
I reproduced the host shape from #7017 — one shared plugin instance, _ensure_started() then close() per request, which is what before_run_callback and Runner.close() do — and counted the CREATE OR REPLACE VIEW statements over 5 cycles:
| view DDL statements | |
|---|---|
origin/main |
125 (25 per cycle) |
| this branch | 25 (once, total) |
That is exactly the intended shape: linear in re-initialisations before, constant after. The existing suite is also green on the branch — 511 passed in tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py.
I also like two choices here more than what we had in our own local mitigation:
- Putting
_schema_readyimmediately afterself._schemain__init__makes the "pure data, survives close" grouping self-evident. - Driving the tests through real
shutdown()cycles rather than poking_starteddirectly. A test that sets the private flag would keep passing if a futureclose()started clearing the memo, so it would not actually guard the regression. Yours does.
The one thing I would still fix before merge is the __getstate__ reset I described in my earlier comment — the memo currently crosses the pickle boundary while every other runtime field is reset, so a plugin restored into an environment whose dataset does not exist yet would skip table creation. One line, and it does not weaken the fix, since the problem being solved is re-verification within one process.
Everything else reads good to me. Thanks for turning this around so fast.
For anyone landing here from the same symptom: the middleware-side trigger — a per-request Runner closing plugins owned by a long-lived App — is filed at ag-ui-protocol/ag-ui#2642. This PR makes the plugin resilient to that; that one stops the churn at source.
Summary
Closes #7017.
BigQueryAgentAnalyticsPluginre-ran its full table/view readiness pass on every re-initialisation instead of once per process. Each pass issues oneCREATE OR REPLACE VIEWstatement per_EVENT_VIEW_DEFSentry (25 at HEAD), and it is awaited frombefore_run_callback, so the DDL sits on the agent request path.Any deployment that shares one plugin instance across short-lived
Runners (Runner.close()→PluginManager.close()→plugin.close()clears_started) re-runs all the view DDL per request. The reporter measured ~19k view-DDL jobs/hour, per-table BigQuery quota exhausted in ~2h, and ~25s added to median latency — a 2.3.0 → 2.8.0 regression.In 2.3.0 the readiness pass was gated behind the schema cache, so it ran at most once per process. At HEAD that gate was deliberately removed so a failed first attempt would still retry — but the comment's load-bearing assumption ("once
_startedis True the steady state pays no extra RPC") only holds when the plugin is initialised once per process, which is false whenever it is closed and reused.Fix
_schema_readyflag that, like_schema, survivesclose()/shutdown(), and gate_ensure_schema_existson it in_lazy_setup.WARNINGwhen a generation mismatch aborts an otherwise-successful setup, so a plugin churning through full setups is no longer completely silent (the issue observed 12 log lines against ~2000 setup runs).Scope is limited to the readiness memoisation and the one diagnostic log line. The view SQL,
_EVENT_VIEW_DEFS, and the_ensure_startedbackoff are untouched (the issue's non-goals).Testing Plan
Added
TestReadinessMemoisedAcrossReinitwith two tests:test_readiness_pass_runs_once_across_reinit_cycles— 3_ensure_started()/shutdown()cycles over one shared plugin issue the view DDL once total (client.querycalledlen(_EVENT_VIEW_DEFS)times), not once per cycle.test_failed_readiness_is_retried_and_not_memoised— a first readiness attempt that raises leaves_schema_readyFalse and is retried on the next setup, which then succeeds and memoises.Evidence the tests prove the bug: with the source change stashed (test kept), both new tests fail (
AttributeError: ... _schema_ready/ would-be 75 vs 25 DDL statements); with the fix restored they pass.Lint/format on the changed files:
pyink(2 files unchanged),isort(exit 0),ruff check(all checks passed).AI assistance disclosure
This change was prepared with AI assistance (Claude). A human reviewed the diff, the tests, and the failing-without-fix evidence before submission.