more replication stuff

Author: ericallam

internal-packages/replication/src/client.ts

@@ -143,25 +143,7 @@ export class LogicalReplicationClient {
       this.ackIntervalTimer = null;
     }
     // Release leader lock if held
-    if (this.leaderLock) {
-      const [releaseError] = await tryCatch(this.leaderLock.release());
-
-      if (releaseError) {
-        this.logger.error("Failed to release leader lock", {
-          name: this.options.name,
-          slotName: this.options.slotName,
-          publicationName: this.options.publicationName,
-          error: releaseError,
-        });
-      } else {
-        this.logger.info("Released leader lock", {
-          name: this.options.name,
-          slotName: this.options.slotName,
-        });
-      }
-
-      this.leaderLock = null;
-    }
+    await this.#releaseLeaderLock();
 
     this.connection?.removeAllListeners();
     this.connection = null;
@@ -198,6 +180,24 @@ export class LogicalReplicationClient {
     return this;
   }
 
+  public async teardown(): Promise<boolean> {
+    await this.stop();
+
+    // Acquire the leaderLock
+    const leaderLockAcquired = await this.#acquireLeaderLock();
+
+    if (!leaderLockAcquired) {
+      return false;
+    }
+
+    // Drop the slot
+    const slotDropped = await this.#dropSlot();
+
+    await this.#releaseLeaderLock();
+
+    return slotDropped;
+  }
+
   public async subscribe(startLsn?: string): Promise<this> {
     await this.stop();
 
@@ -212,28 +212,10 @@ export class LogicalReplicationClient {
     });
 
     // 1. Leader election
-    try {
-      this.leaderLock = await this.redlock.acquire(
-        [`logical-replication-client:${this.options.name}`],
-        this.leaderLockTimeoutMs,
-        {
-          retryCount: 60,
-          retryDelay: 1000,
-          retryJitter: 100,
-        }
-      );
-    } catch (err) {
-      this.logger.error("Leader election failed", {
-        name: this.options.name,
-        table: this.options.table,
-        slotName: this.options.slotName,
-        publicationName: this.options.publicationName,
-        startLsn,
-        error: err,
-      });
+    const leaderLockAcquired = await this.#acquireLeaderLock();
 
+    if (!leaderLockAcquired) {
       this.events.emit("leaderElection", false);
-
       return this.stop();
     }
 
@@ -481,6 +463,31 @@ export class LogicalReplicationClient {
     return res.rows[0].exists;
   }
 
+  async #dropSlot(): Promise<boolean> {
+    if (!this.client) {
+      this.events.emit("error", new LogicalReplicationClientError("Cannot drop slot"));
+      return false;
+    }
+
+    const [dropError] = await tryCatch(
+      this.client.query(`SELECT pg_drop_replication_slot('${this.options.slotName}');`)
+    );
+
+    if (dropError) {
+      this.logger.error("Failed to drop slot", {
+        name: this.options.name,
+        table: this.options.table,
+        slotName: this.options.slotName,
+        publicationName: this.options.publicationName,
+        error: dropError,
+      });
+
+      this.events.emit("error", dropError);
+    }
+
+    return true;
+  }
+
   async #acknowledge(lsn: string): Promise<void> {
     if (!this.autoAcknowledge) return;
     this.events.emit("acknowledge", { lsn });
@@ -520,6 +527,45 @@ export class LogicalReplicationClient {
     return true;
   }
 
