refactor: simplify rank tracking backend (#125)

drizzle/0011_colorful_dark_beast.sql

@@ -0,0 +1,12 @@
+-- Reconcile any stranded pending/running runs before creating the partial
+-- unique index that replaces the rank_check_locks table. With the old
+-- lock-table model, at most one active run per config could exist, so this
+-- mostly protects against orphaned rows that outlived their locks.
+UPDATE `rank_check_runs`
+SET
+  `status` = 'failed',
+  `error_message` = COALESCE(`error_message`, 'Reconciled during lock-table migration'),
+  `completed_at` = COALESCE(`completed_at`, CURRENT_TIMESTAMP)
+WHERE `status` IN ('pending', 'running');--> statement-breakpoint
+DROP TABLE `rank_check_locks`;--> statement-breakpoint
+CREATE UNIQUE INDEX `rank_check_runs_one_active_per_config_idx` ON `rank_check_runs` (`config_id`) WHERE "rank_check_runs"."status" IN ('pending', 'running');

drizzle/meta/_journal.json

@@ -78,6 +78,13 @@
       "when": 1777422740449,
       "tag": "0010_high_liz_osborn",
       "breakpoints": true
+    },
+    {
+      "idx": 11,
+      "version": "6",
+      "when": 1778031161783,
+      "tag": "0011_colorful_dark_beast",
+      "breakpoints": true
     }
   ]
 }

src/db/app.schema.ts

@@ -168,7 +168,11 @@ export const rankTrackingKeywords = sqliteTable(
   ],
 );
 
