Fix advisory lock leak in PostgresMessageQueue

dahlia · claude · dahlia · commit 2fb9d6a6b435 · 2026-01-27T20:21:21.000+09:00
The original implementation called pg_try_advisory_lock() inside a WHERE clause of a subquery. PostgreSQL evaluates WHERE conditions for each row, causing the advisory lock to be acquired multiple times when there were multiple messages with the same ordering_key. However, pg_advisory_unlock() was only called once after processing, leaving the lock partially held due to PostgreSQL's reentrant lock counter. This rewrites the listen() method to: - Process messages WITHOUT ordering key first using a CTE with FOR UPDATE SKIP LOCKED - Process messages WITH ordering key by: 1. Selecting a candidate message 2. Calling pg_try_advisory_lock() in a separate query (exactly once) 3. If lock acquired, delete and process the message 4. Always release the lock in a finally block 5. If lock not acquired, try other candidates with different ordering keys This ensures the advisory lock is acquired exactly once per ordering key, so pg_advisory_unlock() properly releases it. Fixes #538 Co-Authored-By: Claude <noreply@anthropic.com>
diff --git a/packages/postgres/src/mq.test.ts b/packages/postgres/src/mq.test.ts
@@ -1,7 +1,9 @@
 import { test } from "@fedify/fixture";
 import { PostgresMessageQueue } from "@fedify/postgres/mq";
 import { getRandomKey, testMessageQueue } from "@fedify/testing";
+import { deepStrictEqual } from "node:assert/strict";
 import process from "node:process";
+import { test as nodeTest } from "node:test";
 import postgres from "postgres";
 
 const dbUrl = process.env.POSTGRES_URL;
@@ -35,4 +37,121 @@ test("PostgresMessageQueue", { ignore: dbUrl == null }, () => {
   );
 });
 
+// Regression test for advisory lock not being fully released after processing
+// a message with an ordering key.  This test verifies that after processing
+// a message through PostgresMessageQueue.listen(), the advisory lock is fully
+// released and another session can immediately acquire it.
+//
+// The original bug: pg_try_advisory_lock() in a WHERE clause was called
+// multiple times during query execution (once per row with the same
+// ordering_key), causing the lock's reentrant counter to be > 1.  But
+// pg_advisory_unlock() was only called once, leaving the lock partially held.
+//
+// To reproduce the bug, we need MULTIPLE messages with the SAME ordering key.
+// This causes the WHERE clause to evaluate pg_try_advisory_lock() for each row,
+// incrementing the lock counter multiple times.
+//
+// See: https://github.com/fedify-dev/fedify/issues/538
+nodeTest(
+  "PostgresMessageQueue advisory lock release",
+  { skip: dbUrl == null },
+  async () => {
+    const tableName = getRandomKey("message");
+    const channelName = getRandomKey("channel");
+
+    // Use two separate connections to verify lock behavior across sessions
+    const sql1 = postgres(dbUrl!);
+    const sql2 = postgres(dbUrl!);
+
+    const mq = new PostgresMessageQueue(sql1, {
+      tableName,
+      channelName,
+      pollInterval: { milliseconds: 100 },
+    });
+
+    try {
+      await mq.initialize();
+
+      // CRITICAL: Enqueue MULTIPLE messages with the SAME ordering key.
+      // This is what triggers the bug - the WHERE clause evaluates
+      // pg_try_advisory_lock() for each row, incrementing the lock counter
+      // multiple times, but pg_advisory_unlock() is only called once.
+      const orderingKey = "test-ordering-key";
+      await mq.enqueue({ value: 1 }, { orderingKey });
+      await mq.enqueue({ value: 2 }, { orderingKey });
+      await mq.enqueue({ value: 3 }, { orderingKey });
+
+      // Track when the FIRST message is processed
+      let firstMessageProcessed = false;
+      let lockReleasedAfterProcessing = false;
+
+      const controller = new AbortController();
+
+      // Start listening - we only care about processing the FIRST message
+      const listening = mq.listen(
+        () => {
+          if (!firstMessageProcessed) {
+            firstMessageProcessed = true;
+            // Abort after processing first message to stop the listener
+            controller.abort();
+          }
+        },
+        { signal: controller.signal },
+      );
+
+      // Wait for the first message to be processed
+      const startTime = Date.now();
+      while (!firstMessageProcessed && Date.now() - startTime < 10000) {
+        await new Promise((resolve) => setTimeout(resolve, 100));
+      }
+      deepStrictEqual(
+        firstMessageProcessed,
+        true,
+        "First message should be processed",
+      );
+
+      // Wait for listener to fully stop and release locks
+      await listening;
+      await new Promise((resolve) => setTimeout(resolve, 200));
+
+      // THE KEY TEST: After processing ONE message (with multiple messages
+      // having the same ordering key in the queue), the advisory lock should
+      // be FULLY released.  With the bug, the lock counter would be > 0
+      // because pg_try_advisory_lock was called N times but pg_advisory_unlock
+      // was only called once.
+      const lockAfterProcessing = await sql2`
+        SELECT pg_try_advisory_lock(
+          hashtext(${tableName}),
+          hashtext(${orderingKey})
+        ) AS acquired
+      `;
+      lockReleasedAfterProcessing = lockAfterProcessing[0].acquired;
+
+      // Release the lock we just acquired in sql2
+      if (lockReleasedAfterProcessing) {
+        await sql2`
+          SELECT pg_advisory_unlock(
+            hashtext(${tableName}),
+            hashtext(${orderingKey})
+          )
+        `;
+      }
+
+      // THE FIX: After processing, the lock should be fully released
+      // and another session should be able to acquire it
+      deepStrictEqual(
+        lockReleasedAfterProcessing,
+        true,
+        "Lock should be fully released after message processing " +
+          "(bug: lock counter was incremented multiple times but only " +
+          "decremented once)",
+      );
+    } finally {
+      await mq.drop();
+      await sql1.end();
+      await sql2.end();
+    }
+  },
+);
+
 // cspell: ignore sqls
