Fix aborting low-priority syncs

simolus3 · simolus3 · commit 2e17504d69b8 · 2025-02-10T15:35:45.000+01:00
diff --git a/modules/module-mongodb-storage/test/src/__snapshots__/storage_sync.test.ts.snap b/modules/module-mongodb-storage/test/src/__snapshots__/storage_sync.test.ts.snap
@@ -127,6 +127,80 @@ exports[`sync - mongodb > expiring token 2`] = `
 ]
 `;
 
+exports[`sync - mongodb > sync buckets in order 1`] = `
+[
+  {
+    "checkpoint": {
+      "buckets": [
+        {
+          "bucket": "b0[]",
+          "checksum": 920318466,
+          "count": 1,
+          "priority": 2,
+        },
+        {
+          "bucket": "b1[]",
+          "checksum": -1382098757,
+          "count": 1,
+          "priority": 1,
+        },
+      ],
+      "last_op_id": "2",
+      "write_checkpoint": undefined,
+    },
+  },
+  {
+    "data": {
+      "after": "0",
+      "bucket": "b1[]",
+      "data": [
+        {
+          "checksum": 2912868539n,
+          "data": "{"id":"earlier","description":"Test 2"}",
+          "object_id": "earlier",
+          "object_type": "test",
+          "op": "PUT",
+          "op_id": "2",
+          "subkey": "0dfe86bd-d15b-5fd0-9c7b-a31693030ee0",
+        },
+      ],
+      "has_more": false,
+      "next_after": "2",
+    },
+  },
+  {
+    "partial_checkpoint_complete": {
+      "last_op_id": "2",
+      "priority": 1,
+    },
+  },
+  {
+    "data": {
+      "after": "0",
+      "bucket": "b0[]",
+      "data": [
+        {
+          "checksum": 920318466n,
+          "data": "{"id":"t1","description":"Test 1"}",
+          "object_id": "t1",
+          "object_type": "test",
+          "op": "PUT",
+          "op_id": "1",
+          "subkey": "e5aa2ddc-1328-58fa-a000-0b5ed31eaf1a",
+        },
+      ],
+      "has_more": false,
+      "next_after": "1",
+    },
+  },
+  {
+    "checkpoint_complete": {
+      "last_op_id": "2",
+    },
+  },
+]
+`;
+
 exports[`sync - mongodb > sync global data 1`] = `
 [
   {
diff --git a/packages/service-core-tests/src/tests/register-sync-tests.ts b/packages/service-core-tests/src/tests/register-sync-tests.ts
@@ -83,6 +83,159 @@ export function registerSyncTests(factory: storage.TestStorageFactory) {
     expect(lines).toMatchSnapshot();
   });
 
+  test('sync buckets in order', async () => {
+    await using f = await factory();
+
+    const syncRules = await f.updateSyncRules({
+      content: `
+bucket_definitions:
+  b0:
+    priority: 2
+    data:
+      - SELECT * FROM test WHERE LENGTH(id) <= 2;
+  b1:
+    priority: 1
+    data:
+      - SELECT * FROM test WHERE LENGTH(id) > 2;
+    `
+    });
+
+    const bucketStorage = f.getInstance(syncRules);
+    await bucketStorage.autoActivate();
+
+    const result = await bucketStorage.startBatch(test_utils.BATCH_OPTIONS, async (batch) => {
+      await batch.save({
+        sourceTable: TEST_TABLE,
+        tag: storage.SaveOperationTag.INSERT,
+        after: {
+          id: 't1',
+          description: 'Test 1'
+        },
+        afterReplicaId: 't1'
+      });
+
+      await batch.save({
+        sourceTable: TEST_TABLE,
+        tag: storage.SaveOperationTag.INSERT,
+        after: {
+          id: 'earlier',
+          description: 'Test 2'
+        },
+        afterReplicaId: 'earlier'
+      });
+
+      await batch.commit('0/1');
+    });
+
+    const stream = sync.streamResponse({
+      storage: f,
+      params: {
+        buckets: [],
+        include_checksum: true,
+        raw_data: true
+      },
+      parseOptions: test_utils.PARSE_OPTIONS,
+      tracker,
+      syncParams: new RequestParameters({ sub: '' }, {}),
+      token: { exp: Date.now() / 1000 + 10 } as any
+    });
+
+    const lines = await consumeCheckpointLines(stream);
+    expect(lines).toMatchSnapshot();
+  });
+
+  test('sync interrupts low-priority buckets on new checkpoints', async () => {
+    await using f = await factory();
+
+    const syncRules = await f.updateSyncRules({
+      content: `
+bucket_definitions:
+  b0:
+    priority: 2
+    data:
+      - SELECT * FROM test WHERE LENGTH(id) <= 5;
+  b1:
+    priority: 1
+    data:
+      - SELECT * FROM test WHERE LENGTH(id) > 5;
+    `
+    });
+
+    const bucketStorage = f.getInstance(syncRules);
+    await bucketStorage.autoActivate();
+
+    await bucketStorage.startBatch(test_utils.BATCH_OPTIONS, async (batch) => {
+      // Initial data: Add one priority row and 10k low-priority rows.
+      await batch.save({
+        sourceTable: TEST_TABLE,
+        tag: storage.SaveOperationTag.INSERT,
+        after: {
+          id: 'highprio',
+          description: 'High priority row'
+        },
+        afterReplicaId: 'highprio'
+      });
+      for (let i = 0; i < 10_000; i++) {
+        await batch.save({
+          sourceTable: TEST_TABLE,
+          tag: storage.SaveOperationTag.INSERT,
+          after: {
+            id: `${i}`,
+            description: 'low prio'
+          },
+          afterReplicaId: `${i}`
+        });
+      }
+
+      await batch.commit('0/1');
+    });
+
+    const stream = sync.streamResponse({
+      storage: f,
+      params: {
+        buckets: [],
+        include_checksum: true,
+        raw_data: true
+      },
+      parseOptions: test_utils.PARSE_OPTIONS,
+      tracker,
+      syncParams: new RequestParameters({ sub: '' }, {}),
+      token: { exp: Date.now() / 1000 + 10 } as any
+    });
+
+    let sentCheckpoints = 0;
+    for await (const next of stream) {
+      if (typeof next === 'object' && next !== null) {
+        if ('partial_checkpoint_complete' in next) {
+          expect(sentCheckpoints).toBe(1);
+
+          await bucketStorage.startBatch(test_utils.BATCH_OPTIONS, async (batch) => {
+            // Add another high-priority row. This should interrupt the long-running low-priority sync.
+            await batch.save({
+              sourceTable: TEST_TABLE,
+              tag: storage.SaveOperationTag.INSERT,
+              after: {
+                id: 'highprio2',
+                description: 'Another high-priority row'
+              },
+              afterReplicaId: 'highprio2'
+            });
+      
+            await batch.commit('0/2');
+          });
+        }
+        if ('checkpoint' in next || 'checkpoint_diff' in next) {
+          sentCheckpoints += 1;
+        }
+        if ('checkpoint_complete' in next) {
+          break;
+        }
+      }
+    }
+
+    expect(sentCheckpoints).toBe(2);
+  });
+
   test('sync legacy non-raw data', async () => {
     const f = await factory();
 
diff --git a/packages/service-core/src/sync/sync.ts b/packages/service-core/src/sync/sync.ts
@@ -241,7 +241,7 @@ async function* streamResponseInner(
       const abortCheckpointSignal = AbortSignal.any([abortCheckpointController.signal, signal]);
 
       const bucketsByPriority = [...Map.groupBy(bucketsToFetch, (bucket) => bucket.priority).entries()];
-      bucketsByPriority.sort((a, b) => b[0] - a[0]); // Inverting sort order, high priority buckets have smaller priority values
+      bucketsByPriority.sort((a, b) => a[0] - b[0]); // Sort from high to lower priorities
       const lowestPriority = bucketsByPriority.at(-1)?.[0];
 
       function maybeRaceForNewCheckpoint() {
@@ -267,11 +267,11 @@ async function* streamResponseInner(
       // This incrementally updates dataBuckets with each individual bucket position.
       // At the end of this, we can be sure that all buckets have data up to the checkpoint.
       for (const [priority, buckets] of bucketsByPriority) {
+        const isLast = priority === lowestPriority;
         if (abortCheckpointSignal.aborted) {
           break;
         }
 
-        const isLast = priority === lowestPriority;
         yield* bucketDataInBatches({
           storage,
           checkpoint,
@@ -323,7 +323,7 @@ interface BucketDataRequest {
 
 async function* bucketDataInBatches(request: BucketDataRequest) {
   let isDone = false;
-  while (!request.abort_connection.aborted && !isDone) {
+  while (!request.abort_batch.aborted && !isDone) {
     // The code below is functionally the same as this for-await loop below.
     // However, the for-await loop appears to have a memory leak, so we avoid it.
     // for await (const { done, data } of bucketDataBatch(storage, checkpoint, dataBuckets, raw_data, signal)) {
@@ -453,7 +453,7 @@ async function* bucketDataBatch(request: BucketDataRequest): AsyncGenerator<Buck
       // Check if syncing bucket data is supposed to stop before fetching more data
       // from storage.
       if (abort_batch.aborted) {
-        break;
+        return;
       }
     }
 
