Clone workers in streams

Author: neekolas

chaos/streams.ts

@@ -1,23 +1,38 @@
 import type { ChaosProvider } from "@chaos/provider";
-import { typeofStream } from "@workers/main";
+import { typeofStream, type WorkerClient } from "@workers/main";
 import type { WorkerManager } from "@workers/manager";
 
+export type StreamsConfig = {
+  cloned: boolean;
+};
+
 export class StreamsChaos implements ChaosProvider {
-  workers?: WorkerManager;
-  start(workers: WorkerManager) {
+  workers?: WorkerClient[];
+  config: StreamsConfig;
+
+  constructor(config: StreamsConfig) {
+    this.config = config;
+  }
+
+  async start(workers: WorkerManager) {
     console.log("Starting StreamsChaos");
-    this.workers = workers;
-    for (const worker of workers.getAll()) {
-      worker.worker.startStream(typeofStream.Message);
+    let allWorkers = workers.getAll().map((w) => w.worker);
+    if (this.config.cloned) {
+      allWorkers = await Promise.all(allWorkers.map((w) => w.clone()));
+    }
+
+    this.workers = allWorkers;
+    for (const worker of allWorkers) {
+      worker.startStream(typeofStream.Message);
     }
 
     return Promise.resolve();
   }
 
   stop() {
     if (this.workers) {
-      for (const worker of this.workers.getAll()) {
-        worker.worker.stopStreams();
+      for (const worker of this.workers) {
+        worker.stopStreams();
       }
     }
 

forks/cli.ts

@@ -50,7 +50,11 @@ function buildRuntimeConfig(options: ForkOptions): RuntimeConfig {
     network: (options.env || "dev") as "local" | "dev" | "production",
     networkChaos: resolveNetworkChaosConfig(options.networkChaosLevel),
     dbChaos: resolveDbChaosConfig(options.dbChaosLevel),
-    backgroundStreams: options.withBackgroundStreams,
+    backgroundStreams: options.withBackgroundStreams
+      ? {
+          cloned: true,
+        }
+      : null,
   };
 }
 

forks/config.ts

@@ -1,5 +1,6 @@
 import { type DbChaosConfig } from "@chaos/db";
 import type { NetworkChaosConfig } from "@chaos/network";
+import type { StreamsConfig } from "@chaos/streams";
 import { getActiveVersion, type XmtpEnv } from "@helpers/versions";
 
 export const NODE_VERSION = getActiveVersion().nodeBindings; // default to latest version, can be overridden with --nodeBindings=3.1.1
@@ -130,7 +131,7 @@ export type RuntimeConfig = {
   network: XmtpEnv; // XMTP network
   networkChaos: NetworkChaosConfig | null; // Network chaos configuration
   dbChaos: DbChaosConfig | null; // Database chaos configuration
-  backgroundStreams: boolean; //
+  backgroundStreams: StreamsConfig | null; //
 };
 
 export function getConfigFromEnv(): RuntimeConfig {
@@ -161,7 +162,9 @@ export function printConfig(config: RuntimeConfig): void {
   console.info(`randomInboxIdsCount: ${randomInboxIdsCount}`);
   console.info(`installationCount: ${installationCount}`);
   console.info(`testName: ${testName}`);
-  console.info(`backgroundStreams: ${config.backgroundStreams}`);
+  console.info(
+    `backgroundStreams: ${config.backgroundStreams ? "enabled" : "disabled"}. With clones: ${config.backgroundStreams?.cloned}`,
+  );
 
   if (config.networkChaos) {
     console.info("\nNETWORK CHAOS PARAMETERS");

forks/forks.test.ts

@@ -32,8 +32,13 @@ const {
 } = getConfigFromEnv();
 
 const createOperations = (worker: Worker, groupID: string) => {
-  const getGroup = () =>
-    worker.client.conversations.getConversationById(groupID) as Promise<Group>;
+  const getGroup = () => {
+    const group = worker.client.conversations.getConversationById(groupID);
+    if (!group) {
+      throw new Error(`Group ${groupID} not found`);
+    }
+    return group as Promise<Group>;
+  };
 
   return {
     updateName: () =>
@@ -89,7 +94,7 @@ const startChaos = async (workers: WorkerManager): Promise<ChaosProvider[]> => {
   }
 
   if (backgroundStreams) {
-    chaosProviders.push(new StreamsChaos());
+    chaosProviders.push(new StreamsChaos(backgroundStreams));
   }
 
   // Start all chaos providers
@@ -192,7 +197,15 @@ describe(testName, () => {
                 })(),
             );
             try {
-              await Promise.all(parallelOperationsArray);
+              const results = await Promise.allSettled(parallelOperationsArray);
+              for (const result of results) {
+                if (result.status === "rejected") {
+                  console.error(
+                    `Group ${groupIndex + 1} operation failed:`,
+                    result.reason,
+                  );
+                }
+              }
             } catch (e) {
               console.error(`Group ${groupIndex + 1} operation failed:`, e);
             }

workers/main.ts

@@ -113,6 +113,9 @@ interface IWorkerClient {
 
   // Properties
   readonly currentFolder: string;
+
+  // Clone Management
+  clone(): Promise<WorkerClient>;
 }
 
 // Worker thread code as a string
@@ -1254,4 +1257,42 @@ export class WorkerClient extends Worker implements IWorkerClient {
       installationId: newInstallationId,
     };
   }
+
+  /**
+   * Creates a clone of this worker client with a separate underlying client instance
+   * The clone will have the same configuration but a new name: ${original_worker_name}_clone
+   * @returns A new WorkerClient instance with separate client
+   */
+  async clone(): Promise<WorkerClient> {
+    console.debug(`[${this.nameId}] Creating clone of worker`);
+
+    // Create the clone name
+    const cloneName = `${this.name}_clone`;
+
+    // Create a WorkerBase object with the same properties but new name
+    const cloneWorkerBase: WorkerBase = {
+      name: cloneName,
+      sdk: this.sdk,
+      folder: this.folder,
+      walletKey: this.walletKey,
+      encryptionKey: this.encryptionKeyHex,
+    };
+
+    // Create a new WorkerClient instance with the same configuration
+    const clonedWorker = new WorkerClient(
+      cloneWorkerBase,
+      this.env,
+      {}, // Use default worker options
+      this.apiUrl,
+    );
+
+    // Initialize the cloned worker to create its client instance
+    await clonedWorker.initialize();
+
+    console.debug(
+      `[${this.nameId}] Successfully created clone: ${clonedWorker.nameId}`,
+    );
+
+    return clonedWorker;
+  }
 }