feat(worker): Add support for concurrent pollers and more (#1132)

Author: mjameswh

package-lock.json

@@ -294,6 +294,10 @@ impl ObjectHandleConversionsExt for Handle<'_, JsObject> {
             js_value_getter!(cx, self, "maxConcurrentLocalActivityExecutions", JsNumber) as usize;
         let max_outstanding_workflow_tasks =
             js_value_getter!(cx, self, "maxConcurrentWorkflowTaskExecutions", JsNumber) as usize;
+        let max_concurrent_wft_polls =
+            js_value_getter!(cx, self, "maxConcurrentWorkflowTaskPolls", JsNumber) as usize;
+        let max_concurrent_at_polls =
+            js_value_getter!(cx, self, "maxConcurrentActivityTaskPolls", JsNumber) as usize;
         let sticky_queue_schedule_to_start_timeout = Duration::from_millis(js_value_getter!(
             cx,
             self,
@@ -335,6 +339,8 @@ impl ObjectHandleConversionsExt for Handle<'_, JsObject> {
             .max_outstanding_workflow_tasks(max_outstanding_workflow_tasks)
             .max_outstanding_activities(max_outstanding_activities)
             .max_outstanding_local_activities(max_outstanding_local_activities)
+            .max_concurrent_wft_polls(max_concurrent_wft_polls)
+            .max_concurrent_at_polls(max_concurrent_at_polls)
             .max_cached_workflows(max_cached_workflows)
             .sticky_queue_schedule_to_start_timeout(sticky_queue_schedule_to_start_timeout)
             .graceful_shutdown_period(graceful_shutdown_period)

packages/core-bridge/sdk-core

@@ -120,8 +120,9 @@ pub fn start_bridge_loop(
 ) {
     let mut tokio_builder = tokio::runtime::Builder::new_multi_thread();
     tokio_builder.enable_all().thread_name("core");
-    let core_runtime =
-        Arc::new(CoreRuntime::new(telemetry_options, tokio_builder).expect("Failed to create CoreRuntime"));
+    let core_runtime = Arc::new(
+        CoreRuntime::new(telemetry_options, tokio_builder).expect("Failed to create CoreRuntime"),
+    );
 
     core_runtime.tokio_handle().block_on(async {
         loop {

packages/core-bridge/src/conversions.rs

@@ -277,6 +277,16 @@ export interface WorkerOptions {
   maxConcurrentWorkflowTaskExecutions: number;
   maxConcurrentLocalActivityExecutions: number;
 
+  /**
+   * Maximum number of Workflow tasks to poll concurrently.
+   */
+  maxConcurrentWorkflowTaskPolls: number;
+
+  /**
+   * Maximum number of Activity tasks to poll concurrently.
+   */
+  maxConcurrentActivityTaskPolls: number;
+
   /**
    * If set to `false` this worker will only handle workflow tasks and local activities, it will not
    * poll for activity tasks.

packages/core-bridge/src/runtime.rs

@@ -11,11 +11,15 @@ export type Spec = Record<string, () => any>;
 export interface SetupArgSpec extends Spec {
   '--ns': typeof String;
   '--server-address': typeof String;
+  '--client-cert-path': typeof String;
+  '--client-key-path': typeof String;
 }
 
 export const setupArgSpec: SetupArgSpec = {
   '--ns': String,
   '--server-address': String,
+  '--client-cert-path': String,
+  '--client-key-path': String,
 };
 
 export interface StarterArgSpec extends Spec {
@@ -33,6 +37,8 @@ export interface StarterArgSpec extends Spec {
   '--do-query': typeof String;
   '--initial-query-delay-ms': typeof Number;
   '--query-interval-ms': typeof Number;
+  '--client-cert-path': typeof String;
+  '--client-key-path': typeof String;
 }
 
 export const starterArgSpec: StarterArgSpec = {
@@ -50,6 +56,8 @@ export const starterArgSpec: StarterArgSpec = {
   '--do-query': String,
   '--initial-query-delay-ms': Number,
   '--query-interval-ms': Number,
+  '--client-cert-path': String,
+  '--client-key-path': String,
 };
 
 export interface WorkerArgSpec extends Spec {
@@ -59,12 +67,17 @@ export interface WorkerArgSpec extends Spec {
   '--max-concurrent-at-executions': typeof Number;
   '--max-concurrent-wft-executions': typeof Number;
   '--max-concurrent-la-executions': typeof Number;
+  '--max-wft-pollers': typeof Number;
+  '--max-at-pollers': typeof Number;
+  '--wf-thread-pool-size': typeof Number;
   '--log-level': typeof String;
   '--log-file': typeof String;
   '--server-address': typeof String;
   '--otel-url': typeof String;
   '--status-port': typeof Number;
   '--shutdown-grace-time-ms': typeof String;
+  '--client-cert-path': typeof String;
+  '--client-key-path': typeof String;
 }
 
 export const workerArgSpec: WorkerArgSpec = {
@@ -74,13 +87,17 @@ export const workerArgSpec: WorkerArgSpec = {
   '--max-concurrent-at-executions': Number,
   '--max-concurrent-wft-executions': Number,
   '--max-concurrent-la-executions': Number,
-  '--isolate-pool-size': Number,
+  '--max-wft-pollers': Number,
+  '--max-at-pollers': Number,
+  '--wf-thread-pool-size': Number,
   '--log-level': String,
   '--log-file': String,
   '--server-address': String,
   '--otel-url': String,
   '--status-port': Number,
   '--shutdown-grace-time-ms': String,
+  '--client-cert-path': String,
+  '--client-key-path': String,
 };
 
 export interface WrapperArgSpec extends Spec {

packages/core-bridge/ts/index.ts

@@ -1,3 +1,4 @@
+import { readFileSync } from 'fs';
 import arg from 'arg';
 import { Connection } from '@temporalio/client';
 import { createNamespace, waitOnNamespace } from '@temporalio/testing/lib/utils';
@@ -8,8 +9,31 @@ async function main() {
 
   const serverAddress = getRequired(args, '--server-address');
   const namespace = getRequired(args, '--ns');
+  const clientCertPath = args['--client-cert-path'];
+  const clientKeyPath = args['--client-key-path'];
 
-  const connection = await Connection.connect({ address: serverAddress });
+  const tlsConfig =
+    clientCertPath && clientKeyPath
+      ? {
+          tls: {
+            clientCertPair: {
+              crt: readFileSync(clientCertPath),
+              key: readFileSync(clientKeyPath),
+            },
+          },
+        }
+      : {};
+
+  const connection = await Connection.connect({ address: serverAddress, ...tlsConfig });
+
+  console.log('Checking if namespace already exists', { namespace });
+  try {
+    await waitOnNamespace(connection, namespace, 1, 0);
+    // Namespace already exists. Nothing to do
+    return;
+  } catch (e) {
+    // Ignore error. Will create namespace if it does not exist
+  }
 
   await createNamespace(connection, namespace);
   console.log('Registered namespace', { namespace });

packages/test/src/load/args.ts

@@ -1,5 +1,5 @@
 import os from 'node:os';
-import fs from 'node:fs';
+import fs, { readFileSync } from 'node:fs';
 import arg from 'arg';
 import pidusage from 'pidusage';
 import * as grpc from '@grpc/grpc-js';
@@ -201,7 +201,22 @@ async function main() {
   const workerMemoryLogFile = args['--worker-memory-log-file'];
   const workerCPULogFile = args['--worker-cpu-log-file'];
 
-  const connection = await Connection.connect({ address: serverAddress });
+  const clientCertPath = args['--client-cert-path'];
+  const clientKeyPath = args['--client-key-path'];
+
+  const tlsConfig =
+    clientCertPath && clientKeyPath
+      ? {
+          tls: {
+            clientCertPair: {
+              crt: readFileSync(clientCertPath),
+              key: readFileSync(clientKeyPath),
+            },
+          },
+        }
+      : {};
+
+  const connection = await Connection.connect({ address: serverAddress, ...tlsConfig });
   const client = new WorkflowClient({ connection, namespace });
   const stopCondition = runForSeconds ? new UntilSecondsElapsed(runForSeconds) : new NumberOfWorkflows(iterations);
   const queryingOptions = queryName ? { queryName, queryIntervalMs, initialQueryDelayMs } : undefined;

packages/test/src/load/setup.ts

@@ -1,4 +1,4 @@
-import fs from 'node:fs';
+import fs, { readFileSync } from 'node:fs';
 import http from 'node:http';
 import { inspect } from 'node:util';
 import { OTLPTraceExporter } from '@opentelemetry/exporter-trace-otlp-grpc';
@@ -107,10 +107,15 @@ async function main() {
   const maxConcurrentWorkflowTaskExecutions = args['--max-concurrent-wft-executions'] ?? 100;
   const maxConcurrentLocalActivityExecutions = args['--max-concurrent-la-executions'] ?? 100;
   const maxCachedWorkflows = args['--max-cached-wfs'];
+  const maxConcurrentWorkflowTaskPolls = args['--max-wft-pollers'] ?? 2;
+  const maxConcurrentActivityTaskPolls = args['--max-at-pollers'] ?? 2;
+  const workflowThreadPoolSize: number | undefined = args['--wf-thread-pool-size'];
   const oTelUrl = args['--otel-url'];
   const logLevel = (args['--log-level'] || 'INFO').toUpperCase();
   const logFile = args['--log-file'];
   const serverAddress = getRequired(args, '--server-address');
+  const clientCertPath = args['--client-cert-path'];
+  const clientKeyPath = args['--client-key-path'];
   const namespace = getRequired(args, '--ns');
   const taskQueue = getRequired(args, '--task-queue');
   const statusPort = args['--status-port'];
@@ -139,12 +144,26 @@ async function main() {
     logger,
   });
 
+  const tlsConfig =
+    clientCertPath && clientKeyPath
+      ? {
+          tls: {
+            clientCertPair: {
+              crt: readFileSync(clientCertPath),
+              key: readFileSync(clientKeyPath),
+            },
+          },
+        }
+      : {};
+
   const clientConnection = await Connection.connect({
     address: serverAddress,
+    ...tlsConfig,
   });
 
   const connection = await NativeConnection.connect({
     address: serverAddress,
+    ...tlsConfig,
   });
 
   await withOptionalOtel(args, async () => {
@@ -159,6 +178,9 @@ async function main() {
       maxConcurrentWorkflowTaskExecutions,
       maxCachedWorkflows,
       shutdownGraceTime,
+      workflowThreadPoolSize,
+      maxConcurrentActivityTaskPolls,
+      maxConcurrentWorkflowTaskPolls,
       interceptors: {
         activityInbound: [() => new ConnectionInjectorInterceptor(clientConnection)],
       },

packages/test/src/load/starter.ts

@@ -162,8 +162,8 @@ if (RUN_INTEGRATION_TESTS) {
       ...defaultOptions,
       taskQueue,
       sinks,
-      maxCachedWorkflows: 1,
-      maxConcurrentWorkflowTaskExecutions: 1,
+      maxCachedWorkflows: 0,
+      maxConcurrentWorkflowTaskExecutions: 2,
     });
     const client = new WorkflowClient();
     await Promise.all([
@@ -179,7 +179,7 @@ if (RUN_INTEGRATION_TESTS) {
 
     t.deepEqual(recordedMessages, [
       'Workflow execution started, replaying: false, hl: 3',
-      'Workflow execution completed, replaying: false, hl: 12',
+      'Workflow execution completed, replaying: false, hl: 8',
     ]);
   });
 
@@ -201,8 +201,8 @@ if (RUN_INTEGRATION_TESTS) {
       ...defaultOptions,
       taskQueue,
       sinks,
-      maxCachedWorkflows: 1,
-      maxConcurrentWorkflowTaskExecutions: 1,
+      maxCachedWorkflows: 0,
+      maxConcurrentWorkflowTaskExecutions: 2,
     });
     const client = new WorkflowClient();
     await worker.runUntil(client.execute(workflows.logSinkTester, { taskQueue, workflowId: uuid4() }));
@@ -212,6 +212,6 @@ if (RUN_INTEGRATION_TESTS) {
       'Workflow execution started, replaying: false, hl: 3',
       'Workflow execution started, replaying: true, hl: 3',
     ]);
-    t.is(recordedMessages[recordedMessages.length - 1], 'Workflow execution completed, replaying: false, hl: 12');
+    t.is(recordedMessages[recordedMessages.length - 1], 'Workflow execution completed, replaying: false, hl: 8');
   });
 }