chore: import upstream snapshot with attribution

apps/sim/lib/data-drains/dispatcher.ts

@@ -0,0 +1,186 @@
+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<string, boolean>()
+  const isEnterprise = async (orgId: string): Promise<boolean> => {
+    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 }
+}