feat(replication): auto-resubscribe on a failed handoff (self-healing deploys)

d-cs · d-cs · commit 1ae941329e36 · 2026-07-04T15:17:04.000+01:00
Replication clients now optionally re-subscribe (with exponential backoff) after
a lost leader election or a failed START_REPLICATION, instead of logging once and
staying dead. This makes a rolling deploy self-heal: the incoming pod retries
until the draining pod releases the slot, then takes over — no stop-before-start.

Safety:
- #cleanupAttempt() unconditionally ends the pg client (freeing the walsender)
  and releases the leader lock before rescheduling, so retries never leak
  connections/locks (the plain stop() early-returns while _isStopped is set).
- shutdown() (used for intentional stop) sets an intentional-stop latch that is
  re-checked after every await in subscribe() and aborts #acquireLeaderLock's
  spin, so a resubscribe can never race or outlive a shutdown.
- backoff resets only on genuine stream start (replicationStart), so a
  permanently stuck slot backs off to the ceiling and logs loudly rather than
  tight-looping; a subscribeEpoch neutralises stale START_REPLICATION catches.

Runs- and sessions-replication opt in and use shutdown() for all intentional
stops. Adds container tests for the leak, shutdown-race, reset, and self-heal.
diff --git a/apps/webapp/app/services/runsReplicationService.server.ts b/apps/webapp/app/services/runsReplicationService.server.ts
@@ -286,6 +286,7 @@ export class RunsReplicationService {
         table: "TaskRun",
         redisOptions: options.redisOptions,
         autoAcknowledge: false,
+        resubscribeOnFailure: true,
         publicationActions: ["insert", "update", "delete"],
         logger:
           options.logger ?? new Logger("LogicalReplicationClient", options.logLevel ?? "info"),
@@ -428,7 +429,9 @@ export class RunsReplicationService {
 
     if (!hasCurrentTransaction) {
       this.logger.info("No transaction to commit, shutting down immediately");
-      await Promise.all(Array.from(this._sources.values()).map((runtime) => runtime.client.stop()));
+      await Promise.all(
+        Array.from(this._sources.values()).map((runtime) => runtime.client.shutdown())
+      );
       this._isShutDownComplete = true;
       return;
     }
@@ -458,7 +461,7 @@ export class RunsReplicationService {
     for (const runtime of this._sources.values()) {
       this.logger.info("Stopping replication client", { sourceId: runtime.source.id });
 
-      await runtime.client.stop();
+      await runtime.client.shutdown();
 
       if (runtime.acknowledgeInterval) {
         clearInterval(runtime.acknowledgeInterval);
@@ -636,7 +639,7 @@ export class RunsReplicationService {
       // swallow client.stop() rejections so they don't surface as unhandled.
       if (!this._shutdownStopInFlight) {
         this._shutdownStopInFlight = true;
-        Promise.all(Array.from(this._sources.values()).map((r) => r.client.stop()))
+        Promise.all(Array.from(this._sources.values()).map((r) => r.client.shutdown()))
           .catch((error) => {
             this.logger.error("Error stopping replication clients during shutdown", { error });
           })
diff --git a/apps/webapp/app/services/sessionsReplicationService.server.ts b/apps/webapp/app/services/sessionsReplicationService.server.ts
@@ -187,6 +187,7 @@ export class SessionsReplicationService {
       table: "Session",
       redisOptions: options.redisOptions,
       autoAcknowledge: false,
+      resubscribeOnFailure: true,
       publicationActions: ["insert", "update", "delete"],
       logger: options.logger ?? new Logger("LogicalReplicationClient", options.logLevel ?? "info"),
       leaderLockTimeoutMs: options.leaderLockTimeoutMs ?? 30_000,
@@ -265,7 +266,7 @@ export class SessionsReplicationService {
 
     if (!this._currentTransaction) {
       this.logger.info("No transaction to commit, shutting down immediately");
-      await this._replicationClient.stop();
+      await this._replicationClient.shutdown();
       this._isSubscribed = false;
       this._isShutDownComplete = true;
       return;
@@ -294,7 +295,7 @@ export class SessionsReplicationService {
   async stop() {
     this.logger.info("Stopping replication client");
 
-    await this._replicationClient.stop();
+    await this._replicationClient.shutdown();
 
     if (this._acknowledgeInterval) {
       clearInterval(this._acknowledgeInterval);
@@ -430,7 +431,7 @@ export class SessionsReplicationService {
     if (this._isShutDownComplete) return;
 
     if (this._isShuttingDown) {
-      this._replicationClient.stop().finally(() => {
+      this._replicationClient.shutdown().finally(() => {
         this._isSubscribed = false;
         this._isShutDownComplete = true;
       });
diff --git a/internal-packages/replication/src/client.test.ts b/internal-packages/replication/src/client.test.ts
@@ -235,4 +235,234 @@ describe("Replication Client", () => {
       await b.stop();
     }
   );
+
+  postgresAndRedisTest(
+    "resubscribeOnFailure self-heals once the leader releases the slot",
+    async ({ postgresContainer, prisma, redisOptions }) => {
+      await prisma.$executeRawUnsafe(`ALTER TABLE public."TaskRun" REPLICA IDENTITY FULL;`);
+
+      const shared = {
+        slotName: "resub_slot",
+        publicationName: "resub_pub",
+        redisOptions,
+        table: "TaskRun",
+        pgConfig: { connectionString: postgresContainer.getConnectionUri() },
+      };
+
+      // Leader holds the slot.
+      const a = new LogicalReplicationClient({ ...shared, name: "leader-a" });
+      a.events.on("error", () => {});
+      await a.subscribe();
+      await setTimeout(1000);
+
+      // Contender with resubscribe on: loses the election while A holds the slot,
+      // then must self-heal (win) once A releases it — the rolling-deploy handoff.
+      const b = new LogicalReplicationClient({
+        ...shared,
+        name: "contender-b",
+        resubscribeOnFailure: true,
+        resubscribeMinDelayMs: 200,
+        resubscribeMaxDelayMs: 400,
+        leaderLockTimeoutMs: 500,
+        leaderLockAcquireAdditionalTimeMs: 100,
+        leaderLockRetryIntervalMs: 100,
+      });
+      const bElections: boolean[] = [];
+      b.events.on("leaderElection", (won) => bElections.push(won));
+      b.events.on("error", () => {});
+      await b.subscribe();
+      await setTimeout(1500);
+
+      // Still contending, not leader, while A holds the slot.
+      expect(bElections).toContain(false);
+      expect(bElections).not.toContain(true);
+
+      // Release the leader — a scheduled resubscribe should now win.
+      await a.shutdown();
+
+      let becameLeader = false;
+      for (let i = 0; i < 40; i++) {
+        if (bElections.includes(true)) {
+          becameLeader = true;
+          break;
+        }
+        await setTimeout(250);
+      }
+      expect(becameLeader).toBe(true);
+
+      await b.shutdown();
+    }
+  );
+
+  postgresAndRedisTest(
+    "a failing START_REPLICATION retry loop must not leak connections or locks",
+    async ({ postgresContainer, prisma, redisOptions }) => {
+      await prisma.$executeRawUnsafe(`ALTER TABLE public."TaskRun" REPLICA IDENTITY FULL;`);
+
+      const shared = {
+        slotName: "leak_slot",
+        publicationName: "leak_pub",
+        table: "TaskRun",
+        pgConfig: { connectionString: postgresContainer.getConnectionUri() },
+      };
+
+      const a = new LogicalReplicationClient({ ...shared, redisOptions, name: "leak-leader" });
+      a.events.on("error", () => {});
+      await a.subscribe();
+      await setTimeout(1000);
+
+      // B elects on a separate lock namespace so every attempt reaches
+      // START_REPLICATION and dies there ("slot is active") — the stuck-slot shape.
+      const b = new LogicalReplicationClient({
+        ...shared,
+        redisOptions: { ...redisOptions, keyPrefix: `${redisOptions.keyPrefix ?? ""}other:` },
+        name: "leak-contender",
+        resubscribeOnFailure: true,
+        resubscribeMinDelayMs: 200,
+        resubscribeMaxDelayMs: 400,
+        leaderLockTimeoutMs: 1000,
+        leaderLockAcquireAdditionalTimeMs: 300,
+        leaderLockRetryIntervalMs: 100,
+      });
+      const bErrors: Array<unknown> = [];
+      b.events.on("error", (error) => bErrors.push(error));
+      await b.subscribe();
+
+      for (let i = 0; i < 80 && bErrors.length < 3; i++) {
+        await setTimeout(250);
+      }
+      expect(bErrors.length).toBeGreaterThanOrEqual(3);
+
+      // Every failed attempt must end its pg client: at most the one in-flight
+      // attempt's backend may exist, never an accrual across cycles.
+      const backends = await prisma.$queryRaw<{ count: bigint }[]>`
+        SELECT count(*) AS count FROM pg_stat_activity WHERE application_name = 'leak-contender'
+      `;
+      expect(Number(backends[0].count)).toBeLessThanOrEqual(1);
+
+      const active = await prisma.$queryRaw<{ count: bigint }[]>`
+        SELECT count(*) AS count FROM pg_replication_slots WHERE slot_name = 'leak_slot' AND active
+      `;
+      expect(Number(active[0].count)).toBe(1);
+
+      await b.shutdown();
+      await a.shutdown();
+    }
+  );
+
+  postgresAndRedisTest(
+    "shutdown during an in-flight subscribe must not leave a zombie leader",
+    async ({ postgresContainer, prisma, redisOptions }) => {
+      await prisma.$executeRawUnsafe(`ALTER TABLE public."TaskRun" REPLICA IDENTITY FULL;`);
+
+      const shared = {
+        slotName: "zombie_slot",
+        publicationName: "zombie_pub",
+        redisOptions,
+        table: "TaskRun",
+        pgConfig: { connectionString: postgresContainer.getConnectionUri() },
+      };
+
+      const a = new LogicalReplicationClient({ ...shared, name: "zombie-leader" });
+      a.events.on("error", () => {});
+      await a.subscribe();
+      await setTimeout(1000);
+
+      // B's election spins against A's held lock; shut it down mid-subscribe.
+      const b = new LogicalReplicationClient({
+        ...shared,
+        name: "zombie-contender",
+        resubscribeOnFailure: true,
+        leaderLockTimeoutMs: 5000,
+        leaderLockAcquireAdditionalTimeMs: 5000,
+        leaderLockRetryIntervalMs: 100,
+      });
+      const bElections: boolean[] = [];
+      b.events.on("leaderElection", (won) => bElections.push(won));
+      b.events.on("error", () => {});
+
+      const inflight = b.subscribe();
+      await setTimeout(300);
+      await b.shutdown();
+
+      // Release the real leader; a zombie B would now win the lock and the slot.
+      await a.shutdown();
+      await inflight.catch(() => {});
+      await setTimeout(1500);
+
+      const zombieWon = bElections.includes(true);
+      const active = await prisma.$queryRaw<{ count: bigint }[]>`
+        SELECT count(*) AS count FROM pg_replication_slots WHERE slot_name = 'zombie_slot' AND active
+      `;
+      const backends = await prisma.$queryRaw<{ count: bigint }[]>`
+        SELECT count(*) AS count FROM pg_stat_activity WHERE application_name = 'zombie-contender'
+      `;
+      // Reap a zombie (if any) so the test exits cleanly, then assert.
+      await b.shutdown();
+
+      expect(zombieWon).toBe(false);
+      expect(Number(active[0].count)).toBe(0);
+      expect(Number(backends[0].count)).toBe(0);
+    }
+  );
+
+  postgresAndRedisTest(
+    "subscribe after shutdown re-arms resubscribeOnFailure",
+    async ({ postgresContainer, prisma, redisOptions }) => {
+      await prisma.$executeRawUnsafe(`ALTER TABLE public."TaskRun" REPLICA IDENTITY FULL;`);
+
+      const shared = {
+        slotName: "rearm_slot",
+        publicationName: "rearm_pub",
+        redisOptions,
+        table: "TaskRun",
+        pgConfig: { connectionString: postgresContainer.getConnectionUri() },
+      };
+
+      const b = new LogicalReplicationClient({
+        ...shared,
+        name: "rearm-client",
+        resubscribeOnFailure: true,
+        resubscribeMinDelayMs: 200,
+        resubscribeMaxDelayMs: 400,
+        leaderLockTimeoutMs: 500,
+        leaderLockAcquireAdditionalTimeMs: 100,
+        leaderLockRetryIntervalMs: 100,
+      });
+      const bElections: boolean[] = [];
+      b.events.on("leaderElection", (won) => bElections.push(won));
+      b.events.on("error", () => {});
+
+      // Admin stop -> start: shutdown latches the intentional stop...
+      await b.subscribe();
+      await setTimeout(500);
+      await b.shutdown();
+
+      const a = new LogicalReplicationClient({ ...shared, name: "rearm-leader" });
+      a.events.on("error", () => {});
+      await a.subscribe();
+      await setTimeout(1000);
+
+      // ...then an explicit re-subscribe loses the election and must self-heal
+      // once the leader goes away (self-heal re-armed by the subscribe).
+      bElections.length = 0;
+      await b.subscribe();
+      expect(bElections).toContain(false);
+      expect(bElections).not.toContain(true);
+
+      await a.shutdown();
+
+      let becameLeader = false;
+      for (let i = 0; i < 40; i++) {
+        if (bElections.includes(true)) {
+          becameLeader = true;
+          break;
+        }
+        await setTimeout(250);
+      }
+      expect(becameLeader).toBe(true);
+
+      await b.shutdown();
+    }
+  );
 });
diff --git a/internal-packages/replication/src/client.ts b/internal-packages/replication/src/client.ts
