Add initializeWorker in plugin interface.

Author: donald-pinckney

temporal-sdk/src/main/java/io/temporal/common/plugin/SimplePluginBuilder.java

@@ -26,12 +26,14 @@
 import io.temporal.common.interceptors.WorkerInterceptor;
 import io.temporal.common.interceptors.WorkflowClientInterceptor;
 import io.temporal.serviceclient.WorkflowServiceStubsOptions;
+import io.temporal.worker.Worker;
 import io.temporal.worker.WorkerFactoryOptions;
 import io.temporal.worker.WorkerOptions;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import javax.annotation.Nonnull;
 
@@ -69,6 +71,7 @@ public final class SimplePluginBuilder {
   private final List<Consumer<WorkflowClientOptions.Builder>> clientCustomizers = new ArrayList<>();
   private final List<Consumer<WorkerFactoryOptions.Builder>> factoryCustomizers = new ArrayList<>();
   private final List<Consumer<WorkerOptions.Builder>> workerCustomizers = new ArrayList<>();
+  private final List<BiConsumer<String, Worker>> workerInitializers = new ArrayList<>();
   private final List<WorkerInterceptor> workerInterceptors = new ArrayList<>();
   private final List<WorkflowClientInterceptor> clientInterceptors = new ArrayList<>();
   private final List<ContextPropagator> contextPropagators = new ArrayList<>();
@@ -139,6 +142,29 @@ public SimplePluginBuilder customizeWorker(@Nonnull Consumer<WorkerOptions.Build
     return this;
   }
 
+  /**
+   * Adds an initializer that is called after a worker is created. This can be used to register
+   * workflows, activities, and Nexus services on the worker.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * SimplePluginBuilder.newBuilder("my-plugin")
+   *     .initializeWorker((taskQueue, worker) -> {
+   *         worker.registerWorkflowImplementationTypes(MyWorkflow.class);
+   *         worker.registerActivitiesImplementations(new MyActivityImpl());
+   *     })
+   *     .build();
+   * }</pre>
+   *
+   * @param initializer a consumer that receives the task queue name and worker
+   * @return this builder for chaining
+   */
+  public SimplePluginBuilder initializeWorker(@Nonnull BiConsumer<String, Worker> initializer) {
+    workerInitializers.add(Objects.requireNonNull(initializer));
+    return this;
+  }
+
   /**
    * Adds worker interceptors. Interceptors are appended to any existing interceptors in the
    * configuration.
@@ -188,6 +214,7 @@ public PluginBase build() {
         new ArrayList<>(clientCustomizers),
         new ArrayList<>(factoryCustomizers),
         new ArrayList<>(workerCustomizers),
+        new ArrayList<>(workerInitializers),
         new ArrayList<>(workerInterceptors),
         new ArrayList<>(clientInterceptors),
         new ArrayList<>(contextPropagators));
@@ -199,6 +226,7 @@ private static final class SimplePlugin extends PluginBase {
     private final List<Consumer<WorkflowClientOptions.Builder>> clientCustomizers;
     private final List<Consumer<WorkerFactoryOptions.Builder>> factoryCustomizers;
     private final List<Consumer<WorkerOptions.Builder>> workerCustomizers;
+    private final List<BiConsumer<String, Worker>> workerInitializers;
     private final List<WorkerInterceptor> workerInterceptors;
     private final List<WorkflowClientInterceptor> clientInterceptors;
     private final List<ContextPropagator> contextPropagators;
@@ -209,6 +237,7 @@ private static final class SimplePlugin extends PluginBase {
         List<Consumer<WorkflowClientOptions.Builder>> clientCustomizers,
         List<Consumer<WorkerFactoryOptions.Builder>> factoryCustomizers,
         List<Consumer<WorkerOptions.Builder>> workerCustomizers,
+        List<BiConsumer<String, Worker>> workerInitializers,
         List<WorkerInterceptor> workerInterceptors,
         List<WorkflowClientInterceptor> clientInterceptors,
         List<ContextPropagator> contextPropagators) {
@@ -217,6 +246,7 @@ private static final class SimplePlugin extends PluginBase {
       this.clientCustomizers = clientCustomizers;
       this.factoryCustomizers = factoryCustomizers;
       this.workerCustomizers = workerCustomizers;
+      this.workerInitializers = workerInitializers;
       this.workerInterceptors = workerInterceptors;
       this.clientInterceptors = clientInterceptors;
       this.contextPropagators = contextPropagators;
@@ -292,5 +322,12 @@ public WorkerOptions.Builder configureWorker(
       }
       return builder;
     }
+
+    @Override
+    public void initializeWorker(@Nonnull String taskQueue, @Nonnull Worker worker) {
+      for (BiConsumer<String, Worker> initializer : workerInitializers) {
+        initializer.accept(taskQueue, worker);
+      }
+    }
   }
 }

temporal-sdk/src/main/java/io/temporal/common/plugin/WorkerPlugin.java

@@ -21,6 +21,7 @@
 package io.temporal.common.plugin;
 
 import io.temporal.common.Experimental;
+import io.temporal.worker.Worker;
 import io.temporal.worker.WorkerFactory;
 import io.temporal.worker.WorkerFactoryOptions;
 import io.temporal.worker.WorkerOptions;
@@ -103,6 +104,30 @@ default WorkerOptions.Builder configureWorker(
     return builder;
   }
 
+  /**
+   * Called after a worker is created, allowing plugins to register workflows, activities, Nexus
+   * services, and other components on the worker.
+   *
+   * <p>This method is called in forward (registration) order immediately after the worker is
+   * created in {@link WorkerFactory#newWorker}.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * @Override
+   * public void initializeWorker(String taskQueue, Worker worker) {
+   *     worker.registerWorkflowImplementationTypes(MyWorkflow.class);
+   *     worker.registerActivitiesImplementations(new MyActivityImpl());
+   * }
+   * }</pre>
+   *
+   * @param taskQueue the task queue name for the worker
+   * @param worker the newly created worker
+   */
+  default void initializeWorker(@Nonnull String taskQueue, @Nonnull Worker worker) {
+    // Default: no-op
+  }
+
   /**
    * Allows the plugin to wrap worker factory startup. Called during execution phase in reverse
    * order (first plugin wraps all others).

temporal-sdk/src/main/java/io/temporal/worker/WorkerFactory.java

@@ -169,6 +169,14 @@ public synchronized Worker newWorker(String taskQueue, WorkerOptions options) {
               workflowThreadExecutor,
               workflowClient.getOptions().getContextPropagators());
       workers.put(taskQueue, worker);
+
+      // Go through the plugins to call plugin initializeWorker hooks (e.g. register workflows, activities, etc.)
+      for (Object plugin : plugins) {
+        if (plugin instanceof WorkerPlugin) {
+          ((WorkerPlugin) plugin).initializeWorker(taskQueue, worker);
+        }
+      }
+
       return worker;
     } else {
       log.warn(

temporal-sdk/src/test/java/io/temporal/common/plugin/PluginTest.java

@@ -97,6 +97,9 @@ public String getName() {
     final boolean[] called = {false};
     plugin.runWorkerFactory(null, () -> called[0] = true);
     assertTrue("runWorkerFactory should call next", called[0]);
+
+    // Test default initializeWorker is a no-op (doesn't throw)
+    plugin.initializeWorker("test-queue", null);
   }
 
   @Test

temporal-sdk/src/test/java/io/temporal/common/plugin/SimplePluginBuilderTest.java

@@ -192,6 +192,48 @@ public void testInterceptorsAppendToExisting() {
     assertSame(newInterceptor, interceptors[1]);
   }
 
+  @Test
+  public void testInitializeWorker() {
+    AtomicBoolean initialized = new AtomicBoolean(false);
+    String[] capturedTaskQueue = {null};
+
+    PluginBase plugin =
+        SimplePluginBuilder.newBuilder("test")
+            .initializeWorker(
+                (taskQueue, worker) -> {
+                  initialized.set(true);
+                  capturedTaskQueue[0] = taskQueue;
+                })
+            .build();
+
+    // Call initializeWorker with null worker (we're just testing the callback is invoked)
+    ((WorkerPlugin) plugin).initializeWorker("my-task-queue", null);
+
+    assertTrue("Initializer should have been called", initialized.get());
+    assertEquals("my-task-queue", capturedTaskQueue[0]);
+  }
+
+  @Test
+  public void testMultipleWorkerInitializers() {
+    AtomicInteger callCount = new AtomicInteger(0);
+
+    PluginBase plugin =
+        SimplePluginBuilder.newBuilder("test")
+            .initializeWorker((taskQueue, worker) -> callCount.incrementAndGet())
+            .initializeWorker((taskQueue, worker) -> callCount.incrementAndGet())
+            .initializeWorker((taskQueue, worker) -> callCount.incrementAndGet())
+            .build();
+
+    ((WorkerPlugin) plugin).initializeWorker("test-queue", null);
+
+    assertEquals("All initializers should be called", 3, callCount.get());
+  }
+
+  @Test(expected = NullPointerException.class)
+  public void testNullInitializeWorker() {
+    SimplePluginBuilder.newBuilder("test").initializeWorker(null);
+  }
+
   @Test(expected = NullPointerException.class)
   public void testNullName() {
     SimplePluginBuilder.newBuilder(null);