Rewrite concurrent-connections clients - now resumeable.

Author: rkistner

pnpm-lock.yaml

@@ -18,6 +18,7 @@
     "@powersync/service-core": "workspace:*",
     "commander": "^12.0.0",
     "jose": "^4.15.1",
+    "undici": "^7.15.0",
     "ws": "^8.18.0",
     "yaml": "^2.5.0"
   },
@@ -26,4 +27,4 @@
     "@types/ws": "~8.2.0",
     "typescript": "^5.7.3"
   }
-}
+}

test-client/package.json

@@ -41,10 +41,19 @@ program
   .option('-n, --num-clients [num-clients]', 'number of clients to connect')
   .option('-m, --mode [mode]', 'http or websocket')
   .option('-p, --print [print]', 'print a field from the data being downloaded')
+  .option('--once', 'stop after the first checkpoint')
   .action(async (options) => {
     const credentials = await getCredentials(options);
 
-    await concurrentConnections(credentials, options['numClients'] ?? 10, options.mode ?? 'http', options.print);
+    await concurrentConnections(
+      {
+        ...credentials,
+        once: options.once ?? false,
+        mode: options.mode
+      },
+      options['numClients'] ?? 10,
+      options.print
+    );
   });
 
 await program.parseAsync();

test-client/src/bin.ts

