feat(data-warehouse): per-org concurrency budget for the duckgres sink - #68808
Conversation
Migration SQL ChangesHey 👋, we've detected some migrations on this PR. Here's the SQL output for each migration, make sure they make sense:
|
🔍 Migration Risk AnalysisWe've analyzed your migrations for potential risks. Summary: 1 Safe | 0 Needs Review | 0 Blocked ✅ SafeBrief or no lock, backwards compatible Last updated: 2026-07-09 22:24 UTC (b3e1c93) |
|
Reviews (1): Last reviewed commit: "feat(data-warehouse): per-org concurrenc..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| row = await conn.execute( | ||
| f""" | ||
| SELECT count(*) | ||
| FROM ( | ||
| SELECT m.org_id | ||
| FROM {DUCKGRES_LEASE_TABLE} l | ||
| JOIN unnest( | ||
| %(budget_team_ids)s::bigint[], | ||
| %(budget_org_ids)s::varchar[], | ||
| %(budget_values)s::int[] | ||
| ) AS m(team_id, org_id, budget) ON m.team_id = l.team_id | ||
| WHERE l.expires_at > now() | ||
| GROUP BY m.org_id | ||
| HAVING count(*) >= max(m.budget) | ||
| ) saturated | ||
| """, | ||
| { | ||
| "budget_team_ids": [team_id for team_id, _, _ in team_org_budgets], | ||
| "budget_org_ids": [org_id for _, org_id, _ in team_org_budgets], | ||
| "budget_values": [budget for _, _, budget in team_org_budgets], | ||
| }, | ||
| ) | ||
| result = await row.fetchone() |
There was a problem hiding this comment.
conn.execute() in psycopg3 is the cursor itself, not a row — naming it row then calling .fetchone() on it reads as confusing. Every other caller in this file uses cur via async with conn.cursor() as cur:. Aligning the name removes the mismatch.
| row = await conn.execute( | |
| f""" | |
| SELECT count(*) | |
| FROM ( | |
| SELECT m.org_id | |
| FROM {DUCKGRES_LEASE_TABLE} l | |
| JOIN unnest( | |
| %(budget_team_ids)s::bigint[], | |
| %(budget_org_ids)s::varchar[], | |
| %(budget_values)s::int[] | |
| ) AS m(team_id, org_id, budget) ON m.team_id = l.team_id | |
| WHERE l.expires_at > now() | |
| GROUP BY m.org_id | |
| HAVING count(*) >= max(m.budget) | |
| ) saturated | |
| """, | |
| { | |
| "budget_team_ids": [team_id for team_id, _, _ in team_org_budgets], | |
| "budget_org_ids": [org_id for _, org_id, _ in team_org_budgets], | |
| "budget_values": [budget for _, _, budget in team_org_budgets], | |
| }, | |
| ) | |
| result = await row.fetchone() | |
| cur = await conn.execute( | |
| f""" | |
| SELECT count(*) | |
| FROM ( | |
| SELECT m.org_id | |
| FROM {DUCKGRES_LEASE_TABLE} l | |
| JOIN unnest( | |
| %(budget_team_ids)s::bigint[], | |
| %(budget_org_ids)s::varchar[], | |
| %(budget_values)s::int[] | |
| ) AS m(team_id, org_id, budget) ON m.team_id = l.team_id | |
| WHERE l.expires_at > now() | |
| GROUP BY m.org_id | |
| HAVING count(*) >= max(m.budget) | |
| ) saturated | |
| """, | |
| { | |
| "budget_team_ids": [team_id for team_id, _, _ in team_org_budgets], | |
| "budget_org_ids": [org_id for _, org_id, _ in team_org_budgets], | |
| "budget_values": [budget for _, _, budget in team_org_budgets], | |
| }, | |
| ) | |
| result = await cur.fetchone() |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| "budget_team_ids": [team_id for team_id, _, _ in team_org_budgets] if team_org_budgets else [], | ||
| "budget_org_ids": [org_id for _, org_id, _ in team_org_budgets] if team_org_budgets else [], | ||
| "budget_values": [budget for _, _, budget in team_org_budgets] if team_org_budgets else [], |
There was a problem hiding this comment.
The three list comprehensions that unpack team_org_budgets into budget_team_ids / budget_org_ids / budget_values appear verbatim in both get_delta_succeeded_and_lock (with the if team_org_budgets else [] guard) and count_orgs_at_budget (without the guard, because the early-return handles it). A small private helper — e.g. _budget_params(team_org_budgets) → dict — would honour OnceAndOnlyOnce and reduce the surface for future drift if the tuple layout ever changes.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| @dataclass(frozen=True) | ||
| class SinkEnablement: | ||
| """The sink's per-refresh view of who it serves and how hard it may push. | ||
|
|
||
| ``team_org_budgets`` carries one (team_id, org_id, sink_max_concurrency) | ||
| row per enabled team — the queue DB has no teams table, so the claim query | ||
| receives this mapping as unnest'd arrays to enforce the per-org group | ||
| budget fleet-wide. | ||
| """ | ||
|
|
||
| team_ids: list[int] | ||
| team_org_budgets: list[tuple[int, str, int]] |
There was a problem hiding this comment.
frozen=True with mutable list fields gives an incomplete immutability guarantee
frozen=True only blocks attribute reassignment — it does not prevent callers from mutating the list contents (e.g. enablement.team_ids.append(42)). In practice the lists are never mutated here, so this causes no current defect. If true immutability is the intent, the fields could use tuple[int, ...] and tuple[tuple[int, str, int], ...] respectively; otherwise removing frozen=True (and relying on the caller convention) is equally honest about the actual semantics.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
|
Audit note: a session-level adversarial audit of the whole duckgres consumer stack found one HIGH in this PR's original commit — the budget filter ran after the candidate |
|
Reviews (2): Last reviewed commit: "Merge remote-tracking branch 'origin/mas..." | Re-trigger Greptile |
🤖 CI report |
🦔 Hogbox preview · ✅ ready▶ Open the preview
commit |
|
Docs from this PR will be published at posthog.com
Preview will be ready in ~10 minutes. Click Preview link above to access docs at |
The sink's only concurrency bound was pods x max_concurrency, and every in-flight group holds a connection to its org's duckgres server - so scaling replicas (KEDA already scales 3..8 on slot utilization) would concentrate an unbounded number of connections on whichever org has the deepest backlog. That coupling is what blocks raising maxPods. - DuckgresServer.sink_max_concurrency (default 4): per-org cap on concurrently leased sink groups, admin-tunable per org - enablement now returns (team_id, org_id, budget) rows alongside the team set; the queue DB has no teams table, so the claim receives the mapping as unnest'd arrays - the claim query counts live group leases per org (any owner - the fleet-wide usage picture) and ranks eligible groups per org, admitting only the org's remaining slots; saturated orgs' groups are skipped, not deferred, so they never head-of-line block a pod from filling max_groups with other orgs' work - soft cap: overlapping polls can overshoot by a group for one window; the org duckgres max_connections stays the hard backstop - new duckgres_sink_orgs_at_budget gauge + maintenance log field so a budget-limited fleet is distinguishable from a pod-limited one (the KEDA utilization signal stays truthful: capped orgs' skipped work no longer inflates slot usage) Unmapped teams (dev, mid-refresh race) stay uncapped; existing behavior is unchanged when no mapping is passed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… LIMIT A fully saturated org with a deep eligible backlog (e.g. a pending backfill) filled the candidate window with rows the group filter then dropped, so the pod claimed nothing while other orgs had work - the same starvation mode the other-owner lease exclusion already guards against, and the exact scenario the budget targets. Saturated orgs are now excluded pre-LIMIT; partially available orgs keep post-LIMIT per-org ranking. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
e6df3d3 to
d4d05ca
Compare
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
| %(budget_values)s::int[] | ||
| ) AS t(team_id, org_id, budget) | ||
| ), | ||
| org_usage AS ( |
There was a problem hiding this comment.
Medium: Concurrent claims can exceed the org budget
org_usage is computed from a statement snapshot, but no organization-scoped row or advisory lock serializes the subsequent lease inserts. An authenticated tenant can maintain eligible work across many schemas, allowing concurrent pods with different candidate windows to observe the same remaining slots and claim disjoint groups, opening more Duckgres connections than sink_max_concurrency. Serialize claims per organization or atomically allocate slots through a locked per-org counter.
PR overviewThis pull request adds a per-organization concurrency budget for the Duckgres warehouse sink, updating the job-claiming logic to limit how much work each org can run at once. The touched code manages database-backed leasing of Duckgres import jobs across schemas and workers. There is still one open issue in the claim path: concurrent workers can calculate available organization capacity from the same snapshot and then lease separate batches, allowing the intended per-org concurrency limit to be exceeded. An authenticated tenant with enough eligible work could use this to drive more Duckgres connections than configured, creating a resource-exhaustion risk. No issues have been fixed or addressed yet, so the PR still needs a serialization or atomic allocation mechanism before the budget is reliable. Open issues (1)
Fixed/addressed: 0 · PR risk: 5/10 |
2 participants
Problem
The duckgres sink's only concurrency bound is pods x max_concurrency (16), and every in-flight group holds a connection to its org's duckgres server. KEDA already autoscales the deployment (3 to 8 pods on P90 slot utilization), so the failure mode of raising maxPods further is concrete: one org with a deep backlog saturates every pod's slots, utilization pegs, KEDA scales up, and now every new slot also hammers that same org's duckgres server. Scale-up makes the hot org hotter, and there is no fleet-wide limit on how many connections one org can receive.
Changes
Per-org budget, enforced fleet-wide at claim time.
DuckgresServer.sink_max_concurrency(new column, default 4, migration 1248): the max number of the org's groups the sink may process concurrently across all pods. One in-flight group holds at most one duckgres connection (opened per batch inside the group's serial loop), so this bounds the org's sink connection footprint independently of pod count. Admin-tunable per org from day one.duckgres_sink_enablement) now returns(team_id, org_id, budget)rows alongside the enabled-team set, on the same 60s refresh with the same keep-previous-on-error semantics. The queue DB has no teams table, so the claim query receives the mapping as unnest'd arrays.max_groupswith other orgs' work. The pod stays work-conserving.max_connectionsremains the hard backstop. Teams without a mapping row (dev mode, mid-refresh race) are uncapped, preserving existing behavior; passing no mapping at all changes nothing.duckgres_sink_orgs_at_budgetgauge (count only, no org label) plus anorgs_at_budgetfield on the maintenance heartbeat log. Operator contract: backlog growing while pod utilization is low and this gauge is up means the fleet is budget-limited - raise the org's budget or its duckgres capacity, not replicas. Backlog growing with utilization high means raise maxPods.This also makes the existing KEDA signal truthful: a capped org's skipped work no longer inflates slot utilization, so autoscaling only adds pods when more pods genuinely help. Raising the charts
maxPodsbeyond 8 is intended as a follow-up once this lands.How did you test this code?
Automated tests only (I'm an agent; no manual prod testing). Full duckgres
pipeline_v3suite passes locally against a freshly created test DB (initial run had 5 failures that turned out to be a stale shared test DB carrying another branch's migration; all pass with--create-db). New tests, each catching a regression nothing covered before:TestOrgConcurrencyBudgetintest_jobs_db.py: a single fetch cannot claim past an org's budget; another pod's live leases count as usage and expired leases return the slot; a saturated org's older groups are skipped while another org's younger group fillsmax_groups(the work-conservation contract); an unmapped team is uncapped (the mapping must cap, never filter eligibility);count_orgs_at_budgetcounts a saturated org and not an under-budget one.test_enablement.py: per-org budgets ride along with each enabled team and reflect per-orgsink_max_concurrencyvalues (guards the plumbing - an empty mapping silently means uncapped).Existing tests pass unchanged: omitting the mapping preserves current claim behavior, and the consumer tests' keyed call-arg assertions are unaffected by the new kwarg.
Automatic notifications
Docs update
Internal consumer/scheduling behavior; no user-facing docs affected.
🤖 Agent context
Autonomy: Human-driven (agent-assisted)
/django-migrations(additive column onposthog_duckgresserver, not a hot table, no external writers) and/writing-tests(value gate for the new test classes).DuckgresServerrather than a code constant so per-org tuning needs no second migration; scope limited to cap + claim fairness (poll jitter and planner dedup deferred until replica counts actually grow).🤖 Generated with Claude Code