usage-aggregator.ts218 lines · main
| 1 | import { and, eq, isNull, sql as drizzleSql } from 'drizzle-orm'; |
| 2 | |
| 3 | import { newId } from '@briven/shared'; |
| 4 | |
| 5 | import { getDb } from '../db/client.js'; |
| 6 | import { projects, usageEvents, type UsageMetric } from '../db/schema.js'; |
| 7 | import { env } from '../env.js'; |
| 8 | import { log } from '../lib/logger.js'; |
| 9 | import { getAuthMauStats } from '../services/auth-mau.js'; |
| 10 | import { collectConnectionSecondsDeltas } from '../services/connection-seconds.js'; |
| 11 | import { isAuthEnabled } from '../services/tenant-config-store.js'; |
| 12 | import { 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 | |
| 31 | const HOUR_MS = 60 * 60 * 1000; |
| 32 | |
| 33 | /** First millisecond of the UTC hour containing `now`. */ |
| 34 | export 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 | */ |
| 43 | export 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 | |
| 90 | async 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 | |
| 177 | const INTERVAL_MS = HOUR_MS; |
| 178 | let 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 | */ |
| 186 | export 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 | */ |
| 211 | export 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 | } |