@@ -29,7 +29,7 @@ export async function getCheckpointData(options: GetCheckpointOptions) {
   let data: types.StreamingSyncData[] = [];
   let checkpoint: types.StreamingSyncCheckpoint;
 
-  for await (let chunk of ndjsonStream<types.StreamingSyncLine>(response.body!)) {
+  for await (let { chunk } of ndjsonStream<types.StreamingSyncLine>(response.body!)) {
     if (isStreamingSyncData(chunk)) {
       // Collect data
       data.push(chunk);

test-client/src/client.ts

@@ -0,0 +1,37 @@
+import type { StreamingSyncLine, StreamingSyncRequest } from '@powersync/service-core';
+import { Readable } from 'node:stream';
+import { request } from 'undici';
+import { ndjsonStream } from './ndjson.js';
+import { StreamEvent, SyncOptions } from './stream.js';
+
+export async function* openHttpStream(options: SyncOptions): AsyncGenerator<StreamEvent> {
+  const streamRequest: StreamingSyncRequest = {
+    raw_data: true,
+    client_id: options.clientId,
+    buckets: [...(options.bucketPositions ?? new Map()).entries()].map(([bucket, after]) => ({
+      name: bucket,
+      after: after
+    }))
+  };
+  const response = await request(options.endpoint + '/sync/stream', {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      Authorization: `Token ${options.token}`
+    },
+    body: JSON.stringify(streamRequest),
+    signal: options.signal
+  });
+
+  if (response.statusCode != 200) {
+    throw new Error(`Request failed with code: ${response.statusCode}\n${await response.body.text()}`);
+  }
+
+  const stream = Readable.toWeb(response.body) as ReadableStream<Uint8Array>;
+
+  for await (let { chunk, size } of ndjsonStream<StreamingSyncLine>(stream)) {
+    yield { stats: { decodedBytes: size } };
+    yield chunk;
+  }
+  // If we reach this, the connection was closed without error by the server
+}

test-client/src/httpStream.ts

@@ -1,12 +1,12 @@
-import { ndjsonStream } from '../ndjson.js';
-
 import { parentPort, workerData } from 'worker_threads';
+import { openStream, SyncOptions } from '../stream.js';
 
 if (parentPort == null) {
   throw new Error(`Can only run this script in a worker_thread`);
 }
 
-const { i, url, token, print } = workerData;
+const { i, print } = workerData;
+const request: SyncOptions = workerData.request;
 
 let size = 0;
 let numOperations = 0;
@@ -25,51 +25,41 @@ const parseChunk = (chunk: any) => {
   });
 };
 
-const response = await fetch(url + '/sync/stream', {
-  method: 'POST',
-  headers: {
-    'Content-Type': 'application/json',
-    Authorization: `Token ${token}`
-  },
-  body: JSON.stringify({
-    raw_data: true,
-    include_checksums: true
-  })
-});
-
-if (!response.ok || response.body == null) {
-  throw new Error(response.statusText + '\n' + (await response.text()));
-}
-
-for await (let chunk of ndjsonStream<any>(response.body)) {
-  size += JSON.stringify(chunk).length;
-  if (chunk?.checkpoint_complete) {
+for await (let chunk of openStream(request)) {
+  if ('error' in chunk) {
+    // Retried automatically
+    console.error(new Date().toISOString(), i, `Error in stream: ${chunk.error}`);
+  } else if ('checkpoint_complete' in chunk) {
     const duration = performance.now() - lastCheckpointStart;
     let message = `checkpoint_complete op_id: ${chunk.checkpoint_complete.last_op_id}, ops: ${numOperations}, bytes: ${size}, duration: ${duration.toFixed(0)}ms`;
     if (print) {
       message += `, data: [${printData}]`;
     }
     printData = [];
     console.log(new Date().toISOString(), i, message);
-  } else if (chunk?.data) {
+
+    if (request.once) {
+      break;
+    }
+  } else if ('data' in chunk) {
     parseChunk(chunk.data);
     numOperations += chunk.data.data.length;
-  } else if (chunk?.checkpoint) {
+  } else if ('checkpoint' in chunk) {
     lastCheckpointStart = performance.now();
     console.log(new Date().toISOString(), i, `checkpoint buckets: ${chunk.checkpoint.buckets.length}`);
-  } else if (chunk?.checkpoint_diff) {
+  } else if ('checkpoint_diff' in chunk) {
     lastCheckpointStart = performance.now();
     console.log(
       new Date().toISOString(),
       i,
       `checkpoint_diff removed_buckets: ${chunk.checkpoint_diff.removed_buckets.length} updated_buckets: ${chunk.checkpoint_diff.updated_buckets.length}`
     );
+  } else if ('stats' in chunk) {
+    size += chunk.stats.decodedBytes;
   } else {
     const key = Object.keys(chunk)[0];
     if (key != 'token_expires_in' && key != 'data') {
       console.log(new Date().toISOString(), i, key);
     }
   }
 }
-
-parentPort.postMessage({ done: true });

test-client/src/load-testing/http-worker.ts renamed to test-client/src/load-testing/load-test-worker.ts

@@ -1,43 +1,24 @@
 import { Worker } from 'worker_threads';
-import { Credentials } from '../auth.js';
+import { SyncOptions } from '../stream.js';
 
 export type Mode = 'http' | 'websocket';
 
-export async function stream(i: number, credentials: Credentials, mode: Mode, print: string | undefined) {
-  const worker =
-    mode == 'websocket'
-      ? new Worker(new URL('./rsocket-worker.js', import.meta.url), {
-          workerData: { i, token: credentials.token, url: credentials.endpoint.replace(/^http/, 'ws'), print }
-        })
-      : new Worker(new URL('./http-worker.js', import.meta.url), {
-          workerData: { i, token: credentials.token, url: credentials.endpoint, print }
-        });
+export async function stream(i: number, request: SyncOptions, print: string | undefined) {
+  const worker = new Worker(new URL('./load-test-worker.js', import.meta.url), {
+    workerData: { i, request, print }
+  });
   await new Promise((resolve, reject) => {
     worker.on('message', (event) => resolve(event));
     worker.on('error', (err) => reject(err));
+    worker.on('exit', (__code) => {
+      resolve(null);
+    });
   });
   worker.terminate();
 }
 
-export async function streamForever(i: number, credentials: Credentials, mode: Mode, print: string | undefined) {
-  while (true) {
-    try {
-      await stream(i, credentials, mode, print);
-      console.log(new Date().toISOString(), i, 'Stream ended');
-    } catch (e: any) {
-      console.error(new Date().toISOString(), i, e.message);
-      await new Promise((resolve) => setTimeout(resolve, 1000 + Math.random()));
-    }
-  }
-}
-
-export async function concurrentConnections(
-  credentials: Credentials,
-  numClients: number,
-  mode: Mode,
-  print: string | undefined
-) {
+export async function concurrentConnections(options: SyncOptions, numClients: number, print: string | undefined) {
   for (let i = 0; i < numClients; i++) {
-    streamForever(i, credentials, mode, print);
+    stream(i, { ...options, clientId: options.clientId ?? `load-test-${i}` }, print);
   }
 }