usage-aggregator.ts218 lines · main
1import { and, eq, isNull, sql as drizzleSql } from 'drizzle-orm';
2
3import { newId } from '@briven/shared';
4
5import { getDb } from '../db/client.js';
6import { projects, usageEvents, type UsageMetric } from '../db/schema.js';
7import { env } from '../env.js';
8import { log } from '../lib/logger.js';
9import { getAuthMauStats } from '../services/auth-mau.js';
10import { collectConnectionSecondsDeltas } from '../services/connection-seconds.js';
11import { isAuthEnabled } from '../services/tenant-config-store.js';
12import { getInvocationUsage, getStorageUsage } from '../services/usage.js';
13
14/**
15 * Hourly usage aggregator. For each non-deleted project:
16 * 1. invocations — count + duration from function_logs for the hour
17 * window [period_start, period_start + 1h)
18 * 2. storage_bytes — pg_total_relation_size sample at period_end
19 * 3. connection_seconds — scrape briven_realtime_connection_seconds_total
20 * from realtime /metrics, compute delta vs previous period
21 *
22 * One row per (project, hour, metric). Idempotent via the unique
23 * index so re-running the cron for a missed hour just overwrites the
24 * previous attempt — safe for catch-up after a restart.
25 *
26 * Polar push is a separate worker (`push-polar-meters.ts`, follow-up)
27 * that scans WHERE polar_push_status='pending'. Until that lands rows
28 * stay pending and the operator can verify the data via SQL.
29 */
30
31const HOUR_MS = 60 * 60 * 1000;
32
33/** First millisecond of the UTC hour containing `now`. */
34export function currentHourStart(now: Date = new Date()): Date {
35 return new Date(Math.floor(now.getTime() / HOUR_MS) * HOUR_MS);
36}
37
38/**
39 * Roll up the hour that JUST ENDED. Called from the cron. We deliberately
40 * lag by one full hour so the function_logs that landed at xx:59:59 are
41 * always inside the [start, end) window we read.
42 */
43export async function aggregateUsageForCompletedHour(now: Date = new Date()): Promise<{
44 hoursProcessed: number;
45 rowsWritten: number;
46}> {
47 const periodEnd = currentHourStart(now);
48 const periodStart = new Date(periodEnd.getTime() - HOUR_MS);
49
50 const db = getDb();
51 const activeProjects = await db
52 .select({ id: projects.id })
53 .from(projects)
54 .where(isNull(projects.deletedAt));
55
56 // Scrape the realtime /metrics counter once for the whole batch — the
57 // diff returned is the per-project delta since the previous hourly
58 // scrape. Projects not in the map had no realtime activity this hour;
59 // we don't write a zero row for them (saves index churn).
60 const connectionSeconds = await collectConnectionSecondsDeltas();
61
62 let rowsWritten = 0;
63 for (const { id: projectId } of activeProjects) {
64 try {
65 const written = await rollUpProject(
66 projectId,
67 periodStart,
68 periodEnd,
69 connectionSeconds.get(projectId) ?? null,
70 );
71 rowsWritten += written;
72 } catch (err) {
73 log.error('usage_rollup_project_failed', {
74 projectId,
75 periodStart: periodStart.toISOString(),
76 message: err instanceof Error ? err.message : String(err),
77 });
78 }
79 }
80
81 log.info('usage_rollup_done', {
82 periodStart: periodStart.toISOString(),
83 periodEnd: periodEnd.toISOString(),
84 projects: activeProjects.length,
85 rowsWritten,
86 });
87 return { hoursProcessed: 1, rowsWritten };
88}
89
90async function rollUpProject(
91 projectId: string,
92 periodStart: Date,
93 periodEnd: Date,
94 connectionSecondsDelta: number | null,
95): Promise<number> {
96 const [invocations, storage] = await Promise.all([
97 getInvocationUsage(projectId, periodStart, periodEnd),
98 // Storage is a gauge — sampled at period_end. We approximate by
99 // sampling "now"; for a one-hour delay between cron fire and run
100 // the drift is negligible at byte-level granularity.
101 getStorageUsage(projectId),
102 ]);
103
104 const rows: { metric: UsageMetric; value: string }[] = [
105 { metric: 'invocations', value: String(invocations.count) },
106 { metric: 'storage_bytes', value: String(storage.bytes) },
107 ];
108 // Only write a connection_seconds row when realtime had activity for
109 // this project this hour. The collectConnectionSecondsDeltas() helper
110 // returns no entry on a cold scrape (api just booted) and we don't
111 // want to over-count by writing zero rows that the Polar push then
112 // tries to deliver.
113 if (connectionSecondsDelta !== null && connectionSecondsDelta > 0) {
114 rows.push({
115 metric: 'connection_seconds',
116 // Round to integer seconds — Polar meters reject non-finite and
117 // some operators configure their meter as integer-typed. Sub-
118 // second precision wouldn't survive Polar's aggregation anyway.
119 value: String(Math.round(connectionSecondsDelta)),
120 });
121 }
122
123 // MAU is a 30-day gauge — only meaningful when this project has auth
124 // enabled (otherwise the `_briven_auth_sessions` table doesn't exist
125 // and the query would error). Skip silently when off; surface a single
126 // log line so operators can spot tenants that turned auth off mid-month.
127 try {
128 if (await isAuthEnabled(projectId)) {
129 const mau = await getAuthMauStats(projectId);
130 rows.push({ metric: 'auth_mau', value: String(mau.count) });
131 // Phase 5.7 — active SSO connection gauge for per-connection pricing.
132 try {
133 const { countActiveSsoConnections } = await import('../services/auth-sso-pricing.js');
134 const ssoCount = await countActiveSsoConnections(projectId);
135 rows.push({ metric: 'auth_sso_connections', value: String(ssoCount) });
136 } catch (ssoErr) {
137 log.warn('usage_rollup_auth_sso_connections_failed', {
138 projectId,
139 message: ssoErr instanceof Error ? ssoErr.message : String(ssoErr),
140 });
141 }
142 }
143 } catch (err) {
144 log.warn('usage_rollup_auth_mau_failed', {
145 projectId,
146 message: err instanceof Error ? err.message : String(err),
147 });
148 }
149
150 const db = getDb();
151 for (const row of rows) {
152 await db
153 .insert(usageEvents)
154 .values({
155 id: newId('au'),
156 projectId,
157 metric: row.metric,
158 periodStart,
159 value: row.value,
160 polarPushStatus: 'pending',
161 })
162 .onConflictDoUpdate({
163 target: [usageEvents.projectId, usageEvents.periodStart, usageEvents.metric],
164 set: {
165 value: row.value,
166 // Conflict means we re-ran an hour. Reset push status so the
167 // updated value goes out to Polar (the previous attempt was
168 // either stale or never pushed).
169 polarPushStatus: 'pending',
170 polarPushedAt: null,
171 },
172 });
173 }
174 return rows.length;
175}
176
177const INTERVAL_MS = HOUR_MS;
178let timer: ReturnType<typeof setInterval> | null = null;
179
180/**
181 * Start the hourly aggregation cron. Calls aggregateUsageForCompletedHour
182 * ~5 minutes after every wall-clock hour boundary so the previous hour's
183 * function_logs are fully durable before we count them. Idempotent — a
184 * second call is a no-op.
185 */
186export function startUsageAggregator(): void {
187 if (timer) return;
188 if (!env.BRIVEN_DATABASE_URL) {
189 log.warn('usage_aggregator_skipped_no_db');
190 return;
191 }
192 const align = () => {
193 const now = new Date();
194 const nextHour = currentHourStart(new Date(now.getTime() + HOUR_MS));
195 const offset = nextHour.getTime() - now.getTime() + 5 * 60 * 1000; // +5min
196 setTimeout(() => {
197 void aggregateUsageForCompletedHour();
198 timer = setInterval(() => {
199 void aggregateUsageForCompletedHour();
200 }, INTERVAL_MS);
201 }, offset).unref?.();
202 };
203 align();
204 log.info('usage_aggregator_armed');
205}
206
207/**
208 * Number of pending usage rows waiting for the Polar push worker.
209 * Exposed for /metrics + the admin "Polar push" dashboard.
210 */
211export async function countPendingPolarPushes(): Promise<number> {
212 const db = getDb();
213 const [row] = await db
214 .select({ n: drizzleSql<number>`count(*)::int` })
215 .from(usageEvents)
216 .where(and(eq(usageEvents.polarPushStatus, 'pending')));
217 return row?.n ?? 0;
218}