228 lines
9.3 KiB
PL/PgSQL
228 lines
9.3 KiB
PL/PgSQL
-- Daily rollup table for `task_usage`. Background: the dashboard query
|
||
-- ListRuntimeUsage runs `SUM() GROUP BY DATE(created_at), provider, model`
|
||
-- against the raw event stream and is called once per runtime row on the
|
||
-- runtimes list (plus once per detail page load), so it dominates DB load
|
||
-- as event volume grows. We materialise the day-bucketed aggregate here
|
||
-- so reads scan O(days × providers × models) rows instead of O(events).
|
||
--
|
||
-- All query dimensions are denormalised into the table so reads never
|
||
-- need to join `agent_task_queue`. The PK doubles as the upsert key for
|
||
-- the rollup worker.
|
||
CREATE TABLE task_usage_daily (
|
||
bucket_date DATE NOT NULL,
|
||
workspace_id UUID NOT NULL,
|
||
runtime_id UUID NOT NULL,
|
||
provider TEXT NOT NULL,
|
||
model TEXT NOT NULL,
|
||
input_tokens BIGINT NOT NULL DEFAULT 0,
|
||
output_tokens BIGINT NOT NULL DEFAULT 0,
|
||
cache_read_tokens BIGINT NOT NULL DEFAULT 0,
|
||
cache_write_tokens BIGINT NOT NULL DEFAULT 0,
|
||
event_count BIGINT NOT NULL DEFAULT 0,
|
||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||
PRIMARY KEY (bucket_date, workspace_id, runtime_id, provider, model)
|
||
);
|
||
|
||
-- Primary read path: runtime detail page + runtimes-list cost cell, both
|
||
-- filter by runtime_id and order by date DESC. bucket_date DESC in the
|
||
-- index lets the query avoid an extra sort.
|
||
CREATE INDEX idx_task_usage_daily_runtime_date
|
||
ON task_usage_daily (runtime_id, bucket_date DESC);
|
||
|
||
-- Workspace-wide aggregations hit this index instead of fanning out per
|
||
-- runtime.
|
||
CREATE INDEX idx_task_usage_daily_workspace_date
|
||
ON task_usage_daily (workspace_id, bucket_date DESC);
|
||
|
||
-- Single-row state table tracking how far the rollup worker has consumed.
|
||
CREATE TABLE task_usage_rollup_state (
|
||
id SMALLINT PRIMARY KEY DEFAULT 1 CHECK (id = 1),
|
||
watermark_at TIMESTAMPTZ NOT NULL DEFAULT '1970-01-01 00:00:00+00',
|
||
last_run_started_at TIMESTAMPTZ,
|
||
last_run_finished_at TIMESTAMPTZ,
|
||
last_run_rows BIGINT NOT NULL DEFAULT 0,
|
||
last_error TEXT
|
||
);
|
||
INSERT INTO task_usage_rollup_state (id) VALUES (1) ON CONFLICT DO NOTHING;
|
||
|
||
-- Window-based aggregation primitive. Used by both the cron-driven
|
||
-- watermark advancer and the offline backfill command, so they stay
|
||
-- byte-identical in their semantics. Returns the number of output rows
|
||
-- touched.
|
||
--
|
||
-- IDEMPOTENCY CONTRACT (this is the important bit):
|
||
-- For every (bucket_date, workspace_id, runtime_id, provider, model)
|
||
-- key that has at least one task_usage row whose `updated_at` falls in
|
||
-- [p_from, p_to), this function REPLACES the corresponding daily row
|
||
-- with the SUM of *all* task_usage rows for that key (regardless of
|
||
-- their updated_at). It does NOT add a delta.
|
||
--
|
||
-- Consequences:
|
||
-- * Replaying the same window is safe — the row is rebuilt from raw
|
||
-- each time, so the result converges.
|
||
-- * Two callers (cron + backfill) processing overlapping windows is
|
||
-- safe — both write the same value.
|
||
-- * `UpsertTaskUsage` corrections that overwrite token counts are
|
||
-- captured: the corrected row's updated_at gets bumped, the next
|
||
-- window picks up its bucket key, and the bucket is recomputed
|
||
-- from current truth.
|
||
--
|
||
-- Cost: the recompute reads ALL task_usage rows for each dirty bucket,
|
||
-- not just the windowed slice. In steady state only "today" buckets are
|
||
-- dirty (a handful of keys per active runtime), so this stays cheap.
|
||
-- During backfill the entire history's bucket keys become dirty once;
|
||
-- the backfill walks history in monthly slices to bound the working
|
||
-- set per call.
|
||
CREATE OR REPLACE FUNCTION rollup_task_usage_daily_window(
|
||
p_from TIMESTAMPTZ,
|
||
p_to TIMESTAMPTZ
|
||
)
|
||
RETURNS BIGINT
|
||
LANGUAGE plpgsql
|
||
AS $$
|
||
DECLARE
|
||
v_rows BIGINT;
|
||
BEGIN
|
||
IF p_from >= p_to THEN
|
||
RETURN 0;
|
||
END IF;
|
||
|
||
WITH dirty_keys AS (
|
||
SELECT DISTINCT
|
||
DATE(tu.created_at) AS bucket_date,
|
||
i.workspace_id AS workspace_id,
|
||
atq.runtime_id AS runtime_id,
|
||
tu.provider AS provider,
|
||
tu.model AS model
|
||
FROM task_usage tu
|
||
JOIN agent_task_queue atq ON atq.id = tu.task_id
|
||
JOIN issue i ON i.id = atq.issue_id
|
||
WHERE atq.runtime_id IS NOT NULL
|
||
AND (
|
||
-- Steady state: rows updated within the watermark window.
|
||
-- Hits idx_task_usage_updated_at directly.
|
||
(tu.updated_at >= p_from AND tu.updated_at < p_to)
|
||
-- Legacy rows from before migration 072 (updated_at IS NULL)
|
||
-- — discoverable via created_at + the partial index added
|
||
-- in 077. Steady-state windows after backfill never include
|
||
-- historical dates, so this branch is a no-op once the
|
||
-- backfill has swept history.
|
||
OR (tu.updated_at IS NULL
|
||
AND tu.created_at >= p_from
|
||
AND tu.created_at < p_to)
|
||
)
|
||
),
|
||
recomputed AS (
|
||
SELECT
|
||
dk.bucket_date,
|
||
dk.workspace_id,
|
||
dk.runtime_id,
|
||
dk.provider,
|
||
dk.model,
|
||
SUM(tu.input_tokens)::bigint AS input_tokens,
|
||
SUM(tu.output_tokens)::bigint AS output_tokens,
|
||
SUM(tu.cache_read_tokens)::bigint AS cache_read_tokens,
|
||
SUM(tu.cache_write_tokens)::bigint AS cache_write_tokens,
|
||
COUNT(*)::bigint AS event_count
|
||
FROM dirty_keys dk
|
||
JOIN agent_task_queue atq ON atq.runtime_id = dk.runtime_id
|
||
JOIN issue i ON i.id = atq.issue_id
|
||
AND i.workspace_id = dk.workspace_id
|
||
JOIN task_usage tu ON tu.task_id = atq.id
|
||
AND tu.provider = dk.provider
|
||
AND tu.model = dk.model
|
||
AND DATE(tu.created_at) = dk.bucket_date
|
||
GROUP BY 1, 2, 3, 4, 5
|
||
)
|
||
INSERT INTO task_usage_daily AS d (
|
||
bucket_date, workspace_id, runtime_id, provider, model,
|
||
input_tokens, output_tokens, cache_read_tokens, cache_write_tokens,
|
||
event_count
|
||
)
|
||
SELECT
|
||
bucket_date, workspace_id, runtime_id, provider, model,
|
||
input_tokens, output_tokens, cache_read_tokens, cache_write_tokens,
|
||
event_count
|
||
FROM recomputed
|
||
ON CONFLICT (bucket_date, workspace_id, runtime_id, provider, model) DO UPDATE
|
||
SET input_tokens = EXCLUDED.input_tokens,
|
||
output_tokens = EXCLUDED.output_tokens,
|
||
cache_read_tokens = EXCLUDED.cache_read_tokens,
|
||
cache_write_tokens = EXCLUDED.cache_write_tokens,
|
||
event_count = EXCLUDED.event_count,
|
||
updated_at = now();
|
||
|
||
GET DIAGNOSTICS v_rows = ROW_COUNT;
|
||
RETURN v_rows;
|
||
END;
|
||
$$;
|
||
|
||
-- Cron entry point. Advances the watermark by one window each call.
|
||
--
|
||
-- Invariants:
|
||
-- * `pg_try_advisory_lock(4242)` serialises overlapping ticks.
|
||
-- * The window upper bound is `now() - 5 minutes`. The lag exists
|
||
-- because `task_usage` rows are written from a separate transaction;
|
||
-- a row with updated_at = T can become visible to this snapshot at
|
||
-- some t > T. 5 minutes is a generous bound on that visibility delay
|
||
-- and keeps the dashboard "today" bucket at most ~10 min stale
|
||
-- (5 min lag + 5 min cron period).
|
||
-- * On error we record `last_error` and re-raise; the watermark is NOT
|
||
-- advanced because the UPDATE that advances it only runs after the
|
||
-- upsert succeeds.
|
||
-- * SAFE TO RUN CONCURRENTLY WITH BACKFILL: the window primitive is
|
||
-- idempotent (see contract above), so even if cron fires while the
|
||
-- offline backfill is also walking history, the worst case is some
|
||
-- bucket gets written twice with the same value.
|
||
CREATE OR REPLACE FUNCTION rollup_task_usage_daily()
|
||
RETURNS BIGINT
|
||
LANGUAGE plpgsql
|
||
AS $$
|
||
DECLARE
|
||
v_lock_ok BOOLEAN;
|
||
v_from TIMESTAMPTZ;
|
||
v_to TIMESTAMPTZ;
|
||
v_rows BIGINT := 0;
|
||
BEGIN
|
||
SELECT pg_try_advisory_lock(4242) INTO v_lock_ok;
|
||
IF NOT v_lock_ok THEN
|
||
RETURN 0;
|
||
END IF;
|
||
|
||
BEGIN
|
||
UPDATE task_usage_rollup_state
|
||
SET last_run_started_at = now(),
|
||
last_error = NULL
|
||
WHERE id = 1
|
||
RETURNING watermark_at INTO v_from;
|
||
|
||
v_to := now() - INTERVAL '5 minutes';
|
||
|
||
IF v_from < v_to THEN
|
||
v_rows := rollup_task_usage_daily_window(v_from, v_to);
|
||
|
||
UPDATE task_usage_rollup_state
|
||
SET watermark_at = v_to,
|
||
last_run_finished_at = now(),
|
||
last_run_rows = v_rows
|
||
WHERE id = 1;
|
||
ELSE
|
||
UPDATE task_usage_rollup_state
|
||
SET last_run_finished_at = now(),
|
||
last_run_rows = 0
|
||
WHERE id = 1;
|
||
END IF;
|
||
|
||
PERFORM pg_advisory_unlock(4242);
|
||
RETURN v_rows;
|
||
EXCEPTION WHEN OTHERS THEN
|
||
UPDATE task_usage_rollup_state
|
||
SET last_error = SQLERRM,
|
||
last_run_finished_at = now()
|
||
WHERE id = 1;
|
||
PERFORM pg_advisory_unlock(4242);
|
||
RAISE;
|
||
END;
|
||
END;
|
||
$$;
|