147 lines · 5.3 KB
import { env } from "cloudflare:workers";
import { all, first } from "../db";
import { SCRATCH_LIMITS } from "./config";
import { meterSiteRequest } from "./metering";
import type { SiteMeta } from "./config";

const FLAG_TTL_SECONDS = 2 * 24 * 60 * 60;
const HTTP_SWITCHING = 101;

export type UsageDelta = { requests?: number; cpuMs?: number; bytesOut?: number };

export type AccountUsage = {
  day: string;
  requests: number;
  cpuMs: number;
  bytesOut: number;
};

export function usageDay(now = Date.now()): string {
  return new Date(now).toISOString().slice(0, 10);
}

function overKey(siteId: string, day: string): string {
  return `sites:over:${siteId}:${day}`;
}

/**
 * Adds to a site's daily usage row. Unclaimed scratch sites (no account) are the free tier: once a daily cap is
 * crossed a KV flag makes the dispatcher refuse further requests until the next UTC day.
 */
export async function recordUsage(
  meta: Pick<SiteMeta, "id" | "ownerId" | "scratch">,
  delta: UsageDelta,
): Promise<void> {
  const day = usageDay();
  const row = await first<{ requests: number; cpu_ms: number; bytes_out: number }>(
    `INSERT INTO site_usage (site_id, day, account_id, requests, cpu_ms, bytes_out) VALUES (?, ?, ?, ?, ?, ?)
     ON CONFLICT (site_id, day) DO UPDATE SET requests = requests + excluded.requests, cpu_ms = cpu_ms + excluded.cpu_ms,
       bytes_out = bytes_out + excluded.bytes_out, account_id = excluded.account_id
     RETURNING requests, cpu_ms, bytes_out`,