|
| 1 | +/** |
| 2 | + * Completion ordering against real PostgreSQL: a run's usage ledger is written before its |
| 3 | + * log reads terminal, so a reader that sees a finished run always sees its cost. |
| 4 | + */ |
| 5 | +import { db } from '@sim/db' |
| 6 | +import { |
| 7 | + usageLog, |
| 8 | + user, |
| 9 | + workflow, |
| 10 | + workflowExecutionLogs, |
| 11 | + workflowExecutionSnapshots, |
| 12 | + workspace, |
| 13 | +} from '@sim/db/schema' |
| 14 | +import { createDeferred } from '@sim/testing' |
| 15 | +import { generateId } from '@sim/utils/id' |
| 16 | +import { eq, sql } from 'drizzle-orm' |
| 17 | +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' |
| 18 | +import { resolveBillingAttribution } from '@/lib/billing/core/billing-attribution' |
| 19 | +import { acquireAdvisoryXactLock } from '@/lib/db/advisory-locks' |
| 20 | +import { buildCostLedger } from '@/lib/logs/cost-ledger' |
| 21 | +import { executionLogger } from '@/lib/logs/execution/logger' |
| 22 | +import { calculateCostSummary } from '@/lib/logs/execution/logging-factory' |
| 23 | +import type { WorkflowState } from '@/lib/logs/types' |
| 24 | + |
| 25 | +const ids = { |
| 26 | + owner: `ledger-order-owner-${generateId()}`, |
| 27 | + workspace: generateId(), |
| 28 | + workflow: generateId(), |
| 29 | +} |
| 30 | + |
| 31 | +/** The advisory lock the usage ledger write takes for its execution before inserting. */ |
| 32 | +const USAGE_RECONCILE_LOCK = 'execution_usage_reconcile' |
| 33 | +const EXECUTION_FEE = 0.005 |
| 34 | + |
| 35 | +const workflowState: WorkflowState = { |
| 36 | + blocks: { |
| 37 | + start: { |
| 38 | + id: 'start', |
| 39 | + type: 'starter', |
| 40 | + name: 'Start', |
| 41 | + position: { x: 0, y: 0 }, |
| 42 | + subBlocks: {}, |
| 43 | + outputs: {}, |
| 44 | + enabled: true, |
| 45 | + }, |
| 46 | + }, |
| 47 | + edges: [], |
| 48 | + loops: {}, |
| 49 | + parallels: {}, |
| 50 | +} |
| 51 | + |
| 52 | +async function startExecution(executionId: string) { |
| 53 | + await executionLogger.startWorkflowExecution({ |
| 54 | + workflowId: ids.workflow, |
| 55 | + workspaceId: ids.workspace, |
| 56 | + executionId, |
| 57 | + trigger: { type: 'api', source: 'api', timestamp: new Date().toISOString() }, |
| 58 | + environment: { |
| 59 | + variables: {}, |
| 60 | + workflowId: ids.workflow, |
| 61 | + executionId, |
| 62 | + userId: ids.owner, |
| 63 | + workspaceId: ids.workspace, |
| 64 | + }, |
| 65 | + workflowState, |
| 66 | + }) |
| 67 | +} |
| 68 | + |
| 69 | +async function logRow(executionId: string) { |
| 70 | + const [row] = await db |
| 71 | + .select({ status: workflowExecutionLogs.status, costTotal: workflowExecutionLogs.costTotal }) |
| 72 | + .from(workflowExecutionLogs) |
| 73 | + .where(eq(workflowExecutionLogs.executionId, executionId)) |
| 74 | + return row |
| 75 | +} |
| 76 | + |
| 77 | +beforeAll(async () => { |
| 78 | + const now = new Date() |
| 79 | + await db.insert(user).values({ |
| 80 | + id: ids.owner, |
| 81 | + name: 'Ledger Order', |
| 82 | + email: `${ids.owner}@ledger-order.test`, |
| 83 | + emailVerified: true, |
| 84 | + createdAt: now, |
| 85 | + updatedAt: now, |
| 86 | + }) |
| 87 | + await db.insert(workspace).values({ |
| 88 | + id: ids.workspace, |
| 89 | + name: 'Ledger Order', |
| 90 | + ownerId: ids.owner, |
| 91 | + billedAccountUserId: ids.owner, |
| 92 | + }) |
| 93 | + await db.insert(workflow).values({ |
| 94 | + id: ids.workflow, |
| 95 | + userId: ids.owner, |
| 96 | + workspaceId: ids.workspace, |
| 97 | + name: 'Ledger Order', |
| 98 | + lastSynced: now, |
| 99 | + createdAt: now, |
| 100 | + updatedAt: now, |
| 101 | + }) |
| 102 | +}) |
| 103 | + |
| 104 | +afterAll(async () => { |
| 105 | + await db.delete(usageLog).where(eq(usageLog.workflowId, ids.workflow)) |
| 106 | + await db.delete(workflowExecutionLogs).where(eq(workflowExecutionLogs.workflowId, ids.workflow)) |
| 107 | + await db |
| 108 | + .delete(workflowExecutionSnapshots) |
| 109 | + .where(eq(workflowExecutionSnapshots.workflowId, ids.workflow)) |
| 110 | + await db.delete(workspace).where(eq(workspace.id, ids.workspace)) |
| 111 | + await db.delete(user).where(eq(user.id, ids.owner)) |
| 112 | +}) |
| 113 | + |
| 114 | +/** |
| 115 | + * Whether a session waits on this execution's ledger lock. A bigint advisory key is stored |
| 116 | + * split across `classid` (high 32 bits) and `objid` (low 32 bits) with `objsubid = 1`. |
| 117 | + */ |
| 118 | +async function isLedgerLockAwaited(executionId: string) { |
| 119 | + const rows = await db.execute<{ waiting: boolean }>(sql` |
| 120 | + SELECT EXISTS ( |
| 121 | + SELECT 1 FROM pg_locks |
| 122 | + WHERE locktype = 'advisory' AND NOT granted AND objsubid = 1 |
| 123 | + AND ((classid::bigint << 32) | objid::bigint) = hashtextextended(${executionId}, 0) |
| 124 | + ) AS waiting |
| 125 | + `) |
| 126 | + return Boolean(rows[0]?.waiting) |
| 127 | +} |
| 128 | + |
| 129 | +describe('completeWorkflowExecution', () => { |
| 130 | + it('writes the cost ledger before the run reads finished', async () => { |
| 131 | + const billingAttribution = await resolveBillingAttribution({ |
| 132 | + actorUserId: ids.owner, |
| 133 | + workspaceId: ids.workspace, |
| 134 | + }) |
| 135 | + const executionId = generateId() |
| 136 | + await startExecution(executionId) |
| 137 | + |
| 138 | + /** Holds the ledger write at its lock, so the log can be read while it waits there. */ |
| 139 | + const lockHeld = createDeferred<void>() |
| 140 | + const releaseLock = createDeferred<void>() |
| 141 | + const holder = db.transaction(async (tx) => { |
| 142 | + await acquireAdvisoryXactLock(tx, USAGE_RECONCILE_LOCK, executionId) |
| 143 | + lockHeld.resolve() |
| 144 | + await releaseLock.promise |
| 145 | + }) |
| 146 | + await lockHeld.promise |
| 147 | + |
| 148 | + const completion = executionLogger.completeWorkflowExecution({ |
| 149 | + executionId, |
| 150 | + endedAt: new Date().toISOString(), |
| 151 | + totalDurationMs: 5, |
| 152 | + costSummary: calculateCostSummary([], { baseExecutionCharge: EXECUTION_FEE }), |
| 153 | + finalOutput: {}, |
| 154 | + traceSpans: [], |
| 155 | + status: 'completed', |
| 156 | + actorUserId: ids.owner, |
| 157 | + billingAttribution, |
| 158 | + }) |
| 159 | + |
| 160 | + /** A completion that settles before blocking surfaces its own outcome instead of a timeout. */ |
| 161 | + const settledWithoutBlocking = completion.then(() => { |
| 162 | + throw new Error('Completion finished without waiting on the ledger lock') |
| 163 | + }) |
| 164 | + |
| 165 | + let statusWhileLedgerBlocked: string | undefined |
| 166 | + try { |
| 167 | + await Promise.race([ |
| 168 | + vi.waitFor(async () => { |
| 169 | + expect(await isLedgerLockAwaited(executionId)).toBe(true) |
| 170 | + }), |
| 171 | + settledWithoutBlocking, |
| 172 | + ]) |
| 173 | + statusWhileLedgerBlocked = (await logRow(executionId))?.status |
| 174 | + } finally { |
| 175 | + releaseLock.resolve() |
| 176 | + await holder |
| 177 | + } |
| 178 | + await completion |
| 179 | + |
| 180 | + expect(statusWhileLedgerBlocked).toBe('running') |
| 181 | + expect(await logRow(executionId)).toMatchObject({ status: 'completed' }) |
| 182 | + expect((await buildCostLedger(executionId))?.total).toBeCloseTo(EXECUTION_FEE, 8) |
| 183 | + expect(Number((await logRow(executionId))?.costTotal)).toBeCloseTo(EXECUTION_FEE, 8) |
| 184 | + }) |
| 185 | +}) |
0 commit comments