diff --git a/apps/sim/lib/billing/storage/index.ts b/apps/sim/lib/billing/storage/index.ts index 9203c24b455..8820af428ba 100644 --- a/apps/sim/lib/billing/storage/index.ts +++ b/apps/sim/lib/billing/storage/index.ts @@ -1,6 +1,8 @@ export { checkStorageQuota, getUserStorageLimit, getUserStorageUsage } from './limits' export { + checkAndIncrementStorageUsageInTx, decrementStorageUsage, decrementStorageUsageInTx, incrementStorageUsage, + maybeNotifyStorageLimit, } from './tracking' diff --git a/apps/sim/lib/billing/storage/tracking.ts b/apps/sim/lib/billing/storage/tracking.ts index 6311a881846..3c31bb8eb3b 100644 --- a/apps/sim/lib/billing/storage/tracking.ts +++ b/apps/sim/lib/billing/storage/tracking.ts @@ -7,11 +7,12 @@ import { db } from '@sim/db' import { organization, userStats } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { eq, sql } from 'drizzle-orm' +import { generateId } from '@sim/utils/id' +import { and, eq, sql } from 'drizzle-orm' import { maybeNotifyLimit } from '@/lib/billing/core/limit-notifications' import type { HighestPrioritySubscription } from '@/lib/billing/core/plan' import { getUserStorageLimit, getUserStorageUsage } from '@/lib/billing/storage/limits' -import { isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils' +import { getFreeTierLimit, isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils' import { isBillingEnabled } from '@/lib/core/config/env-flags' const logger = createLogger('StorageTracking') @@ -33,7 +34,7 @@ function formatGb(bytes: number, decimals: number): string { * @param rearmOnly - True on decrements, so a shrink that leaves usage above a * threshold re-arms but never sends (a drop is not a fresh crossing). */ -async function maybeNotifyStorageLimit( +export async function maybeNotifyStorageLimit( userId: string, workspaceId: string, sub: HighestPrioritySubscription | null, @@ -195,3 +196,94 @@ export async function decrementStorageUsageInTx( .where(eq(userStats.userId, userId)) } } + +/** + * Atomically check quota and increment a user's (or their org's) storage + * counter inside an existing transaction, using a pre-resolved subscription. + * The check and the increment are a single conditional `UPDATE`, so two + * concurrent callers can no longer both read the same pre-increment usage, + * both pass the check, and both commit past the limit — the second caller's + * `UPDATE` re-evaluates the WHERE clause against the first caller's already + * -committed-within-the-same-DB-round-trip row and correctly fails. Replaces + * the old read-then-decide-then-increment-after-commit split (`checkStorageQuota` + * + a fire-and-forget `incrementStorageUsage` after the transaction), which left + * a window between the read and the increment. + * + * On success, callers should best-effort call {@link maybeNotifyStorageLimit} + * after the transaction commits (mirrors the existing post-increment threshold + * check) — this helper doesn't do it itself since it runs mid-transaction. + * + * For a personal (non-org-scoped) `userId`, this first upserts the + * `userStats` row on `tx` — a documented possibility for OAuth account + * linking (see `ensureUserStatsExists` in `lib/billing/core/usage.ts`, whose + * insert values this mirrors) — because the conditional `UPDATE` below + * matches 0 rows, and therefore reads as "quota exceeded", if that row + * doesn't exist yet. `ensureUserStatsExists` itself isn't reused here since + * it writes through the standalone `db` client, which would open a second + * pooled connection while this transaction's is held. + */ +export async function checkAndIncrementStorageUsageInTx( + tx: StorageTransaction, + sub: HighestPrioritySubscription | null, + userId: string, + bytes: number +): Promise<{ allowed: boolean; currentUsage: number; limit: number; error?: string }> { + if (!isBillingEnabled) { + return { allowed: true, currentUsage: 0, limit: Number.MAX_SAFE_INTEGER } + } + + const limit = await getUserStorageLimit(userId, sub) + + if (bytes <= 0) { + return { allowed: true, currentUsage: await getUserStorageUsage(userId, sub), limit } + } + + const orgScoped = isOrgScopedSubscription(sub, userId) && sub + + if (!orgScoped) { + await tx + .insert(userStats) + .values({ + id: generateId(), + userId, + currentUsageLimit: getFreeTierLimit().toString(), + usageLimitUpdatedAt: new Date(), + }) + .onConflictDoNothing({ target: userStats.userId }) + } + + const [updated] = orgScoped + ? await tx + .update(organization) + .set({ storageUsedBytes: sql`${organization.storageUsedBytes} + ${bytes}` }) + .where( + and( + eq(organization.id, sub.referenceId), + sql`${organization.storageUsedBytes} + ${bytes} <= ${limit}` + ) + ) + .returning({ storageUsedBytes: organization.storageUsedBytes }) + : await tx + .update(userStats) + .set({ storageUsedBytes: sql`${userStats.storageUsedBytes} + ${bytes}` }) + .where( + and( + eq(userStats.userId, userId), + sql`${userStats.storageUsedBytes} + ${bytes} <= ${limit}` + ) + ) + .returning({ storageUsedBytes: userStats.storageUsedBytes }) + + if (updated) { + return { allowed: true, currentUsage: updated.storageUsedBytes - bytes, limit } + } + + const currentUsage = await getUserStorageUsage(userId, sub) + const newUsage = currentUsage + bytes + return { + allowed: false, + currentUsage, + limit, + error: `Storage limit exceeded. Used: ${(newUsage / (1024 * 1024 * 1024)).toFixed(2)}GB, Limit: ${(limit / (1024 * 1024 * 1024)).toFixed(0)}GB`, + } +} diff --git a/apps/sim/lib/knowledge/chunks/service.ts b/apps/sim/lib/knowledge/chunks/service.ts index fe34a40b543..6660027d470 100644 --- a/apps/sim/lib/knowledge/chunks/service.ts +++ b/apps/sim/lib/knowledge/chunks/service.ts @@ -17,6 +17,8 @@ import { estimateTokenCount } from '@/lib/tokenization/estimators' const logger = createLogger('ChunksService') +const KB_CHUNK_LOCK_TIMEOUT_MS = 5_000 + /** * Query chunks for a document with filtering and pagination */ @@ -101,7 +103,23 @@ export async function queryChunks( } /** - * Create a new chunk for a document + * Create a new chunk for a document. + * + * Assigns `chunkIndex` as `max(chunkIndex) + 1` under a transactional + * `pg_advisory_xact_lock` keyed on the document, so concurrent calls for the + * same document serialize instead of computing the same index and colliding + * on the `(document_id, chunk_index)` unique constraint. A `SELECT ... FOR + * UPDATE` on the current max row doesn't prevent that collision, since the + * row it would lock is unrelated to the not-yet-inserted next row. A row + * lock on `document` instead of an advisory lock would also work, but would + * invert the embedding-before-document lock order every other chunk + * mutation path uses (see lock-order.test.ts) — the advisory lock is a + * separate namespace, so it can't deadlock against that convention. + * + * `pg_advisory_xact_lock` auto-releases at transaction end, so there's no + * session lock to leak onto a pooled connection, and `lock_timeout` bounds + * the wait (it raises SQLSTATE 55P03 instead of hanging a pooled connection) + * if a same-document holder is stuck. */ export async function createChunk( knowledgeBaseId: string, @@ -135,8 +153,14 @@ export async function createChunk( const chunkId = generateId() const now = new Date() - // Use transaction to atomically get next index and insert chunk const newChunk = await db.transaction(async (tx) => { + await tx.execute( + sql`select set_config('lock_timeout', ${`${KB_CHUNK_LOCK_TIMEOUT_MS}ms`}, true)` + ) + await tx.execute( + sql`select pg_advisory_xact_lock(hashtextextended(${`kb_chunk_seq:${documentId}`}, 0))` + ) + const activeDocument = await tx .select({ id: document.id }) .from(document) @@ -156,7 +180,6 @@ export async function createChunk( throw new Error('Document not found') } - // Get the next chunk index atomically within the transaction const lastChunk = await tx .select({ chunkIndex: embedding.chunkIndex }) .from(embedding) diff --git a/apps/sim/lib/knowledge/documents/service.ts b/apps/sim/lib/knowledge/documents/service.ts index 56300b1d837..88b33ff7a62 100644 --- a/apps/sim/lib/knowledge/documents/service.ts +++ b/apps/sim/lib/knowledge/documents/service.ts @@ -18,9 +18,9 @@ import type { HighestPrioritySubscription } from '@/lib/billing/core/plan' import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription' import { recordUsage } from '@/lib/billing/core/usage-log' import { - checkStorageQuota, + checkAndIncrementStorageUsageInTx, decrementStorageUsageInTx, - incrementStorageUsage, + maybeNotifyStorageLimit, } from '@/lib/billing/storage' import { checkAndBillOverageThreshold } from '@/lib/billing/threshold-billing' import type { ChunkingStrategy, StrategyOptions } from '@/lib/chunkers/types' @@ -864,6 +864,33 @@ export function isTriggerAvailable(): boolean { return Boolean(env.TRIGGER_SECRET_KEY) && isTriggerDevEnabled } +/** + * Resolves the subscription to bill for a storage-metered upload, ahead of + * the FOR UPDATE transaction that will insert the document row. + * + * Must run before that transaction opens: `knowledgeBase.userId` is + * immutable once a KB is created, so resolving it here can't go stale + * relative to the transaction's own read of it, and resolving it inside the + * transaction would open a second pooled-database-connection checkout while + * the first is held. + */ +async function resolveQuotaSubscription( + knowledgeBaseId: string, + uploadedBy: string | null +): Promise { + const billedUserId = + uploadedBy ?? + ( + await db + .select({ userId: knowledgeBase.userId }) + .from(knowledgeBase) + .where(eq(knowledgeBase.id, knowledgeBaseId)) + .limit(1) + )[0]?.userId + + return billedUserId ? getHighestPrioritySubscription(billedUserId) : null +} + export async function createDocumentRecords( documents: Array<{ filename: string @@ -883,7 +910,15 @@ export async function createDocumentRecords( requestId: string, uploadedBy: string | null = null ): Promise { - let storageBilling: { userId: string; workspaceId: string | null; bytes: number } | null = null + let storageBilling: { + userId: string + workspaceId: string | null + bytes: number + sub: HighestPrioritySubscription | null + } | null = null + + const totalBytes = documents.reduce((sum, docData) => sum + (docData.fileSize || 0), 0) + const sub = totalBytes > 0 ? await resolveQuotaSubscription(knowledgeBaseId, uploadedBy) : null const returnData = await db.transaction(async (tx) => { await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR UPDATE`) @@ -912,13 +947,12 @@ export async function createDocumentRecords( ) const billedUserId = uploadedBy ?? kb[0].userId - const totalBytes = documents.reduce((sum, docData) => sum + (docData.fileSize || 0), 0) if (totalBytes > 0) { - const quotaCheck = await checkStorageQuota(billedUserId, totalBytes) + const quotaCheck = await checkAndIncrementStorageUsageInTx(tx, sub, billedUserId, totalBytes) if (!quotaCheck.allowed) { throw new Error(quotaCheck.error || 'Storage limit exceeded') } - storageBilling = { userId: billedUserId, workspaceId: kbWorkspaceId, bytes: totalBytes } + storageBilling = { userId: billedUserId, workspaceId: kbWorkspaceId, bytes: totalBytes, sub } } // One load per batch (was N+1); skip entirely if no doc carries tags. @@ -1011,11 +1045,14 @@ export async function createDocumentRecords( }) if (storageBilling) { - const billing: { userId: string; workspaceId: string | null; bytes: number } = storageBilling - try { - await incrementStorageUsage(billing.userId, billing.bytes, billing.workspaceId ?? undefined) - } catch (storageError) { - logger.error(`[${requestId}] Failed to update storage tracking:`, storageError) + const billing: { + userId: string + workspaceId: string | null + bytes: number + sub: HighestPrioritySubscription | null + } = storageBilling + if (billing.workspaceId) { + void maybeNotifyStorageLimit(billing.userId, billing.workspaceId, billing.sub) } } @@ -1338,7 +1375,15 @@ export async function createSingleDocument( ...processedTags, } - let storageBilling: { userId: string; workspaceId: string | null; bytes: number } | null = null + let storageBilling: { + userId: string + workspaceId: string | null + bytes: number + sub: HighestPrioritySubscription | null + } | null = null + + const sub = + documentData.fileSize > 0 ? await resolveQuotaSubscription(knowledgeBaseId, uploadedBy) : null await db.transaction(async (tx) => { await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR UPDATE`) @@ -1367,7 +1412,12 @@ export async function createSingleDocument( const billedUserId = uploadedBy ?? kb[0].userId if (documentData.fileSize > 0) { - const quotaCheck = await checkStorageQuota(billedUserId, documentData.fileSize) + const quotaCheck = await checkAndIncrementStorageUsageInTx( + tx, + sub, + billedUserId, + documentData.fileSize + ) if (!quotaCheck.allowed) { throw new Error(quotaCheck.error || 'Storage limit exceeded') } @@ -1375,6 +1425,7 @@ export async function createSingleDocument( userId: billedUserId, workspaceId: kb[0].workspaceId, bytes: documentData.fileSize, + sub, } } @@ -1387,11 +1438,14 @@ export async function createSingleDocument( }) if (storageBilling) { - const billing: { userId: string; workspaceId: string | null; bytes: number } = storageBilling - try { - await incrementStorageUsage(billing.userId, billing.bytes, billing.workspaceId ?? undefined) - } catch (storageError) { - logger.error(`[${requestId}] Failed to update storage tracking:`, storageError) + const billing: { + userId: string + workspaceId: string | null + bytes: number + sub: HighestPrioritySubscription | null + } = storageBilling + if (billing.workspaceId) { + void maybeNotifyStorageLimit(billing.userId, billing.workspaceId, billing.sub) } }