diff --git a/packages/postgres/src/mq.ts b/packages/postgres/src/mq.ts
@@ -161,50 +161,111 @@ export class PostgresMessageQueue implements MessageQueue {
     const { signal } = options;
     const poll = async () => {
       while (!signal?.aborted) {
-        // Use PostgreSQL advisory locks for distributed ordering key locking.
-        // pg_try_advisory_lock returns true if the lock was acquired, false
-        // otherwise.  We use hashtext() to convert the table name and ordering
-        // key to integers for the lock ID.  Messages without an ordering key
-        // (null) can always be processed.
+        let processed = false;
+
+        // Step 1: Try to process messages without ordering key first.
+        // These don't need advisory locks.
         const query = this.#sql`
-          DELETE FROM ${this.#sql(this.#tableName)}
-          WHERE id = (
-            SELECT id
+          WITH candidate AS (
+            SELECT id, ordering_key
             FROM ${this.#sql(this.#tableName)}
             WHERE created + delay < CURRENT_TIMESTAMP
-              AND (ordering_key IS NULL
-                   OR pg_try_advisory_lock(
-                        hashtext(${this.#tableName}),
-                        hashtext(ordering_key)
-                      ))
+              AND ordering_key IS NULL
             ORDER BY created
             LIMIT 1
+            FOR UPDATE SKIP LOCKED
           )
+          DELETE FROM ${this.#sql(this.#tableName)}
+          WHERE id IN (SELECT id FROM candidate)
           RETURNING message, ordering_key;
         `.execute();
         const cancel = query.cancel.bind(query);
         signal?.addEventListener("abort", cancel);
-        let i = 0;
         for (const row of await query) {
           if (signal?.aborted) return;
-          const orderingKey = row.ordering_key as string | null;
-          try {
-            await handler(row.message);
-          } finally {
-            // Release the distributed advisory lock if we acquired one
-            if (orderingKey != null) {
+          await handler(row.message);
+          processed = true;
+        }
+        signal?.removeEventListener("abort", cancel);
+
+        // If we processed a message without ordering key, continue the loop
+        if (processed) continue;
+
+        // Step 2: Try to process a message with an ordering key.
+        // We do this separately to ensure pg_try_advisory_lock is called
+        // exactly once per attempt.
+        // We loop through candidates until we find one we can lock, or run out.
+        const attemptedOrderingKeys = new Set<string>();
+        while (!signal?.aborted) {
+          // Find a candidate with ordering key that we haven't tried yet
+          const candidateResult = await this.#sql`
+            SELECT id, ordering_key
+            FROM ${this.#sql(this.#tableName)}
+            WHERE created + delay < CURRENT_TIMESTAMP
+              AND ordering_key IS NOT NULL
+              ${
+            attemptedOrderingKeys.size > 0
+              ? this.#sql`AND ordering_key NOT IN ${
+                this.#sql([...attemptedOrderingKeys])
+              }`
+              : this.#sql``
+          }
+            ORDER BY created
+            LIMIT 1
+          `;
+
+          if (candidateResult.length === 0) {
+            // No more candidates to try
+            break;
+          }
+
+          const candidate = candidateResult[0];
+          const candidateId = candidate.id as string;
+          const orderingKey = candidate.ordering_key as string;
+          attemptedOrderingKeys.add(orderingKey);
+
+          // Try to acquire the advisory lock (exactly once)
+          const lockResult = await this.#sql`
+            SELECT pg_try_advisory_lock(
+              hashtext(${this.#tableName}),
+              hashtext(${orderingKey})
+            ) AS acquired
+          `;
+
+          if (lockResult[0].acquired) {
+            try {
+              // We have the lock, now delete and process the message
+              const deleteResult = await this.#sql`
+                DELETE FROM ${this.#sql(this.#tableName)}
+                WHERE id = ${candidateId}
+                RETURNING message, ordering_key
+              `;
+
+              for (const row of deleteResult) {
+                if (signal?.aborted) return;
+                await handler(row.message);
+                processed = true;
+              }
+            } finally {
+              // Always release the advisory lock
               await this.#sql`
                 SELECT pg_advisory_unlock(
                   hashtext(${this.#tableName}),
                   hashtext(${orderingKey})
-                );
+                )
               `;
             }
+            // If we processed a message, continue the outer loop
+            if (processed) break;
           }
-          i++;
+          // Lock not acquired, try next candidate with different ordering key
         }
-        signal?.removeEventListener("abort", cancel);
-        if (i < 1) break;
+
+        // If we processed a message, continue the outer loop
+        if (processed) continue;
+
+        // No messages to process, exit the loop
+        break;
       }
     };
     const timeouts = new Set<ReturnType<typeof setTimeout>>();
