From 1e1c528b0efb6d4a63cb9da03acf0b9b591f03f4 Mon Sep 17 00:00:00 2001
From: George Su <georgewsu@users.noreply.github.com>
Date: Thu, 1 Aug 2024 14:26:34 -0700
Subject: [PATCH] Issue 330: Support removeDocument from synchronizer and
 isDocSubscribedTo (#2)

Co-authored-by: George Su <gsu@blend.com>
---
 .../src/RemoteHeadsSubscriptions.ts              |  7 +++++++
 .../src/synchronizer/CollectionSynchronizer.ts   |  5 ++---
 .../test/CollectionSynchronizer.test.ts          | 16 ++++++++++++++++
 .../test/RemoteHeadsSubscriptions.test.ts        | 12 ++++++++----
 4 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/packages/automerge-repo/src/RemoteHeadsSubscriptions.ts b/packages/automerge-repo/src/RemoteHeadsSubscriptions.ts
index 6862c7d5a..fc897c65b 100644
--- a/packages/automerge-repo/src/RemoteHeadsSubscriptions.ts
+++ b/packages/automerge-repo/src/RemoteHeadsSubscriptions.ts
@@ -325,6 +325,13 @@ export class RemoteHeadsSubscriptions extends EventEmitter<RemoteHeadsSubscripti
     }
   }
 
+  isDocSubscribedTo(documentId: DocumentId) {
+    for (const [peerId, subscribedDocs] of this.#subscribedDocsByPeer) {
+      if (subscribedDocs.has(documentId)) return true
+    }
+    return false
+  }
+
   #isPeerSubscribedToDoc(peerId: PeerId, documentId: DocumentId) {
     const subscribedDocs = this.#subscribedDocsByPeer.get(peerId)
     return subscribedDocs && subscribedDocs.has(documentId)
diff --git a/packages/automerge-repo/src/synchronizer/CollectionSynchronizer.ts b/packages/automerge-repo/src/synchronizer/CollectionSynchronizer.ts
index ddfb4820f..47b377dba 100644
--- a/packages/automerge-repo/src/synchronizer/CollectionSynchronizer.ts
+++ b/packages/automerge-repo/src/synchronizer/CollectionSynchronizer.ts
@@ -116,10 +116,9 @@ export class CollectionSynchronizer extends Synchronizer {
     })
   }
 
-  // TODO: implement this
-  // eslint-disable-next-line @typescript-eslint/no-unused-vars
   removeDocument(documentId: DocumentId) {
-    throw new Error("not implemented")
+    delete this.#docSynchronizers[documentId]
+    delete this.#docSetUp[documentId]
   }
 
   /** Adds a peer and maybe starts synchronizing with them */
diff --git a/packages/automerge-repo/test/CollectionSynchronizer.test.ts b/packages/automerge-repo/test/CollectionSynchronizer.test.ts
index 7bbf44a24..1005d4f13 100644
--- a/packages/automerge-repo/test/CollectionSynchronizer.test.ts
+++ b/packages/automerge-repo/test/CollectionSynchronizer.test.ts
@@ -73,6 +73,22 @@ describe("CollectionSynchronizer", () => {
 
       synchronizer.addDocument(handle.documentId)
 
+      setTimeout(done)
+    }))
+
+  it("should not synchronize a document removed from synchronizer", () =>
+    new Promise<void>((done, reject) => {
+      const handle = repo.create()
+
+      synchronizer.addDocument(handle.documentId)
+      synchronizer.removeDocument(handle.documentId)
+
+      synchronizer.once("message", () => {
+        reject(new Error("Should not have sent a message"))
+      })
+
+      synchronizer.addPeer("peer2" as PeerId)
+
       setTimeout(done)
     }))
 })
diff --git a/packages/automerge-repo/test/RemoteHeadsSubscriptions.test.ts b/packages/automerge-repo/test/RemoteHeadsSubscriptions.test.ts
index 13d68040e..790e14a29 100644
--- a/packages/automerge-repo/test/RemoteHeadsSubscriptions.test.ts
+++ b/packages/automerge-repo/test/RemoteHeadsSubscriptions.test.ts
@@ -175,7 +175,7 @@ describe("RepoHeadsSubscriptions", () => {
       event: "change-remote-subs",
     })
 
-    // unsubsscribe from storage B
+    // unsubscribe from storage B
     remoteHeadsSubscriptions.unsubscribeFromRemotes([storageB])
 
     // should forward unsubscribe to generous peer
@@ -223,8 +223,12 @@ describe("RepoHeadsSubscriptions", () => {
       emitter: remoteHeadsSubscriptions,
       event: "notify-remote-heads",
     })
+    assert(!remoteHeadsSubscriptions.isDocSubscribedTo(docA))
+    assert(!remoteHeadsSubscriptions.isDocSubscribedTo(docC))
     remoteHeadsSubscriptions.subscribePeerToDoc(peerC, docA)
     remoteHeadsSubscriptions.subscribePeerToDoc(peerC, docC)
+    assert(remoteHeadsSubscriptions.isDocSubscribedTo(docA))
+    assert(remoteHeadsSubscriptions.isDocSubscribedTo(docC))
 
     // change message for docA in storageB
     remoteHeadsSubscriptions.handleRemoteHeads(docAHeadsChangedForStorageB)
@@ -248,7 +252,7 @@ describe("RepoHeadsSubscriptions", () => {
 
     // unsubscribe peer C
     remoteHeadsSubscriptions.handleControlMessage(unsubscribePeerCFromStorageB)
-    const messagesAfteUnsubscribePromise = collectMessages({
+    const messagesAfterUnsubscribePromise = collectMessages({
       emitter: remoteHeadsSubscriptions,
       event: "notify-remote-heads",
     })
@@ -256,8 +260,8 @@ describe("RepoHeadsSubscriptions", () => {
     // heads of docB for storageB change
     remoteHeadsSubscriptions.handleRemoteHeads(docBHeadsChangedForStorageB)
 
-    // expect not to be be notified
-    messages = await messagesAfteUnsubscribePromise
+    // expect not to be notified
+    messages = await messagesAfterUnsubscribePromise
     assert.strictEqual(messages.length, 0)
   })