+  async #acquireLeaderLock(): Promise<boolean> {
+    try {
+      this.leaderLock = await this.redlock.acquire(
+        [`logical-replication-client:${this.options.name}`],
+        this.leaderLockTimeoutMs,
+        {
+          retryCount: 60,
+          retryDelay: 1000,
+          retryJitter: 100,
+        }
+      );
+    } catch (err) {
+      this.logger.error("Leader election failed", {
+        name: this.options.name,
+        table: this.options.table,
+        slotName: this.options.slotName,
+        publicationName: this.options.publicationName,
+        error: err,
+      });
+
+      return false;
+    }
+
+    return true;
+  }
+
+  async #releaseLeaderLock() {
+    if (!this.leaderLock) return;
+    const [releaseError] = await tryCatch(this.leaderLock.release());
+    this.leaderLock = null;
+
+    if (releaseError) {
+      this.logger.error("Failed to release leader lock", {
+        name: this.options.name,
+        error: releaseError,
+      });
+    }
+  }
+
   async #startLeaderLockHeartbeat() {
     if (this.leaderLockHeartbeatTimer) {
       clearInterval(this.leaderLockHeartbeatTimer);

internal-packages/replication/src/pgoutput.ts

@@ -134,11 +134,21 @@ class BinaryReader {
     const lower = this.readInt32();
     return upper.toString(16).toUpperCase() + "/" + lower.toString(16).toUpperCase();
   }
+  readUint32(): number {
+    // >>> 0 ensures unsigned
+    return this.readInt32() >>> 0;
+  }
+
+  readUint64(): bigint {
+    // Combine two unsigned 32-bit ints into a 64-bit bigint
+    return (BigInt(this.readUint32()) << 32n) | BigInt(this.readUint32());
+  }
+
   readTime(): bigint {
-    // microseconds since 2000-01-01
-    const high = this.readInt32();
-    const low = this.readInt32();
-    return BigInt(high) * 4294967296n + BigInt(low);
+    // (POSTGRES_EPOCH_JDATE - UNIX_EPOCH_JDATE) * USECS_PER_DAY == 946684800000000
+    const microsSinceUnixEpoch = this.readUint64() + 946684800000000n;
+    // Convert to milliseconds for JS Date compatibility
+    return microsSinceUnixEpoch;
   }
 }
 

internal-packages/replication/src/stream.test.ts

@@ -81,6 +81,7 @@ describe("LogicalReplicationStream", () => {
 
         resolve(undefined);
       }).then(() => {});
+
       // Now we want to read from the stream
       for await (const transaction of subscription.stream) {
         received.push(transaction);
@@ -97,4 +98,106 @@ describe("LogicalReplicationStream", () => {
       await subscription.client.stop();
     }
   );
+
+  postgresAndRedisTest(
+    "should respect highWaterMark and not pull more data than allowed",
+    async ({ postgresContainer, prisma, redisOptions }) => {
+      await prisma.$executeRawUnsafe(`ALTER TABLE public."TaskRun" REPLICA IDENTITY FULL;`);
+
+      type TaskRunData = {
+        friendlyId: string;
+        taskIdentifier: string;
+        payload: string;
+        traceId: string;
+        spanId: string;
+        queue: string;
+        runtimeEnvironmentId: string;
+        projectId: string;
+      };
+
+      const subscription = createSubscription<TaskRunData>({
+        name: "test_stream",
+        publicationName: "test_publication_stream",
+        slotName: "test_slot_stream",
+        pgConfig: {
+          connectionString: postgresContainer.getConnectionUri(),
+        },
+        table: "TaskRun",
+        redisOptions,
+        filterTags: ["insert"],
+        abortSignal: AbortSignal.timeout(10000),
+      });
+
+      const organization = await prisma.organization.create({
+        data: {
+          title: "test",
+          slug: "test",
+        },
+      });
+
+      const project = await prisma.project.create({
+        data: {
+          name: "test",
+          slug: "test",
+          organizationId: organization.id,
+          externalRef: "test",
+        },
+      });
+
+      const runtimeEnvironment = await prisma.runtimeEnvironment.create({
+        data: {
+          slug: "test",
+          type: "DEVELOPMENT",
+          projectId: project.id,
+          organizationId: organization.id,
+          apiKey: "test",
+          pkApiKey: "test",
+          shortcode: "test",
+        },
+      });
+
+      // Insert a row into the table
+      new Promise(async (resolve) => {
+        await setTimeout(2000);
+
+        for (let i = 0; i < 5; i++) {
+          await prisma.taskRun.create({
+            data: {
+              friendlyId: `run_${i}`,
+              taskIdentifier: "my-task",
+              payload: JSON.stringify({ foo: "bar" }),
+              traceId: `${i}`,
+              spanId: `${i}`,
+              queue: "test",
+              runtimeEnvironmentId: runtimeEnvironment.id,
+              projectId: project.id,
+            },
+          });
+        }
+
+        resolve(undefined);
+      }).then(() => {});
+
+      const received: Transaction<TaskRunData>[] = [];
+      const iterator = subscription.stream[Symbol.asyncIterator]();
+
+      // Pull the first item, then wait before pulling the next
+      const first = await iterator.next();
+      received.push(first.value);
+
+      // Wait to simulate slow consumer
+      await setTimeout(2000);
+
+      // Pull the next item
+      const second = await iterator.next();
+      received.push(second.value);
+
+      // Optionally, check internal state or spy on client.subscribe/data to ensure only 1 item was buffered at a time
+
+      expect(received.length).toBe(2);
+
+      // Clean up
+      await subscription.client.stop();
+    }
+  );
 });