fix(stash): maintain update atomicity in subscribeQuery (#3663)

Author: alvrs

.changeset/strange-rabbits-matter.md

@@ -0,0 +1,6 @@
+---
+"@latticexyz/stash": patch
+---
+
+Stash now preserves batch updates when subscribing to query results.
+Previously, while Stash supported batching table updates for atomic onchain changes, subscribing to query results would split these updates by table.

packages/stash/src/actions/subscribeQuery.test.ts

@@ -217,6 +217,8 @@ describe("defineQuery", () => {
     // Update multiple records but only advance timer once
     setRecord({ stash, table: Health, key: { player: `0x2` }, value: { health: 2 } });
     setRecord({ stash, table: Health, key: { player: `0x1` }, value: { health: 1 } });
+    setRecord({ stash, table: Position, key: { player: `0x1` }, value: { x: 1, y: 2 } });
+    setRecord({ stash, table: Position, key: { player: `0x2` }, value: { x: 2, y: 3 } });
     vi.advanceTimersToNextTimer();
 
     expect(subscriber).toBeCalledTimes(2);
@@ -263,6 +265,50 @@ describe("defineQuery", () => {
         current: { player: "0x1", health: 1 },
         type: "enter",
       },
+      {
+        table: {
+          label: "Position",
+          type: "table",
+          namespace: "namespace1",
+          namespaceLabel: "namespace1",
+          name: "Position",
+          tableId: "0x74626e616d6573706163653100000000506f736974696f6e0000000000000000",
+          schema: {
+            player: { type: "bytes32", internalType: "bytes32" },
+            x: { type: "int32", internalType: "int32" },
+            y: { type: "int32", internalType: "int32" },
+          },
+          key: ["player"],
+          codegen: { outputDirectory: "tables", tableIdArgument: false, storeArgument: false, dataStruct: true },
+          deploy: { disabled: false },
+        },
+        key: { player: "0x1" },
+        previous: { player: "0x1", x: 1, y: 4 },
+        current: { player: "0x1", x: 1, y: 2 },
+        type: "update",
+      },
+      {
+        table: {
+          label: "Position",
+          type: "table",
+          namespace: "namespace1",
+          namespaceLabel: "namespace1",
+          name: "Position",
+          tableId: "0x74626e616d6573706163653100000000506f736974696f6e0000000000000000",
+          schema: {
+            player: { type: "bytes32", internalType: "bytes32" },
+            x: { type: "int32", internalType: "int32" },
+            y: { type: "int32", internalType: "int32" },
+          },
+          key: ["player"],
+          codegen: { outputDirectory: "tables", tableIdArgument: false, storeArgument: false, dataStruct: true },
+          deploy: { disabled: false },
+        },
+        key: { player: "0x2" },
+        previous: { player: "0x2", x: 2, y: 3 },
+        current: { player: "0x2", x: 2, y: 3 },
+        type: "update",
+      },
     ]);
   });
 });

packages/stash/src/actions/subscribeQuery.ts

@@ -16,9 +16,9 @@ import {
   Key,
   TableRecord,
 } from "../common";
-import { getTable } from "./getTable";
 import { runQuery } from "./runQuery";
 import { encodeKey } from "./encodeKey";
+import { subscribeStash } from "./subscribeStash";
 
 export type SubscribeQueryOptions = CommonQueryOptions;
 
@@ -119,13 +119,16 @@ export function subscribeQuery<query extends Query>({
     subscribers.forEach((subscriber) => subscriber(updates));
   };
 
-  // Subscribe to each table's update stream and return the unsubscribers
-  const unsubsribers = query.map((fragment) =>
-    getTable({ stash, table: fragment.table }).subscribe({
-      subscriber: (updates) => updateQueryResult(updates),
-    }),
-  );
-  const unsubscribe = () => unsubsribers.forEach((unsub) => unsub());
+  // Too maintain atomicity of stash updates we don't subscribe to tables individually but filter global stash updates.
+  const relevantTables = new Set(query.map((fragment) => fragment.table.tableId));
+  const unsubscribe = subscribeStash({
+    stash,
+    subscriber: (event) => {
+      if (event.type !== "records") return;
+      const relevantUpdates = event.updates.filter((update) => relevantTables.has(update.table.tableId));
+      updateQueryResult(relevantUpdates);
+    },
+  });
 
   // Notify initial subscribers
   const updates: QueryUpdates = Object.values(matching).map((key) => ({ key, type: "enter" }));