Repository: electric-sql/pglite
Pull request #592: fix(pglite/live) handle race condition in live query unsubscription
Author: login: samwillis
Created: 2025-04-12T21:37:16Z | Merged: 2025-04-14T08:13:02Z

This fixes a race condition in live query unsubscription that could result in live queries failing to update.

The accompanying test fails on the previous version, and passes with this change.

In essence, if you start crating a new live query, but don't await it's subscription - so don't await the query it runs internally at the end - you can interweave a unsubscription of a previous live query. This can result in the unsubscribe removing the listener that has just been added for the new live query.

Technically this is a bug not in the live query (or the react useLiveQuery), but in the `listen` and `unlisted` methods, but it makes sense to test via a live query.


Diff:
diff --git a/.changeset/strange-news-burn.md b/.changeset/strange-news-burn.md
new file mode 100644
index 000000000..837b53454
--- /dev/null
+++ b/.changeset/strange-news-burn.md
@@ -0,0 +1,5 @@
+---
+'@electric-sql/pglite': patch
+---
+
+Fix a race condition in live query unsubscription that could result in live queries failing to update.
\ No newline at end of file
diff --git a/packages/pglite/src/pglite.ts b/packages/pglite/src/pglite.ts
index 111297604..d123f9df3 100644
--- a/packages/pglite/src/pglite.ts
+++ b/packages/pglite/src/pglite.ts
@@ -55,6 +55,7 @@ export class PGlite
 
   #queryMutex = new Mutex()
   #transactionMutex = new Mutex()
+  #listenMutex = new Mutex()
   #fsSyncMutex = new Mutex()
   #fsSyncScheduled = false
 
@@ -728,6 +729,10 @@ export class PGlite
    * @param callback The callback to call when a notification is received
    */
   async listen(channel: string, callback: (payload: string) => void) {
+    return this._runExclusiveListen(() => this.#listen(channel, callback))
+  }
+
+  async #listen(channel: string, callback: (payload: string) => void) {
     const pgChannel = toPostgresName(channel)
     if (!this.#notifyListeners.has(pgChannel)) {
       this.#notifyListeners.set(pgChannel, new Set())
@@ -753,16 +758,26 @@ export class PGlite
    * @param callback The callback to remove
    */
   async unlisten(channel: string, callback?: (payload: string) => void) {
+    return this._runExclusiveListen(() => this.#unlisten(channel, callback))
+  }
+
+  async #unlisten(channel: string, callback?: (payload: string) => void) {
     const pgChannel = toPostgresName(channel)
+    const cleanUp = async () => {
+      await this.exec(`UNLISTEN ${channel}`)
+      // While that query was running, another query might have subscribed
+      // so we need to check again
+      if (this.#notifyListeners.get(pgChannel)?.size === 0) {
+        this.#notifyListeners.delete(pgChannel)
+      }
+    }
     if (callback) {
       this.#notifyListeners.get(pgChannel)?.delete(callback)
       if (this.#notifyListeners.get(pgChannel)?.size === 0) {
-        await this.exec(`UNLISTEN ${channel}`)
-        this.#notifyListeners.delete(pgChannel)
+        await cleanUp()
       }
     } else {
-      await this.exec(`UNLISTEN ${channel}`)
-      this.#notifyListeners.delete(pgChannel)
+      await cleanUp()
     }
   }
 
@@ -821,4 +836,8 @@ export class PGlite
     const dump = await this.dumpDataDir('none')
     return PGlite.create({ loadDataDir: dump })
   }
+
+  _runExclusiveListen<T>(fn: () => Promise<T>): Promise<T> {
+    return this.#listenMutex.runExclusive(fn)
+  }
 }
diff --git a/packages/pglite/tests/live.test.ts b/packages/pglite/tests/live.test.ts
index 2acacacdd..e57bf4415 100644
--- a/packages/pglite/tests/live.test.ts
+++ b/packages/pglite/tests/live.test.ts
@@ -1225,5 +1225,52 @@ await testEsmAndCjs(async (importType) => {
         }),
       ).rejects.toThrow('offset and limit must be numbers')
     })
