fix(store-sync): setup message listener before setting up subscription (#3765)

Author: alvrs

.changeset/sharp-shrimps-dream.md

@@ -0,0 +1,5 @@
+---
+"@latticexyz/store-sync": patch
+---
+
+Fixed a race condition in the preconfirmed logs stream by setting up the message listener before setting up the subscription.

packages/store-sync/package.json

@@ -70,7 +70,7 @@
     "clean:js": "shx rm -rf dist",
     "dev": "tsup --watch",
     "lint": "eslint .",
-    "playground": "DEBUG=mud:* tsx playground/index.ts",
+    "playground": "DEBUG=mud:* tsx playground/index.ts | tee playground.log",
     "test": "vitest",
     "test:ci": "vitest --run"
   },

packages/store-sync/src/createPreconfirmedBlockStream.ts

@@ -73,11 +73,13 @@ export function createPreconfirmedBlockStream(opts: PreconfirmedBlockStreamOptio
 
   let preconfirmedTransactionLogs: { [txHash: string]: Partial<StoreEventsLog>[] | undefined } = {};
   let preconfirmedLogsState: "initializing" | "initialized" | "waiting" = "waiting";
+  let firstPreconfirmedBlockNumber: bigint | undefined = undefined;
   let attempt = 0;
   const preconfirmedBlockLogs$ = recreatePreconfirmedStream$.pipe(
     tap(() => {
       if (attempt !== 0) debug(`waiting ${attempt * 500}ms before initializing preconfirmed logs stream`);
       preconfirmedLogsState = "initializing";
+      firstPreconfirmedBlockNumber = undefined;
       preconfirmedTransactionLogs = {};
     }),
     switchMap(() => timer(attempt * 500)),
@@ -114,6 +116,10 @@ export function createPreconfirmedBlockStream(opts: PreconfirmedBlockStreamOptio
       return !isProcessedBlock;
     }),
     tap((block) => {
+      if (preconfirmedLogsState !== "initialized") {
+        firstPreconfirmedBlockNumber = block.blockNumber;
+        debug("first preconfirmed block number", firstPreconfirmedBlockNumber);
+      }
       debug("preconfirmed block", block.blockNumber, "with", block.logs.length, "logs");
       preconfirmedLogsState = "initialized";
       attempt = 0;
@@ -125,6 +131,16 @@ export function createPreconfirmedBlockStream(opts: PreconfirmedBlockStreamOptio
         }
         (preconfirmedTransactionLogs[txHash] ??= []).push(log);
       });
+      debug(
+        "preconfirmed logs state",
+        Object.fromEntries(
+          Object.entries(preconfirmedTransactionLogs).map(([txHash, logs]) => [
+            txHash,
+            { numLogs: logs?.length, blockNumber: logs?.[0]?.blockNumber },
+          ]),
+        ),
+        "transactions",
+      );
     }),
   );
 
@@ -133,7 +149,12 @@ export function createPreconfirmedBlockStream(opts: PreconfirmedBlockStreamOptio
       processedLatestBlockNumber = block.blockNumber;
 
       const mismatchingTransactions: string[] = [];
-      if (preconfirmedLogsState === "initialized") {
+      const confirmPreconfirmedLogs =
+        preconfirmedLogsState === "initialized" &&
+        firstPreconfirmedBlockNumber &&
+        block.blockNumber >= firstPreconfirmedBlockNumber;
+
+      if (confirmPreconfirmedLogs) {
         const logsByTransaction = groupBy(
           block.logs.filter((log) => log.transactionHash) as StoreEventsLog[],
           (log) => log.transactionHash,
@@ -148,7 +169,7 @@ export function createPreconfirmedBlockStream(opts: PreconfirmedBlockStreamOptio
               JSON.stringify(
                 {
                   txHash,
-                  numPreconfirmedLogs: preconfirmedLogs?.length,
+                  numPreconfirmedLogs: preconfirmedLogs?.length ?? "none",
                   numLatestLogs: latestLogs.length,
                   missingLogs: latestLogs.filter(
                     (log) => !preconfirmedLogs?.find((preconfirmedLog) => log.logIndex === preconfirmedLog.logIndex),
@@ -190,14 +211,14 @@ export function createPreconfirmedBlockStream(opts: PreconfirmedBlockStreamOptio
       }
 
       // While the preconfirmed logs stream is initializing, don't recreate it and pass the block through
-      if (preconfirmedLogsState === "initializing") {
-        debug("preconfirmed logs stream is initializing, not recreating");
+      if (!confirmPreconfirmedLogs) {
+        debug("block is before first preconfirmed block, not recreating preconfirmed stream");
         return block;
       }
 
       // If the preconfirmed logs stream is initialized but there are mismatching logs, recreate it and pass the block through.
       // Pass all logs from this block, not just the mismatching ones, to make sure they appear in the right order.
-      if (preconfirmedLogsState === "initialized" && mismatchingTransactions.length > 0) {
+      if (mismatchingTransactions.length > 0) {
         debug("mismatching transactions found in latest block", block.blockNumber, "recreating preconfirmed stream");
         recreatePreconfirmedStream$.next();
         return block;

packages/store-sync/src/watchLogs.ts

@@ -43,6 +43,7 @@ export function watchLogs({ url, address, fromBlock }: WatchLogsInput): WatchLog
       // Buffer the live logs received until the gap from `fromBlock` to `currentBlock` is closed
       let caughtUp = false;
       const logBuffer: StoreEventsLog[] = [];
+      let subscriptionId: Hex | undefined = undefined;
 
       ws = new WebSocket(url);
 
@@ -78,13 +79,6 @@ export function watchLogs({ url, address, fromBlock }: WatchLogsInput): WatchLog
         debug(`ws${wsId} upgrade`);
       });
 
-      const subscriptionId = await request<Hex>({
-        ws,
-        method: "wiresaw_watchLogs",
-        params: [{ address, topics }],
-        wsId,
-      });
-
       ws.on("message", (message) => {
         const data = JSON.parse(message.toString());
 
@@ -112,6 +106,13 @@ export function watchLogs({ url, address, fromBlock }: WatchLogsInput): WatchLog
         debug(`ws${wsId} message`);
       });
 
+      subscriptionId = await request<Hex>({
+        ws,
+        method: "wiresaw_watchLogs",
+        params: [{ address, topics }],
+        wsId,
+      });
+
       // Catch up to the pending logs
       fetchInitialLogs({ ws, address, fromBlock, topics, wsId })
         .then((initialLogs) => {