-// One row per check execution (manual or scheduled)
+// One row per check execution (manual or scheduled).
+// A partial unique index on `config_id WHERE status IN ('pending','running')`
+// enforces at most one in-flight run per config at the DB level, which is how
+// duplicate-trigger protection is implemented — INSERT of a second pending run
+// for the same config fails with a unique-constraint violation.
 export const rankCheckRuns = sqliteTable(
   "rank_check_runs",
   {
@@ -198,24 +202,12 @@ export const rankCheckRuns = sqliteTable(
   (table) => [
     index("rank_check_runs_config_idx").on(table.configId, table.startedAt),
     index("rank_check_runs_project_idx").on(table.projectId, table.startedAt),
+    uniqueIndex("rank_check_runs_one_active_per_config_idx")
+      .on(table.configId)
+      .where(sql`${table.status} IN ('pending', 'running')`),
   ],
 );
 
-// One active lock per rank tracking config to prevent overlapping runs
-export const rankCheckLocks = sqliteTable(
-  "rank_check_locks",
-  {
-    configId: text("config_id")
-      .primaryKey()
-      .references(() => rankTrackingConfigs.id, { onDelete: "cascade" }),
-    runId: text("run_id").notNull(),
-    acquiredAt: text("acquired_at")
-      .notNull()
-      .default(sql`(current_timestamp)`),
-  },
-  (table) => [index("rank_check_locks_run_idx").on(table.runId)],
-);
-
 // One row per keyword per device per check run
 export const rankSnapshots = sqliteTable(
   "rank_snapshots",

src/server/features/rank-tracking/repositories/RankTrackingRepository.ts

@@ -4,7 +4,6 @@ import { db } from "@/db";
 import {
   rankTrackingConfigs,
   rankCheckRuns,
-  rankCheckLocks,
   rankSnapshots,
   rankTrackingKeywords,
   projects,
@@ -137,17 +136,27 @@ async function getDueConfigsWithOrganization(nowIso: string) {
 // Run CRUD
 // ---------------------------------------------------------------------------
 
-async function createRun(data: {
+/**
+ * Try to insert a new pending run. Returns true if inserted, false if blocked
+ * by the partial unique index on (config_id) WHERE status IN ('pending',
+ * 'running') — i.e. another active run exists for this config.
+ *
+ * This is how duplicate-trigger protection is enforced: the DB rejects the
+ * second insert rather than a separate lock table.
+ */
+async function tryCreateRun(data: {
   id: string;
   configId: string;
   projectId: string;
   keywordsTotal: number;
   isSubsetRun?: boolean;
-}) {
+}): Promise<boolean> {
-  await db.insert(rankCheckRuns).values({
+  const inserted = await db
-    ...data,
+    .insert(rankCheckRuns)
-    status: "pending",
+    .values({ ...data, status: "pending" })
-  });
+    .onConflictDoNothing()
+    .returning({ id: rankCheckRuns.id });
+  return inserted.length > 0;
 }
 
 async function updateRun(
@@ -176,38 +185,24 @@ async function getLatestRunForConfig(configId: string) {
   return rows[0] ?? null;
 }
 
-async function tryCreateRunLock(configId: string, runId: string) {
+/**
-  const inserted = await db
+ * Returns the currently active (pending or running) run for a config, if any.
-    .insert(rankCheckLocks)
+ * At most one such row exists, enforced by the partial unique index.
-    .values({ configId, runId })
+ */
-    .onConflictDoNothing({ target: rankCheckLocks.configId })
+async function getActiveRunForConfig(configId: string) {
-    .returning({ runId: rankCheckLocks.runId });
-
-  return inserted.length > 0;
-}
-
-async function getRunLock(configId: string) {
   const rows = await db
     .select()
-    .from(rankCheckLocks)
+    .from(rankCheckRuns)
-    .where(eq(rankCheckLocks.configId, configId))
+    .where(
+      and(
+        eq(rankCheckRuns.configId, configId),
+        inArray(rankCheckRuns.status, ["pending", "running"]),
+      ),
+    )
     .limit(1);
   return rows[0] ?? null;
 }
 
-async function deleteRunLock(configId: string, runId?: string) {
-  await db
-    .delete(rankCheckLocks)
-    .where(
-      runId
-        ? and(
-            eq(rankCheckLocks.configId, configId),
-            eq(rankCheckLocks.runId, runId),
-          )
-        : eq(rankCheckLocks.configId, configId),
-    );
-}
-
 // ---------------------------------------------------------------------------
 // Snapshots
 // ---------------------------------------------------------------------------
@@ -369,13 +364,11 @@ export const RankTrackingRepository = {
   createConfig,
   updateConfig,
   getDueConfigsWithOrganization,
-  createRun,
+  tryCreateRun,
   updateRun,
   getRunById,
   getLatestRunForConfig,
-  tryCreateRunLock,
+  getActiveRunForConfig,
-  getRunLock,
-  deleteRunLock,
   insertSnapshots,
   getSnapshotsForRun,
   getKeywordsForConfig,

src/server/features/rank-tracking/services/RankTrackingService.ts

@@ -218,8 +218,8 @@ async function getLatestRun(configId: string, projectId: string) {
   if (!run) return null;
 
   // If the DB says the run is still active, check the workflow instance.
-  // We only report staleness here — the cron handler cleans up stale locks
+  // We only report staleness here — the next call to beginRankCheckRun will
-  // when it next tries to acquire one (via cleanupStaleLock). Mutating from
+  // mark a stale blocker as failed before retrying its insert. Mutating from
   // this read path caused a race where the original workflow kept running
   // while a replacement was started.
   const reconciliation = await reconcileActiveRankCheckRun(run);

src/server/features/rank-tracking/services/rankCheckRunGuards.ts

@@ -7,16 +7,16 @@ import type {
 } from "@/types/schemas/rank-tracking";
 
 type RunRow = Awaited<ReturnType<typeof RankTrackingRepository.getRunById>>;
-type RunLockRow = Awaited<ReturnType<typeof RankTrackingRepository.getRunLock>>;
 
-// Coordination invariants for rank checks:
+// Coordination model:
-// - `workflow id === run id`, so the workflow instance is the authoritative
+// - workflow id === run id (workflow instance is the authoritative runtime).
-//   runtime identity for a stored run.
+// - A partial unique index on rank_check_runs(config_id) WHERE status IN
-// - `rank_check_locks` enforces at most one active runner per config.
+//   ('pending','running') enforces at most one active run per config at the
-// - Only the lock owner is allowed to spend credits, write snapshots, or
+//   DB level. A failed INSERT *is* the "already running" signal — no
-//   finalize the run.
+//   separate lock table is needed.
-// - Missing/unknown workflow state is tolerated briefly during startup before
+// - Flipping status to 'completed'/'failed' is what frees the slot.
-//   we treat the run as stale and repair it.
+// - Missing/unknown workflow state is tolerated briefly during startup
+//   before we treat a run as stale and mark it failed.
 
 type RankCheckWorkflowStatus = {
   status:
@@ -77,85 +77,11 @@ function getStaleReason(
     return workflowStatus.error?.message ?? `Workflow ${workflowStatus.status}`;
   }
   if (workflowStatus.status === "complete") {
-    return "Workflow completed without releasing the run lock";
+    return "Workflow completed without finalizing the run";
   }
   return `Workflow is no longer active (${workflowStatus.status})`;
 }
 
-export async function beginRankCheckRun(input: {
-  workflow: Env["RANK_CHECK_WORKFLOW"];
-  config: RankCheckConfigForStart;
-  projectId: string;
-  billingCustomer: BillingCustomerContext;
-  keywordsTotal: number;
-  keywordIds?: string[];
-  trigger: "manual" | "scheduled";
-  workflowStartErrorMessage: string;
-}): Promise<RankCheckTriggerResult> {
-  const runId = crypto.randomUUID();
-  const lockResult = await acquireRankCheckRunLock(input.config.id, runId);
-  if (!lockResult.acquired) {
-    return {
-      ok: false,
-      reason: "already_running",
-      blockingRunId: lockResult.blockingRunId,
-    };
-  }
-
-  try {
-    await RankTrackingRepository.createRun({
-      id: runId,
-      configId: input.config.id,
-      projectId: input.projectId,
-      keywordsTotal: input.keywordsTotal,
-      isSubsetRun: (input.keywordIds?.length ?? 0) > 0,
-    });
-
-    await input.workflow.create({
-      id: runId,
-      params: {
-        runId,
-        configId: input.config.id,
-        billingCustomer: {
-          userId: input.billingCustomer.userId,
-          userEmail: input.billingCustomer.userEmail,
-          organizationId: input.billingCustomer.organizationId,
-          projectId: input.billingCustomer.projectId,
-        },
-        projectId: input.projectId,
-        domain: input.config.domain,
-        locationCode: input.config.locationCode,
-        languageCode: input.config.languageCode,
-        devices: input.config.devices,
-        serpDepth: input.config.serpDepth,
-        trigger: input.trigger,
-        keywordIds: input.keywordIds,
-      },
-    });
-  } catch (error) {
-    try {
-      await failRunAndReleaseRankCheckLock(
-        input.config.id,
-        runId,
-        input.workflowStartErrorMessage,
-      );
-    } catch {
-      await releaseRankCheckRunLock(input.config.id, runId);
-    }
-
-    try {
-      const instance = await input.workflow.get(runId);
-      await instance.terminate();
-    } catch {
-      // Workflow may not have been created
-    }
-
-    throw error;
-  }
-
-  return { ok: true, runId };
-}
-
 async function getStaleRankCheckRunReason(input: {
   run: RunRow;
   runId: string;
@@ -181,8 +107,23 @@ async function getStaleRankCheckRunReason(input: {
   return getStaleReason(workflowStatus, input.run);
 }
 
-async function failRunIfNeeded(runId: string, reason: string, run: RunRow) {
+/**
-  if (!run || run.status === "completed" || run.status === "failed") return;
+ * Mark a run as failed if it's still in an active state. Idempotent — safe to
+ * call on runs that are already completed/failed.
+ */
+export async function failRunIfActive(
+  runId: string,
+  reason: string,
+  run?: RunRow,
+) {
+  const current = run ?? (await RankTrackingRepository.getRunById(runId));
+  if (
+    !current ||
+    current.status === "completed" ||
+    current.status === "failed"
+  ) {
+    return;
+  }
   await RankTrackingRepository.updateRun(runId, {
     status: "failed",
     errorMessage: reason,
@@ -190,23 +131,95 @@ async function failRunIfNeeded(runId: string, reason: string, run: RunRow) {
   });
 }
 
-async function cleanupStaleLock(lock: NonNullable<RunLockRow>) {
+export async function beginRankCheckRun(input: {
-  const run = await RankTrackingRepository.getRunById(lock.runId);
+  workflow: Env["RANK_CHECK_WORKFLOW"];
-  if (run?.status === "completed" || run?.status === "failed") {
+  config: RankCheckConfigForStart;
-    await RankTrackingRepository.deleteRunLock(lock.configId, lock.runId);
+  projectId: string;
-    return true;
+  billingCustomer: BillingCustomerContext;
+  keywordsTotal: number;
+  keywordIds?: string[];
+  trigger: "manual" | "scheduled";
+  workflowStartErrorMessage: string;
+}): Promise<RankCheckTriggerResult> {
+  // At most two attempts: once normally, once after clearing a stale blocker.
+  for (let attempt = 0; attempt < 2; attempt++) {
+    const runId = crypto.randomUUID();
+    const inserted = await RankTrackingRepository.tryCreateRun({
+      id: runId,
+      configId: input.config.id,
+      projectId: input.projectId,
+      keywordsTotal: input.keywordsTotal,
+      isSubsetRun: (input.keywordIds?.length ?? 0) > 0,
+    });
+
+    if (inserted) {
+      try {
+        await input.workflow.create({
+          id: runId,
+          params: {
+            runId,
+            configId: input.config.id,
+            billingCustomer: input.billingCustomer,
+            projectId: input.projectId,
+            domain: input.config.domain,
+            locationCode: input.config.locationCode,
+            languageCode: input.config.languageCode,
+            devices: input.config.devices,
+            serpDepth: input.config.serpDepth,
+            trigger: input.trigger,
+            keywordIds: input.keywordIds,
+          },
+        });
+      } catch (error) {
+        // Workflow couldn't start — flip the run to failed so the
+        // partial-index slot is released. Best-effort cleanup of any
+        // zombie instance.
+        await failRunIfActive(runId, input.workflowStartErrorMessage);
+        try {
+          const instance = await input.workflow.get(runId);
+          await instance.terminate();
+        } catch {
+          // Workflow may not have been created.
+        }
+        throw error;
+      }
+      return { ok: true, runId };
+    }
+
+    // INSERT was blocked by the partial unique index — another active run
+    // exists. Inspect it to decide whether to retry or return already_running.
+    const blocker = await RankTrackingRepository.getActiveRunForConfig(
+      input.config.id,
+    );
+    if (!blocker) {
+      // Raced: blocker's status flipped between insert and select. Loop.
+      continue;
+    }
+
+    if (attempt === 0) {
+      const staleReason = await getStaleRankCheckRunReason({
+        run: blocker,
+        runId: blocker.id,
+        ageMs: Date.now() - new Date(blocker.startedAt).getTime(),
+      });
+      if (staleReason) {
+        await failRunIfActive(blocker.id, staleReason, blocker);
+        continue; // slot is free now — retry insert
+      }
+    }
+
+    return { ok: false, reason: "already_running", blockingRunId: blocker.id };
   }
 
-  const staleReason = await getStaleRankCheckRunReason({
+  // Exhausted attempts (rapid churn). Report whatever's blocking now.
-    runId: lock.runId,
+  const final = await RankTrackingRepository.getActiveRunForConfig(
-    run,
+    input.config.id,
-    ageMs: Date.now() - new Date(lock.acquiredAt).getTime(),
+  );
-  });
+  return {
-  if (!staleReason) return false;
+    ok: false,
-
+    reason: "already_running",
-  await failRunIfNeeded(lock.runId, staleReason, run);
+    blockingRunId: final?.id ?? null,
-  await RankTrackingRepository.deleteRunLock(lock.configId, lock.runId);
+  };
-  return true;
 }
 
 export async function reconcileActiveRankCheckRun(run: NonNullable<RunRow>) {
@@ -226,53 +239,3 @@ export async function reconcileActiveRankCheckRun(run: NonNullable<RunRow>) {
     completedAt: new Date().toISOString(),
   };
 }
-
-async function acquireRankCheckRunLock(configId: string, runId: string) {
-  // Optimistic: try to grab the lock immediately
-  if (await RankTrackingRepository.tryCreateRunLock(configId, runId)) {
-    return { acquired: true as const };
-  }
-
-  // Lock exists — check if it's stale and can be cleaned up
-  const existingLock = await RankTrackingRepository.getRunLock(configId);
-  if (!existingLock) {
-    // Lock was released between our insert and select — retry once
-    if (await RankTrackingRepository.tryCreateRunLock(configId, runId)) {
-      return { acquired: true as const };
-    }
-    const blocker = await RankTrackingRepository.getRunLock(configId);
-    return { acquired: false as const, blockingRunId: blocker?.runId ?? null };
-  }
-
-  const cleaned = await cleanupStaleLock(existingLock);
-  if (!cleaned) {
-    return { acquired: false as const, blockingRunId: existingLock.runId };
-  }
-
-  // Stale lock cleaned — retry once
-  if (await RankTrackingRepository.tryCreateRunLock(configId, runId)) {
-    return { acquired: true as const };
-  }
-
-  const blocker = await RankTrackingRepository.getRunLock(configId);
-  return { acquired: false as const, blockingRunId: blocker?.runId ?? null };
-}
-
-export async function runOwnsRankCheckLock(configId: string, runId: string) {
-  const lock = await RankTrackingRepository.getRunLock(configId);
-  return lock?.runId === runId;
-}
-
-export async function releaseRankCheckRunLock(configId: string, runId: string) {
-  await RankTrackingRepository.deleteRunLock(configId, runId);
-}
-
-export async function failRunAndReleaseRankCheckLock(
-  configId: string,
-  runId: string,
-  errorMessage: string,
-) {
-  const run = await RankTrackingRepository.getRunById(runId);
-  await failRunIfNeeded(runId, errorMessage, run);
-  await RankTrackingRepository.deleteRunLock(configId, runId);
-}

src/server/workflows/RankCheckWorkflow.ts

@@ -6,11 +6,7 @@ import {
 import { NonRetryableError } from "cloudflare:workflows";
 import type { BillingCustomerContext } from "@/server/billing/subscription";
 import { RankTrackingRepository } from "@/server/features/rank-tracking/repositories/RankTrackingRepository";
-import {
+import { failRunIfActive } from "@/server/features/rank-tracking/services/rankCheckRunGuards";
-  failRunAndReleaseRankCheckLock,
-  releaseRankCheckRunLock,
-  runOwnsRankCheckLock,
-} from "@/server/features/rank-tracking/services/rankCheckRunGuards";
 import { runLiveCheck } from "@/server/workflows/rankCheckPaths";
 import { createDataforseoClient } from "@/server/lib/dataforseoClient";
 import { captureServerEvent } from "@/server/lib/posthog";
@@ -50,10 +46,12 @@ async function prepareRankCheckKeywords(input: {
   serpDepth: number;
   keywordIds?: string[];
 }) {
-  const ownsLock = await runOwnsRankCheckLock(input.configId, input.runId);
+  // If stale-cleanup marked our run failed before we got here, bail out
-  if (!ownsLock) {
+  // rather than resurrecting a superseded run.
+  const run = await RankTrackingRepository.getRunById(input.runId);
+  if (!run || run.status === "failed" || run.status === "completed") {
     throw new NonRetryableError(
-      `Rank check lock is not held by run ${input.runId}`,
+      `Run ${input.runId} is no longer active (status=${run?.status ?? "missing"})`,
     );
   }
 
@@ -122,13 +120,13 @@ async function finalizeRankCheckRun(input: {
   trigger: RankCheckParams["trigger"];
   batchError: string | null;
 }) {
-  // Re-check lock ownership before finalizing. If the lock was stolen
+  // If stale-cleanup already marked our run failed, don't overwrite that
-  // (stale cleanup raced with a slow workflow), bail out to avoid
+  // decision with a completed status — a replacement run may already be
-  // overwriting the replacement run's state.
+  // underway.
-  const ownsLock = await runOwnsRankCheckLock(input.configId, input.runId);
+  const run = await RankTrackingRepository.getRunById(input.runId);
-  if (!ownsLock) {
+  if (!run || run.status === "failed" || run.status === "completed") {
     console.warn(
-      `[rank-check] ${input.runId} lost lock ownership, skipping finalization`,
+      `[rank-check] ${input.runId} no longer active (status=${run?.status ?? "missing"}), skipping finalization`,
     );
     return;
   }
@@ -143,9 +141,7 @@ async function finalizeRankCheckRun(input: {
   const keywordsChecked = new Set(snapshots.map((s) => s.trackingKeywordId))
     .size;
 
-  // Derive incompleteCount from the run's keywordsTotal (set in prepare step)
+  const keywordsTotal = run.keywordsTotal || keywordsChecked;
-  const run = await RankTrackingRepository.getRunById(input.runId);
-  const keywordsTotal = run?.keywordsTotal ?? keywordsChecked;
   const incompleteCount = keywordsTotal - keywordsChecked;
 
   let errorMessage: string | undefined;
@@ -155,6 +151,8 @@ async function finalizeRankCheckRun(input: {
     errorMessage = `${incompleteCount} keyword(s) could not be checked`;
   }
 
+  // Flipping status away from 'pending'/'running' is what releases the
+  // partial-index slot for the next run.
   await RankTrackingRepository.updateRun(input.runId, {
     status: "completed",
     keywordsChecked,
@@ -170,8 +168,6 @@ async function finalizeRankCheckRun(input: {
     lastSkipReason: null,
   });
 
-  await releaseRankCheckRunLock(input.configId, input.runId);
-
   await captureServerEvent({
     distinctId: input.billingCustomer.userId,
     event: "rank_tracking:check_complete",
@@ -194,11 +190,7 @@ async function markRankCheckRunFailed(input: {
 }) {
   const errorMessage =
     input.error instanceof Error ? input.error.message : "Unknown error";
-  await failRunAndReleaseRankCheckLock(
+  await failRunIfActive(input.runId, errorMessage);
-    input.configId,
-    input.runId,
-    errorMessage,
-  );
 
   // Flag the config so the UI can show why the scheduled check was skipped
   const isInsufficientCredits =
@@ -256,11 +248,7 @@ export class RankCheckWorkflow extends WorkflowEntrypoint<
       },
     );
     if (!configCheck.isActive) {
-      await failRunAndReleaseRankCheckLock(
+      await failRunIfActive(runId, "Config has been archived");
-        configId,
-        runId,
-        "Config has been archived",
-      );
       return;
     }