| 1 | import type { Env } from "../env.js"; |
| 2 | import { HttpError, now } from "./util.js"; |
| 3 | |
| 4 | /** |
| 5 | * Per-account quotas on what Diffinery runs and keeps for an account. The account that pays is the |
| 6 | * owner of the repository (a person or an organization), whoever started the work. |
| 7 | * |
| 8 | * - hosted_minutes: minutes of hosted Actions jobs and agent sandboxes per calendar month (UTC). |
| 9 | * - hosted_jobs: hosted Actions jobs running at once across the account's repositories; more wait. |
| 10 | * - sandboxes: agent sandboxes open at once across the account's repositories. |
| 11 | * - repositories: repositories the account owns. |
| 12 | * |
| 13 | * Defaults come from the environment (QUOTA_*) and otherwise DEFAULT_QUOTAS; a site administrator |
| 14 | * can override them per account (account_quotas). -1 means unlimited. Accounts run by a site |
| 15 | * administrator (one, or an organization with one as owner) have no quotas, so the instance's own |
| 16 | * work never stops on them. |
| 17 | */ |
| 18 | export type QuotaKey = "hosted_minutes" | "hosted_jobs" | "sandboxes" | "repositories"; |
| 19 | export const QUOTA_KEYS: readonly QuotaKey[] = ["hosted_minutes", "hosted_jobs", "sandboxes", "repositories"]; |
| 20 | export type Quotas = Record<QuotaKey, number>; |
| 21 | |
| 22 | export const DEFAULT_QUOTAS: Quotas = { hosted_minutes: 3000, hosted_jobs: 5, sandboxes: 5, repositories: 500 }; |
| 23 | const ENV_KEYS: Record<QuotaKey, string> = { hosted_minutes: "QUOTA_HOSTED_MINUTES", hosted_jobs: "QUOTA_HOSTED_JOBS", sandboxes: "QUOTA_SANDBOXES", repositories: "QUOTA_REPOSITORIES" }; |
| 24 | |
| 25 | /** |
| 26 | * The longest a job or sandbox can run (a job's 6 hours plus start-up, a sandbox's 24-hour token): |
| 27 | * a row still open after that is counted as ended then, so a lost end can't hold a slot forever. |
| 28 | */ |
| 29 | export const MAX_RUN_MS: Record<"job" | "sandbox", number> = { job: 7 * 3600_000, sandbox: 24 * 3600_000 }; |
| 30 | |
| 31 | const UNLIMITED: Quotas = { hosted_minutes: -1, hosted_jobs: -1, sandboxes: -1, repositories: -1 }; |
| 32 | |
| 33 | function envDefault(env: Env, k: QuotaKey): number { |
| 34 | const raw = (env as unknown as Record<string, string | undefined>)[ENV_KEYS[k]]; |
| 35 | const n = raw === undefined || raw === "" ? NaN : Math.trunc(Number(raw)); |
| 36 | return Number.isFinite(n) && n >= -1 ? n : DEFAULT_QUOTAS[k]; |
| 37 | } |
| 38 | |
| 39 | /** Whether the account is run by a site administrator (itself, or an organization with one as owner). */ |
| 40 | async function adminRun(env: Env, accountId: string): Promise<boolean> { |
| 41 | const r = await env.DB.prepare( |
| 42 | "SELECT 1 FROM accounts a WHERE a.id = ? AND (a.site_admin = 1 OR EXISTS (SELECT 1 FROM org_members m JOIN accounts o ON o.id = m.user_id WHERE m.org_id = a.id AND m.role = 'owner' AND o.site_admin = 1 AND o.status = 'active'))", |
| 43 | ) |
| 44 | .bind(accountId) |
| 45 | .first(); |
| 46 | return !!r; |
| 47 | } |
| 48 | |
| 49 | /** The quotas that apply to an account, and whether they're its own (overridden) or the defaults. */ |
| 50 | export async function quotasFor(env: Env, accountId: string): Promise<Quotas & { exempt: boolean }> { |
| 51 | if (await adminRun(env, accountId)) return { ...UNLIMITED, exempt: true }; |
| 52 | const o = await env.DB.prepare("SELECT * FROM account_quotas WHERE account_id = ?").bind(accountId).first<Record<QuotaKey, number | null>>(); |
| 53 | const q = {} as Quotas; |
| 54 | for (const k of QUOTA_KEYS) q[k] = o?.[k] ?? envDefault(env, k); |
| 55 | return { ...q, exempt: false }; |
| 56 | } |
| 57 | |
| 58 | /** Start of the current calendar month (UTC). */ |
| 59 | export function monthStart(at = new Date()): string { |
| 60 | return new Date(Date.UTC(at.getUTCFullYear(), at.getUTCMonth(), 1)).toISOString(); |
| 61 | } |
| 62 | |
| 63 | export interface Usage { |
| 64 | /** Whole minutes of hosted jobs and sandboxes this month, rounded up per job or sandbox. */ |
| 65 | hosted_minutes: number; |
| 66 | hosted_jobs: number; |
| 67 | sandboxes: number; |
| 68 | repositories: number; |
| 69 | } |
| 70 | |
| 71 | /** What an account uses now. */ |
| 72 | export async function usageOf(env: Env, accountId: string, at = new Date()): Promise<Usage> { |
| 73 | const t = at.getTime(); |
| 74 | const month = monthStart(at); |
| 75 | // Anything that started earlier ended before the month began (or is counted as ended by then). |
| 76 | const rows = await env.DB.prepare("SELECT kind, started_at, ended_at FROM compute_usage WHERE account_id = ? AND started_at >= ?") |
| 77 | .bind(accountId, new Date(Math.min(Date.parse(month), t) - MAX_RUN_MS.sandbox).toISOString()) |
| 78 | .all<{ kind: "job" | "sandbox"; started_at: string; ended_at: string | null }>(); |
| 79 | let minutes = 0; |
| 80 | let jobs = 0; |
| 81 | let sandboxes = 0; |
| 82 | for (const r of rows.results) { |
| 83 | const start = Date.parse(r.started_at); |
| 84 | const cap = start + MAX_RUN_MS[r.kind]; |
| 85 | const end = Math.min(r.ended_at ? Date.parse(r.ended_at) : t, cap); |
| 86 | const from = Math.max(start, Date.parse(month)); |
| 87 | if (end > from) minutes += Math.ceil((end - from) / 60_000); |
| 88 | if (!r.ended_at && t < cap) { |
| 89 | if (r.kind === "job") jobs++; |
| 90 | else sandboxes++; |
| 91 | } |
| 92 | } |
| 93 | const repos = await env.DB.prepare("SELECT count(*) AS n FROM repos WHERE owner_id = ?").bind(accountId).first<{ n: number }>(); |
| 94 | return { hosted_minutes: minutes, hosted_jobs: jobs, sandboxes, repositories: repos?.n ?? 0 }; |
| 95 | } |
| 96 | |
| 97 | const within = (limit: number, used: number) => limit < 0 || used < limit; |
| 98 | |
| 99 | /** |
| 100 | * Whether the account may start a hosted job or a sandbox now: "ok", "wait" (a job waits for one of |
| 101 | * the account's running jobs to end), or an error message (minutes used up, or too many sandboxes). |
| 102 | */ |
| 103 | export async function mayStart(env: Env, accountId: string, kind: "job" | "sandbox"): Promise<{ ok: true } | { ok: false; wait: boolean; message: string }> { |
| 104 | const q = await quotasFor(env, accountId); |
| 105 | if (q.exempt) return { ok: true }; |
| 106 | const u = await usageOf(env, accountId); |
| 107 | if (!within(q.hosted_minutes, u.hosted_minutes)) |
| 108 | return { ok: false, wait: false, message: `this account has used its ${q.hosted_minutes} hosted minutes for this month (Actions jobs on hosted runners and agent sandboxes); they renew on the 1st (UTC). Self-hosted runners still work` }; |
| 109 | if (kind === "job" && !within(q.hosted_jobs, u.hosted_jobs)) return { ok: false, wait: true, message: `waiting for one of this account's ${q.hosted_jobs} hosted jobs to finish` }; |
| 110 | if (kind === "sandbox" && !within(q.sandboxes, u.sandboxes)) return { ok: false, wait: false, message: `this account already has ${q.sandboxes} agent sandboxes open; close one or wait for one to finish` }; |
| 111 | return { ok: true }; |
| 112 | } |
| 113 | |
| 114 | /** Record that a hosted job or sandbox started; `id` is the job or sandbox id. */ |
| 115 | export async function startCompute(env: Env, row: { id: string; accountId: string; repoId: string; kind: "job" | "sandbox" }): Promise<void> { |
| 116 | await env.DB.prepare("INSERT OR REPLACE INTO compute_usage (id, account_id, repo_id, kind, started_at) VALUES (?, ?, ?, ?, ?)").bind(row.id, row.accountId, row.repoId, row.kind, now()).run(); |
| 117 | } |
| 118 | |
| 119 | /** Record that a hosted job or sandbox ended (idempotent). */ |
| 120 | export async function endCompute(env: Env, id: string): Promise<void> { |
| 121 | await env.DB.prepare("UPDATE compute_usage SET ended_at = ? WHERE id = ? AND ended_at IS NULL").bind(now(), id).run(); |
| 122 | } |
| 123 | |
| 124 | /** Refuse a new repository past the owner's quota. */ |
| 125 | export async function checkRepositoryQuota(env: Env, ownerId: string, ownerLogin: string): Promise<void> { |
| 126 | const q = await quotasFor(env, ownerId); |
| 127 | if (q.exempt || q.repositories < 0) return; |
| 128 | const n = (await env.DB.prepare("SELECT count(*) AS n FROM repos WHERE owner_id = ?").bind(ownerId).first<{ n: number }>())?.n ?? 0; |
| 129 | if (n >= q.repositories) throw new HttpError(403, `${ownerLogin} has reached its limit of ${q.repositories} repositories; delete one or ask a site administrator to raise it`); |
| 130 | } |
| 131 | |
| 132 | /** A site administrator's overrides for an account; null clears one back to the default. */ |
| 133 | export async function setQuotas(env: Env, accountId: string, patch: Partial<Record<QuotaKey, number | null>>, by: string): Promise<void> { |
| 134 | for (const [k, v] of Object.entries(patch)) { |
| 135 | if (!QUOTA_KEYS.includes(k as QuotaKey)) throw new HttpError(422, `unknown quota ${k}; use ${QUOTA_KEYS.join(", ")}`); |
| 136 | if (v !== null && (typeof v !== "number" || !Number.isInteger(v) || v < -1 || v > 1e9)) throw new HttpError(422, `${k} must be a whole number, -1 for unlimited, or null for the default`); |
| 137 | } |
| 138 | const cur = (await env.DB.prepare("SELECT * FROM account_quotas WHERE account_id = ?").bind(accountId).first<Record<QuotaKey, number | null>>()) ?? null; |
| 139 | const next = Object.fromEntries(QUOTA_KEYS.map((k) => [k, k in patch ? (patch[k] ?? null) : (cur?.[k] ?? null)])) as Record<QuotaKey, number | null>; |
| 140 | await env.DB.prepare( |
| 141 | "INSERT INTO account_quotas (account_id, hosted_minutes, hosted_jobs, sandboxes, repositories, updated_at, updated_by) VALUES (?, ?, ?, ?, ?, ?, ?) ON CONFLICT(account_id) DO UPDATE SET hosted_minutes = excluded.hosted_minutes, hosted_jobs = excluded.hosted_jobs, sandboxes = excluded.sandboxes, repositories = excluded.repositories, updated_at = excluded.updated_at, updated_by = excluded.updated_by", |
| 142 | ) |
| 143 | .bind(accountId, next.hosted_minutes, next.hosted_jobs, next.sandboxes, next.repositories, now(), by) |
| 144 | .run(); |
| 145 | } |