feat: implement storage quota management for uploads

- Added storage quota enforcement for audio and image uploads in the respective routes. - Introduced reservation system to manage concurrent uploads and prevent quota overages. - Enhanced comment creation to account for audio and image attachment sizes against user quotas. - Created new UploadReservation model to track in-flight upload reservations. - Backfilled existing video assets with size information from R2. - Added progress component for UI feedback during uploads. - Updated API responses to include reservation IDs for better quota management. - Adjusted error handling to return appropriate storage limit exceeded messages.
2026-09-11 17:46:06 +00:00 · 2026-04-15 19:53:43 +03:00
parent acf31b3d6b
commit 873945464d
21 changed files with 753 additions and 85 deletions
@@ -0,0 +1,176 @@
+import type { NextResponse } from 'next/server';
+import { db } from '@/lib/db';
+import { apiErrors } from '@/lib/api-response';
+import { isStripeFeatureEnabled } from '@/lib/feature-flags';
+import { getCachedUserBunnyStorage } from '@/lib/admin-stats';
+
+// 200 GB expressed in bytes
+export const PLAN_STORAGE_LIMIT_BYTES = BigInt(200) * BigInt(1024) * BigInt(1024) * BigInt(1024);
+
+// TTL for upload reservations: 30 minutes is enough for R2 image/audio uploads
+const RESERVATION_TTL_MS = 30 * 60 * 1000;
+
+// Sentinel error thrown inside a Prisma transaction to signal quota exceeded
+class QuotaExceededError extends Error {}
+
+/**
+ * Returns total bytes used by a given billed user across R2 (image + audio),
+ * Bunny Stream, and any active (non-expired) upload reservations.
+ * Uses the cached Bunny stats (10-min TTL) to avoid calling the Bunny API on
+ * every upload.
+ */
+export async function getUserTotalStorageBytes(userId: string): Promise<bigint> {
+  const [r2Rows, bunnyByUser, reservationRows] = await Promise.all([
+    db.$queryRaw<[{ total: bigint }]>`
+      SELECT COALESCE(SUM(size_bytes), 0)::bigint AS total
+      FROM video_assets
+      WHERE "billedUserId" = ${userId}
+        AND provider IN ('R2_IMAGE', 'R2_AUDIO')
+    `,
+    getCachedUserBunnyStorage(),
+    db.$queryRaw<[{ total: bigint }]>`
+      SELECT COALESCE(SUM("sizeBytes"), 0)::bigint AS total
+      FROM upload_reservations
+      WHERE "billedUserId" = ${userId}
+        AND "expiresAt" > NOW()
+    `,
+  ]);
+
+  const r2Bytes = r2Rows[0]?.total ?? BigInt(0);
+  const bunnyBytes = BigInt(bunnyByUser[userId] ?? 0);
+  const reservedBytes = reservationRows[0]?.total ?? BigInt(0);
+
+  return r2Bytes + bunnyBytes + reservedBytes;
+}
+
+/**
+ * Returns storage usage info for a user in a UI-friendly shape.
+ */
+export async function getUserStorageInfo(userId: string): Promise<{
+  usedBytes: bigint;
+  limitBytes: bigint;
+  percentage: number;
+}> {
+  const usedBytes = await getUserTotalStorageBytes(userId);
+  const limitBytes = PLAN_STORAGE_LIMIT_BYTES;
+  const percentage = limitBytes > BigInt(0)
+    ? Math.min(100, Number((usedBytes * BigInt(10000)) / limitBytes) / 100)
+    : 0;
+
+  return { usedBytes, limitBytes, percentage };
+}
+
+/**
+ * Checks whether the user can upload `incomingSizeBytes` more data.
+ *
+ * Returns a 507 response if the quota would be exceeded, or `null` if the
+ * upload is allowed. When Stripe is disabled the check is always skipped so
+ * self-hosted instances without billing still work.
+ *
+ * Uses `>=` so a user at exactly the limit cannot initiate new uploads.
+ */
+export async function enforceStorageQuota(
+  userId: string,
+  incomingSizeBytes: bigint,
+): Promise<NextResponse | null> {
+  if (!isStripeFeatureEnabled()) {
+    return null;
+  }
+
+  const usedBytes = await getUserTotalStorageBytes(userId);
+
+  if (usedBytes + incomingSizeBytes >= PLAN_STORAGE_LIMIT_BYTES) {
+    return apiErrors.storageExceeded() as NextResponse;
+  }
+
+  return null;
+}
+
+/**
+ * Atomically checks the quota and records an in-flight upload reservation.
+ *
+ * Uses a PostgreSQL advisory transaction lock (per user) so concurrent callers
+ * are serialised: the second request sees the first reservation in the sum and
+ * cannot double-book the same headroom.
+ *
+ * Returns `{ reservationId }` on success or `{ error }` (a 507 NextResponse)
+ * when the quota would be exceeded. Call `releaseStorageReservation` to delete
+ * the reservation once the paired asset is committed (or if the upload fails).
+ *
+ * When Stripe is disabled the check is skipped and `reservationId` is `null`.
+ */
+export async function reserveStorageQuota(
+  userId: string,
+  incomingSizeBytes: bigint,
+): Promise<{ reservationId: string | null } | { error: NextResponse }> {
+  if (!isStripeFeatureEnabled()) {
+    return { reservationId: null };
+  }
+
+  const expiresAt = new Date(Date.now() + RESERVATION_TTL_MS);
+
+  // Fetch Bunny storage BEFORE entering the transaction to avoid holding the
+  // advisory lock during a potentially slow/failing HTTP call on cache miss.
+  const bunnyData = await getCachedUserBunnyStorage();
+  const bunnyBytes = BigInt(bunnyData[userId] ?? 0);
+
+  try {
+    const reservationId = await db.$transaction(async (tx) => {
+      // Serialise quota checks for this user via a per-user advisory lock.
+      // Combine two 32-bit hashtext() halves into a single 64-bit bigint to
+      // eliminate the 32-bit hash-space collision risk of plain hashtext().
+      // Use $executeRaw — the function returns void which $queryRaw cannot deserialize.
+      await tx.$executeRaw`
+        SELECT pg_advisory_xact_lock(
+          ('x' || left(md5(${userId}), 16))::bit(64)::bigint
+        )
+      `;
+
+      // Read committed R2 storage under the lock
+      const [r2Row] = await tx.$queryRaw<[{ total: bigint }]>`
+        SELECT COALESCE(SUM(size_bytes), 0)::bigint AS total
+        FROM video_assets
+        WHERE "billedUserId" = ${userId}
+          AND provider IN ('R2_IMAGE', 'R2_AUDIO')
+      `;
+      const r2Bytes = r2Row?.total ?? BigInt(0);
+
+      // Read active (non-expired) reservations under the same lock
+      const [resRow] = await tx.$queryRaw<[{ total: bigint }]>`
+        SELECT COALESCE(SUM("sizeBytes"), 0)::bigint AS total
+        FROM upload_reservations
+        WHERE "billedUserId" = ${userId}
+          AND "expiresAt" > NOW()
+      `;
+      const reservedBytes = resRow?.total ?? BigInt(0);
+
+      const totalUsed = r2Bytes + reservedBytes + bunnyBytes;
+      if (totalUsed + incomingSizeBytes >= PLAN_STORAGE_LIMIT_BYTES) {
+        throw new QuotaExceededError();
+      }
+
+      const reservation = await tx.uploadReservation.create({
+        data: { billedUserId: userId, sizeBytes: incomingSizeBytes, expiresAt },
+        select: { id: true },
+      });
+
+      return reservation.id;
+    });
+
+    return { reservationId };
+  } catch (e) {
+    if (e instanceof QuotaExceededError) {
+      return { error: apiErrors.storageExceeded() as NextResponse };
+    }
+    throw e;
+  }
+}
+
+/**
+ * Deletes an upload reservation created by `reserveStorageQuota`.
+ * Safe to call with `null` (no-op) for flows where billing is disabled.
+ */
+export async function releaseStorageReservation(reservationId: string | null): Promise<void> {
+  if (!reservationId) return;
+  await db.uploadReservation.deleteMany({ where: { id: reservationId } });
+}