From c6a52f79f258dbabd50f57401cea8103be756b36 Mon Sep 17 00:00:00 2001
From: Erik Arvidsson <arv@roci.dev>
Date: Tue, 7 Dec 2021 16:22:28 -0800
Subject: [PATCH] fix: Subscriptions with errors never recovered (#757)

Even if we get an exception calling the subscription query body we need
to keep track of the keys.

Closes #754
---
 src/replicache.test.ts | 56 ++++++++++++++++++++++++++++++++++++++++++
 src/replicache.ts      | 10 +++++---
 2 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/src/replicache.test.ts b/src/replicache.test.ts
index cae87286..46910651 100644
--- a/src/replicache.test.ts
+++ b/src/replicache.test.ts
@@ -3042,3 +3042,59 @@ testWithBothStores('client ID is set correctly on transactions', async () => {
 
   await rep.mutate.expectClientID(repClientID);
 });
+
+test('subscription with error in body', async () => {
+  const rep = await replicacheForTesting('subscription-with-error-in-body', {
+    mutators: {
+      addData,
+    },
+  });
+
+  let bodyCallCounter = 0;
+  let errorCounter = 0;
+  const letters = 'abc';
+
+  rep.subscribe(
+    async tx => {
+      bodyCallCounter++;
+      const a = await tx.get('a');
+      if (a === undefined) {
+        throw new Error('a is undefined');
+      }
+      const b = await tx.get('b');
+      if (b === undefined) {
+        throw new Error('b is undefined');
+      }
+      const c = await tx.get('c');
+      if (c === undefined) {
+        throw new Error('c is undefined');
+      }
+      return {a, b, c};
+    },
+    {
+      onData(data) {
+        expect(data).to.deep.equal({a: 1, b: 2, c: 3});
+      },
+      onError(err) {
+        expect(err)
+          .to.be.instanceOf(Error)
+          .with.property(
+            'message',
+            letters[errorCounter++] + ' is undefined',
+            `Error for ${errorCounter} is incorrect`,
+          );
+      },
+    },
+  );
+
+  await tickUntil(() => bodyCallCounter === 1);
+
+  await rep.mutate.addData({a: 1});
+  expect(bodyCallCounter).to.equal(2);
+
+  await rep.mutate.addData({b: 2});
+  expect(bodyCallCounter).to.equal(3);
+
+  await rep.mutate.addData({c: 3});
+  expect(bodyCallCounter).to.equal(4);
+});
diff --git a/src/replicache.ts b/src/replicache.ts
index 78e7df7e..ee0bca82 100644
--- a/src/replicache.ts
+++ b/src/replicache.ts
@@ -914,14 +914,18 @@ export class Replicache<MD extends MutatorDefs = {}> {
     const results = await this._queryInternal(async tx => {
       const promises = subs.map(async s => {
         // Tag the result so we can deal with success vs error below.
+        const stx = new SubscriptionTransactionWrapper(tx);
         try {
-          const stx = new SubscriptionTransactionWrapper(tx);
           const value = await s.body(stx);
-          s.keys = stx.keys;
-          s.scans = stx.scans;
           return {ok: true, value} as R;
         } catch (error) {
           return {ok: false, error} as R;
+        } finally {
+          // We need to keep track of the subscription keys even if there was an
+          // exception because changes to the keys can make the subscription
+          // body succeeed.
+          s.keys = stx.keys;
+          s.scans = stx.scans;
         }
       });
       return await Promise.all(promises);