move document subscription to be managed by subscribeToFindMany

Author: nikgraf

packages/hypergraph-react/src/HypergraphSpaceContext.tsx

@@ -3,7 +3,7 @@
 import type { AnyDocumentId, DocHandle, Repo } from '@automerge/automerge-repo';
 import { useRepo } from '@automerge/automerge-repo-react-hooks';
 import { Entity, Utils } from '@graphprotocol/hypergraph';
-import { type DocumentContent, subscribeToDocumentChanges } from '@graphprotocol/hypergraph/Entity';
+import type { DocumentContent } from '@graphprotocol/hypergraph/Entity';
 import * as Schema from 'effect/Schema';
 import { type ReactNode, createContext, useContext, useEffect, useRef, useState, useSyncExternalStore } from 'react';
 
@@ -12,7 +12,6 @@ export type HypergraphContext = {
   repo: Repo;
   id: AnyDocumentId;
   handle: DocHandle<Entity.DocumentContent>;
-  unsubscribeChangeListener: () => void;
 };
 
 export const HypergraphReactContext = createContext<HypergraphContext | undefined>(undefined);
@@ -32,28 +31,17 @@ export function HypergraphProvider({ space, children }: { space: string; childre
 
   let current = ref.current;
   if (current === undefined || space !== current.space || repo !== current.repo) {
-    current?.unsubscribeChangeListener(); // unsubscribe from the previous space when switching to a new space
-
     const id = Utils.idToAutomergeId(space) as AnyDocumentId;
     const handle = repo.find<DocumentContent>(id);
-    const unsubscribeChangeListener = subscribeToDocumentChanges(handle);
 
     current = ref.current = {
       space,
       repo,
       id,
       handle,
-      unsubscribeChangeListener,
     };
   }
 
-  // biome-ignore lint/correctness/useExhaustiveDependencies: no need for dependencies as the unsubscribe is called from the ref
-  useEffect(() => {
-    return () => {
-      current?.unsubscribeChangeListener(); // unsubscribe from the previous space when the component unmounts
-    };
-  }, []);
-
   return <HypergraphReactContext.Provider value={current}>{children}</HypergraphReactContext.Provider>;
 }
 

packages/hypergraph/src/Entity.ts

@@ -101,6 +101,14 @@ type DecodedEntitiesCache = Map<
 
 const decodedEntitiesCache: DecodedEntitiesCache = new Map();
 
+const documentChangeListener: {
+  subscribedQueriesCount: number;
+  unsubscribe: undefined | (() => void);
+} = {
+  subscribedQueriesCount: 0,
+  unsubscribe: undefined,
+};
+
 export const subscribeToDocumentChanges = (handle: DocHandle<DocumentContent>) => {
   const onChange = ({ patches, doc }: { patches: Array<Patch>; doc: DocumentContent }) => {
     const changedEntities = new Set<string>();
@@ -327,9 +335,16 @@ export function subscribeToFindMany<const S extends AnyNoContext>(
 
   const unsubscribe = () => {
     const query = decodedEntitiesCache.get(typeName)?.queries.get(queryKey);
-    if (!query || !query.listeners) return;
-    query.listeners = query.listeners.filter((cachedListener) => cachedListener !== listener);
-    console.log('unsubscribe query', query.listeners);
+    if (query?.listeners) {
+      query.listeners = query.listeners.filter((cachedListener) => cachedListener !== listener);
+      console.log('unsubscribe query', query.listeners);
+    }
+
+    documentChangeListener.subscribedQueriesCount--;
+    if (documentChangeListener.subscribedQueriesCount === 0) {
+      documentChangeListener.unsubscribe?.();
+      documentChangeListener.unsubscribe = undefined;
+    }
   };
 
   const entities = findMany(handle, type);
@@ -375,6 +390,11 @@ export function subscribeToFindMany<const S extends AnyNoContext>(
     });
   }
 
+  if (documentChangeListener.subscribedQueriesCount === 0) {
+    documentChangeListener.unsubscribe = subscribeToDocumentChanges(handle);
+  }
+  documentChangeListener.subscribedQueriesCount++;
+
   return { listener, getEntities, unsubscribe };
 }