+
+    it("doesn't have a race condition when unsubscribing from a live query", async () => {
+      const db = await PGlite.create({
+        extensions: { live },
+      })
+
+      await db.exec(`
+        CREATE TABLE IF NOT EXISTS testTable (
+          id SERIAL PRIMARY KEY,
+          number INT
+        );
+      `)
+
+      let results
+      const eventTarget = new EventTarget()
+
+      // Create a first live query
+      const { unsubscribe } = await db.live.query({
+        query: 'SELECT * FROM testTable WHERE number > 1',
+      })
+
+      // Create a second live query that listens to the same channel, but don't
+      // await it's subscription to complete.
+      // This creates a race condition with the following unsubscribe, both are
+      // potentially interlaced with each other.
+      db.live.query({
+        query: 'SELECT * FROM testTable WHERE number > 2',
+        callback: (result) => {
+          results = result
+          eventTarget.dispatchEvent(new Event('change'))
+        },
+      })
+
+      // Unsubscribe from the first live query
+      await unsubscribe()
+
+      const promise = new Promise((resolve) =>
+        eventTarget.addEventListener('change', resolve, { once: true }),
+      )
+
+      await db.exec('INSERT INTO testTable (number) VALUES (3);')
+
+      // In a failed test, this will never resolve.
+      await promise
+
+      expect(results.rows).toEqual([{ id: 1, number: 3 }])
+    }, 3000)
   })
 })


Reviews:
user: login: tdrz
body: Unsure about the `unlisten` from all. We should either update the docs to state the there is undefined behavior if you don't wait for the `unlisten` to finish; or not allow any new subscriptions to the channel until any pending `unlisten` actions have been completed.
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-04-13T18:14:23Z
user: login: samwillis
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-13T18:18:34Z
user: login: tdrz
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-13T18:23:57Z
user: login: samwillis
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-13T18:39:19Z
user: login: tdrz
body: LGTM!
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-14T08:06:31Z

