import { db } from '@sim/db' import { dataDrainRuns, dataDrains } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { toError } from '@sim/utils/errors' import { and, eq, isNull, lt, or } from 'drizzle-orm' import { isOrganizationOnEnterprisePlan } from '@/lib/billing/core/subscription' import { getJobQueue } from '@/lib/core/async-jobs' import { isBillingEnabled } from '@/lib/core/config/env-flags' const logger = createLogger('DataDrainsDispatcher') const HOUR_MS = 60 * 60 * 1000 const DAY_MS = 24 * HOUR_MS /** * Cron fires hourly. Without a buffer, a drain that finishes a few minutes * after the tick (lastRunAt = 10:05) won't satisfy `lastRunAt < now - cadence` * at the next tick (10:05 < 10:00 is false), so an "hourly" drain effectively * runs every two hours. Subtracting a small buffer from the cadence absorbs * normal run duration plus cron jitter without allowing back-to-back runs * within the same tick. */ const CADENCE_BUFFER_MS = 5 * 60 * 1000 /** * Maximum wall-clock duration any single drain run is allowed before its * `data_drain_runs` row is considered orphaned. Runs that exceed this are * almost certainly the result of a Trigger.dev worker crash mid-run — there * is no live process still updating them. */ const ORPHAN_THRESHOLD_MS = 60 * 60 * 1000 /** * Marks `running` rows older than the orphan threshold as `failed`. Without * this, a worker crash leaves run history permanently misleading and (worse) * the drain row's `lastRunAt` reflects a successful claim that never finished * — but the drain `cursor` never advanced, so re-running is safe. */ export async function reapOrphanedRuns(now: Date = new Date()): Promise<{ reaped: number }> { const cutoff = new Date(now.getTime() - ORPHAN_THRESHOLD_MS) const reaped = await db .update(dataDrainRuns) .set({ status: 'failed', finishedAt: now, error: `Orphaned run reaped after exceeding ${ORPHAN_THRESHOLD_MS / 60_000}m without completion`, }) .where(and(eq(dataDrainRuns.status, 'running'), lt(dataDrainRuns.startedAt, cutoff))) .returning({ id: dataDrainRuns.id }) if (reaped.length > 0) { logger.warn('Reaped orphaned data drain runs', { count: reaped.length }) } return { reaped: reaped.length } } /** * Selects every enabled drain whose schedule is due (or has never run) and * fans out one `run-data-drain` job per drain. Each drain is atomically * claimed via a conditional UPDATE before being enqueued — two concurrent * dispatcher invocations cannot both win the same row, and a manual run that * lands between the SELECT and the UPDATE will lose the race cleanly. Drains * belonging to orgs that have lapsed off the enterprise plan are skipped. */ export async function dispatchDueDrains(now: Date = new Date()): Promise<{ candidates: number dispatched: number skipped: number reaped: number }> { const { reaped } = await reapOrphanedRuns(now) const hourlyCutoff = new Date(now.getTime() - HOUR_MS + CADENCE_BUFFER_MS) const dailyCutoff = new Date(now.getTime() - DAY_MS + CADENCE_BUFFER_MS) const duePredicate = and( eq(dataDrains.enabled, true), or( isNull(dataDrains.lastRunAt), and(eq(dataDrains.scheduleCadence, 'hourly'), lt(dataDrains.lastRunAt, hourlyCutoff)), and(eq(dataDrains.scheduleCadence, 'daily'), lt(dataDrains.lastRunAt, dailyCutoff)) ) ) const candidates = await db .select({ id: dataDrains.id, organizationId: dataDrains.organizationId, lastRunAt: dataDrains.lastRunAt, }) .from(dataDrains) .where(duePredicate) if (candidates.length === 0) { return { candidates: 0, dispatched: 0, skipped: 0, reaped } } // Self-hosted deployments have no subscription infra; `DATA_DRAINS_ENABLED` // is the global on/off there. Cache per-org so a multi-drain org pays one // billing lookup. const enterpriseCache = new Map() const isEnterprise = async (orgId: string): Promise => { if (!isBillingEnabled) return true const cached = enterpriseCache.get(orgId) if (cached !== undefined) return cached const result = await isOrganizationOnEnterprisePlan(orgId) enterpriseCache.set(orgId, result) return result } const queue = await getJobQueue() let dispatched = 0 let skipped = 0 for (const candidate of candidates) { let enterprise: boolean try { enterprise = await isEnterprise(candidate.organizationId) } catch (error) { // A billing-API failure for one org must not abort the whole batch — // skip this drain and let the next cron tick retry it. logger.warn('Enterprise check failed; skipping drain', { drainId: candidate.id, organizationId: candidate.organizationId, error, }) skipped++ continue } if (!enterprise) { skipped++ continue } // Conditional claim — re-asserts the due predicate to lose to any other // dispatcher or manual-run path that's already moved this drain forward. const claimed = await db .update(dataDrains) .set({ lastRunAt: now, updatedAt: now }) .where(and(eq(dataDrains.id, candidate.id), duePredicate)) .returning({ id: dataDrains.id }) if (claimed.length === 0) continue try { // concurrencyKey serializes runs of the same drain on the job queue, so // a manual run-now racing a cron claim can never execute in parallel. await queue.enqueue( 'run-data-drain', { drainId: candidate.id, trigger: 'cron' }, { concurrencyKey: `data-drain:${candidate.id}` } ) dispatched++ } catch (error) { // Roll back the claim so a transient queue outage doesn't delay this // drain by a full cadence. Scoped to our own claim timestamp so it // can't trample a concurrent advance. The rollback itself is guarded // so a DB error here doesn't abort the rest of the batch. try { await db .update(dataDrains) .set({ lastRunAt: candidate.lastRunAt, updatedAt: now }) .where(and(eq(dataDrains.id, candidate.id), eq(dataDrains.lastRunAt, now))) } catch (rollbackError) { logger.error('Failed to roll back data-drain claim after enqueue failure', { drainId: candidate.id, enqueueError: toError(error).message, rollbackError: toError(rollbackError).message, }) continue } logger.error('Failed to enqueue data-drain job; rolled back claim', { drainId: candidate.id, error, }) } } logger.info('Data drain dispatch complete', { candidates: candidates.length, dispatched, skipped, reaped, }) return { candidates: candidates.length, dispatched, skipped, reaped } }