Fix: DataStream Race Condition (and flaky unit tests) (#819)

Author: stevensJourney

.changeset/nice-pillows-juggle.md

@@ -0,0 +1,5 @@
+---
+'@powersync/common': patch
+---
+
+Fixed delayed streaming processing, due to a race condition, when connecting via the HTTP connection method (could potentially also affect WebSockets).

packages/common/src/client/sync/stream/AbstractRemote.ts

@@ -573,7 +573,11 @@ export abstract class AbstractRemote {
 
     const stream = new DataStream<T, string>({
       logger: this.logger,
-      mapLine: mapLine
+      mapLine: mapLine,
+      pressure: {
+        highWaterMark: 20,
+        lowWaterMark: 10
+      }
     });
 
     abortSignal?.addEventListener('abort', () => {
@@ -585,46 +589,54 @@ export abstract class AbstractRemote {
     let buffer = '';
 
 
+    const consumeStream = async () => {
+      while (!stream.closed && !abortSignal?.aborted && !readerReleased) {
+        const { done, value } = await reader.read();
+        if (done) {
+          const remaining = buffer.trim();
+          if (remaining.length != 0) {
+            stream.enqueueData(remaining);
+          }
 
-    const l = stream.registerListener({
-      lowWater: async () => {
-        if (stream.closed || abortSignal?.aborted || readerReleased) {
-          return
-        } 
-        try {
-          let didCompleteLine = false;
-          while (!didCompleteLine) {
-            const { done, value } = await reader.read();
-            if (done) {
-              const remaining = buffer.trim();
-              if (remaining.length != 0) {
-                stream.enqueueData(remaining);
-              }
+          stream.close();
+          await closeReader();
+          return;
+        }
 
-              stream.close();
-              await closeReader();
-              return;
-            }
+        const data = decoder.decode(value, { stream: true });
+        buffer += data;
 
-            const data = decoder.decode(value, { stream: true });
-            buffer += data;
+        const lines = buffer.split('\n');
+        for (var i = 0; i < lines.length - 1; i++) {
+          var l = lines[i].trim();
+          if (l.length > 0) {
+            stream.enqueueData(l);
+          }
+        }
 
-            const lines = buffer.split('\n');
-            for (var i = 0; i < lines.length - 1; i++) {
-              var l = lines[i].trim();
-              if (l.length > 0) {
-                stream.enqueueData(l);
-                didCompleteLine = true;
+        buffer = lines[lines.length - 1];
+
+        // Implement backpressure by waiting for the low water mark to be reached
+        if (stream.dataQueue.length > stream.highWatermark) {
+          await new Promise<void>((resolve) => {
+            const dispose = stream.registerListener({
+              lowWater: async () => {
+                resolve();
+                dispose();
+              },
+              closed: () => {
+                resolve();
+                dispose();
               }
-            }
-
-            buffer = lines[lines.length - 1];
-          }
-        } catch (ex) {
-          stream.close();
-          throw ex;
+            })
+          })
         }
-      },
+      }
+    }
+
+    consumeStream().catch(ex => this.logger.error('Error consuming stream', ex));
+
+    const l = stream.registerListener({
       closed: () => {
         closeReader();
         l?.();

packages/common/src/client/sync/stream/AbstractStreamingSyncImplementation.ts

@@ -194,8 +194,7 @@ export interface RequiredAdditionalConnectionOptions extends Required<Additional
 }
 
 export interface StreamingSyncImplementation
-  extends BaseObserverInterface<StreamingSyncImplementationListener>,
-    Disposable {
+  extends BaseObserverInterface<StreamingSyncImplementationListener>, Disposable {
   /**
    * Connects to the sync service
    */

packages/common/src/utils/DataStream.ts

@@ -110,6 +110,18 @@ export class DataStream<ParsedData, SourceData = any> extends BaseObserver<DataS
       return null;
     }
 
+    // Wait for any pending processing to complete first.
+    // This ensures we register our listener before calling processQueue(),
+    // avoiding a race where processQueue() sees no reader and returns early.
+    if (this.processingPromise) {
+      await this.processingPromise;
+    }
+
+    // Re-check after await - stream may have closed while we were waiting
+    if (this.closed) {
+      return null;
+    }
+
     return new Promise((resolve, reject) => {
       const l = this.registerListener({
         data: async (data) => {
@@ -151,7 +163,7 @@ export class DataStream<ParsedData, SourceData = any> extends BaseObserver<DataS
 
     const promise = (this.processingPromise = this._processQueue());
     promise.finally(() => {
-      return (this.processingPromise = null);
+      this.processingPromise = null;
     });
     return promise;
   }
@@ -190,7 +202,6 @@ export class DataStream<ParsedData, SourceData = any> extends BaseObserver<DataS
     }
 
     if (this.dataQueue.length > 0) {
-      // Next tick
       setTimeout(() => this.processQueue());
     }
   }

packages/node/tests/sync.test.ts

@@ -257,6 +257,7 @@ function defineSyncTests(impl: SyncClientImplementation) {
         await waitForProgress(database, [0, 2]);
         pushDataLine(syncService, 'a', 2);
         await waitForProgress(database, [2, 2]);
+
         pushCheckpointComplete(syncService);
         await waitForSyncStatus(database, (s) => s.downloadProgress == null);
       },
@@ -973,8 +974,6 @@ async function waitForProgress(
       return false;
     }
 
-    //console.log('checking', progress);
-
     const check = (expected: [number, number], actual: ProgressWithOperations): boolean => {
       return actual.downloadedOperations == expected[0] && actual.totalOperations == expected[1];
     };
@@ -985,7 +984,6 @@ async function waitForProgress(
 
     for (const [priority, expected] of forPriorities) {
       if (!check(expected, progress.untilPriority(priority))) {
-        //console.log('failed for', priority, expected, progress);
         return false;
       }
     }

packages/web/tests/mockSyncServiceExample.test.ts

@@ -63,17 +63,23 @@ describe('Mock Sync Service Example', { timeout: 100000 }, () => {
         }
       });
 
+      // Wait for sync to complete (hasSynced becomes true) before closing the response
+      await vi.waitFor(
+        () => {
+          expect(database.currentStatus.hasSynced).toBe(true);
+        },
+        { timeout: 5000 }
+      );
+
       // Complete the response
       await mockService.completeResponse(syncRequestId);
 
-      // Wait for sync to complete and verify the data was saved
-      await vi.waitFor(async () => {
-        const rows = await database.getAll('SELECT * FROM lists WHERE id = ?', ['1']);
-        expect(rows).toHaveLength(1);
-        expect(rows[0]).toMatchObject({
-          id: '1',
-          name: 'from server'
-        });
+      // Verify the data was saved
+      const rows = await database.getAll('SELECT * FROM lists WHERE id = ?', ['1']);
+      expect(rows).toHaveLength(1);
+      expect(rows[0]).toMatchObject({
+        id: '1',
+        name: 'from server'
       });
 
       // Verify the data by querying the database