Review comments:
diff_hunk: @@ -754,15 +754,21 @@ export class PGlite
    */
   async unlisten(channel: string, callback?: (payload: string) => void) {
     const pgChannel = toPostgresName(channel)
+    const cleanUp = async () => {
+      await this.exec(`UNLISTEN ${channel}`)
+      // While that query was running, another query might have subscribed
+      // so we need to check again
+      if (this.#notifyListeners.get(pgChannel)?.size === 0) {
+        this.#notifyListeners.delete(pgChannel)
+      }
+    }
     if (callback) {
       this.#notifyListeners.get(pgChannel)?.delete(callback)
       if (this.#notifyListeners.get(pgChannel)?.size === 0) {
-        await this.exec(`UNLISTEN ${channel}`)
-        this.#notifyListeners.delete(pgChannel)
+        await cleanUp()
       }
     } else {
-      await this.exec(`UNLISTEN ${channel}`)
-      this.#notifyListeners.delete(pgChannel)
+      await cleanUp()
path: packages/pglite/src/pglite.ts
user: login: tdrz
body: Unsure if this will yield the expected result in every situation. Depending on when the "UNLISTEN" is actually executed, if another query subscribes for the same channel at the worst moment, it might never get notified (because of the unlisten having been run after the listen).
created_at: 2025-04-13T18:12:20Z
author_association: COLLABORATOR
diff_hunk: @@ -754,15 +754,21 @@ export class PGlite
    */
   async unlisten(channel: string, callback?: (payload: string) => void) {
     const pgChannel = toPostgresName(channel)
+    const cleanUp = async () => {
+      await this.exec(`UNLISTEN ${channel}`)
+      // While that query was running, another query might have subscribed
+      // so we need to check again
+      if (this.#notifyListeners.get(pgChannel)?.size === 0) {
+        this.#notifyListeners.delete(pgChannel)
+      }
+    }
     if (callback) {
       this.#notifyListeners.get(pgChannel)?.delete(callback)
       if (this.#notifyListeners.get(pgChannel)?.size === 0) {
-        await this.exec(`UNLISTEN ${channel}`)
-        this.#notifyListeners.delete(pgChannel)
+        await cleanUp()
       }
     } else {
-      await this.exec(`UNLISTEN ${channel}`)
-      this.#notifyListeners.delete(pgChannel)
+      await cleanUp()
path: packages/pglite/src/pglite.ts
user: login: samwillis
body: Hmm, yep. I wander if we should wrap the function in a mutex so that there is only aver one concurrent call?
created_at: 2025-04-13T18:18:34Z
author_association: COLLABORATOR
diff_hunk: @@ -754,15 +754,21 @@ export class PGlite
    */
   async unlisten(channel: string, callback?: (payload: string) => void) {
     const pgChannel = toPostgresName(channel)
+    const cleanUp = async () => {
+      await this.exec(`UNLISTEN ${channel}`)
+      // While that query was running, another query might have subscribed
+      // so we need to check again
+      if (this.#notifyListeners.get(pgChannel)?.size === 0) {
+        this.#notifyListeners.delete(pgChannel)
+      }
+    }
     if (callback) {
       this.#notifyListeners.get(pgChannel)?.delete(callback)
       if (this.#notifyListeners.get(pgChannel)?.size === 0) {
-        await this.exec(`UNLISTEN ${channel}`)
-        this.#notifyListeners.delete(pgChannel)
+        await cleanUp()
       }
     } else {
-      await this.exec(`UNLISTEN ${channel}`)
-      this.#notifyListeners.delete(pgChannel)
+      await cleanUp()
path: packages/pglite/src/pglite.ts
user: login: tdrz
body: Yes, sounds reasonable. I think just doing it per channel is an overkill, so having a global mutex sounds like a good solution to me.
created_at: 2025-04-13T18:23:57Z
author_association: COLLABORATOR
diff_hunk: @@ -754,15 +754,21 @@ export class PGlite
    */
   async unlisten(channel: string, callback?: (payload: string) => void) {
     const pgChannel = toPostgresName(channel)
+    const cleanUp = async () => {
+      await this.exec(`UNLISTEN ${channel}`)
+      // While that query was running, another query might have subscribed
+      // so we need to check again
+      if (this.#notifyListeners.get(pgChannel)?.size === 0) {
+        this.#notifyListeners.delete(pgChannel)
+      }
+    }
     if (callback) {
       this.#notifyListeners.get(pgChannel)?.delete(callback)
       if (this.#notifyListeners.get(pgChannel)?.size === 0) {
-        await this.exec(`UNLISTEN ${channel}`)
-        this.#notifyListeners.delete(pgChannel)
+        await cleanUp()
       }
     } else {
-      await this.exec(`UNLISTEN ${channel}`)
-      this.#notifyListeners.delete(pgChannel)
+      await cleanUp()
path: packages/pglite/src/pglite.ts
user: login: samwillis
body: Agreed, we have the mutex implementation already for the query and transaction mutexes.

We could maybe just use the transaction mutex?

If this was a new mutex, and then run within a transaction that rolled back, it would leave the subscription ini a broken state.

Not being able to call this api from within the transaction seems fair.

If you would like to pick this up and finish it, please do (first thing Monday). We can then get it published ASAP.
created_at: 2025-04-13T18:39:19Z
author_association: COLLABORATOR

PR comments:
user: login: samwillis
created_at: 2025-04-12T21:44:23Z
body: There is an argument that we should debounce that cleanup, so that in the case of a live query changing it's params we don't unnecessarily call the `UNSUBSCRIBE`, but I think that maybe should be a follow up PR as it brings some complexity of keeping track of timeouts.
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-04-12T21:53:33Z
body: - PGlite: https://github.com/electric-sql/pglite/actions/runs/14423595749/artifacts/2933619915
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-04-12T21:58:09Z
body: - Demos: https://github.com/electric-sql/pglite/actions/runs/14423595749/artifacts/2933626112
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-04-13T11:23:06Z
body: - PGlite: https://github.com/electric-sql/pglite/actions/runs/14428890166/artifacts/2934964545
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-04-13T11:28:07Z
body: - Demos: https://github.com/electric-sql/pglite/actions/runs/14428890166/artifacts/2934971079
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-04-13T11:28:18Z
body: <!-- NETLIFY DEPLOY COMMENT GENERATED BY ACTIONS_NETLIFY - APP ID SHA256: c536feb4306d76e78986ae3b1c9b4e09f389826631a6c0c01aa5e2ce843d4a67 -->
🚀 Deployed on https://67fcc1cfe4f132a9cc6d7c4d--pglite.netlify.app
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-04-14T08:01:00Z
body: - PGlite: https://github.com/electric-sql/pglite/actions/runs/14440185058/artifacts/2937888398
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-04-14T08:05:34Z
body: - Demos: https://github.com/electric-sql/pglite/actions/runs/14440185058/artifacts/2937914910
author_association: CONTRIBUTOR