add startWorker and shutdownWorker

Author: donald-pinckney

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

@@ -71,6 +71,8 @@ public final class SimplePluginBuilder {
   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<BiConsumer<String, Worker>> workerStartCallbacks = new ArrayList<>();
+  private final List<BiConsumer<String, Worker>> workerShutdownCallbacks = new ArrayList<>();
   private final List<WorkerInterceptor> workerInterceptors = new ArrayList<>();
   private final List<WorkflowClientInterceptor> clientInterceptors = new ArrayList<>();
   private final List<ContextPropagator> contextPropagators = new ArrayList<>();
@@ -164,6 +166,59 @@ public SimplePluginBuilder initializeWorker(@Nonnull BiConsumer<String, Worker>
     return this;
   }
 
+  /**
+   * Adds a callback that is invoked when a worker starts. This can be used to start per-worker
+   * resources or record metrics.
+   *
+   * <p>Note: For registering workflows and activities, use {@link #initializeWorker} instead, as
+   * registrations must happen before the worker starts polling.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * SimplePluginBuilder.newBuilder("my-plugin")
+   *     .onWorkerStart((taskQueue, worker) -> {
+   *         logger.info("Worker started for task queue: {}", taskQueue);
+   *         perWorkerResources.put(taskQueue, new ResourcePool());
+   *     })
+   *     .build();
+   * }</pre>
+   *
+   * @param callback a consumer that receives the task queue name and worker when the worker starts
+   * @return this builder for chaining
+   */
+  public SimplePluginBuilder onWorkerStart(@Nonnull BiConsumer<String, Worker> callback) {
+    workerStartCallbacks.add(Objects.requireNonNull(callback));
+    return this;
+  }
+
+  /**
+   * Adds a callback that is invoked when a worker shuts down. This can be used to clean up
+   * per-worker resources initialized in {@link #initializeWorker} or {@link #onWorkerStart}.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * SimplePluginBuilder.newBuilder("my-plugin")
+   *     .onWorkerShutdown((taskQueue, worker) -> {
+   *         logger.info("Worker shutting down for task queue: {}", taskQueue);
+   *         ResourcePool pool = perWorkerResources.remove(taskQueue);
+   *         if (pool != null) {
+   *             pool.close();
+   *         }
+   *     })
+   *     .build();
+   * }</pre>
+   *
+   * @param callback a consumer that receives the task queue name and worker when the worker shuts
+   *     down
+   * @return this builder for chaining
+   */
+  public SimplePluginBuilder onWorkerShutdown(@Nonnull BiConsumer<String, Worker> callback) {
+    workerShutdownCallbacks.add(Objects.requireNonNull(callback));
+    return this;
+  }
+
   /**
    * Adds worker interceptors. Interceptors are appended to any existing interceptors in the
    * configuration.
@@ -214,6 +269,8 @@ public PluginBase build() {
         new ArrayList<>(factoryCustomizers),
         new ArrayList<>(workerCustomizers),
         new ArrayList<>(workerInitializers),
+        new ArrayList<>(workerStartCallbacks),
+        new ArrayList<>(workerShutdownCallbacks),
         new ArrayList<>(workerInterceptors),
         new ArrayList<>(clientInterceptors),
         new ArrayList<>(contextPropagators));
@@ -226,6 +283,8 @@ private static final class SimplePlugin extends PluginBase {
     private final List<Consumer<WorkerFactoryOptions.Builder>> factoryCustomizers;
     private final List<Consumer<WorkerOptions.Builder>> workerCustomizers;
     private final List<BiConsumer<String, Worker>> workerInitializers;
+    private final List<BiConsumer<String, Worker>> workerStartCallbacks;
+    private final List<BiConsumer<String, Worker>> workerShutdownCallbacks;
     private final List<WorkerInterceptor> workerInterceptors;
     private final List<WorkflowClientInterceptor> clientInterceptors;
     private final List<ContextPropagator> contextPropagators;
@@ -237,6 +296,8 @@ private static final class SimplePlugin extends PluginBase {
         List<Consumer<WorkerFactoryOptions.Builder>> factoryCustomizers,
         List<Consumer<WorkerOptions.Builder>> workerCustomizers,
         List<BiConsumer<String, Worker>> workerInitializers,
+        List<BiConsumer<String, Worker>> workerStartCallbacks,
+        List<BiConsumer<String, Worker>> workerShutdownCallbacks,
         List<WorkerInterceptor> workerInterceptors,
         List<WorkflowClientInterceptor> clientInterceptors,
         List<ContextPropagator> contextPropagators) {
@@ -246,6 +307,8 @@ private static final class SimplePlugin extends PluginBase {
       this.factoryCustomizers = factoryCustomizers;
       this.workerCustomizers = workerCustomizers;
       this.workerInitializers = workerInitializers;
+      this.workerStartCallbacks = workerStartCallbacks;
+      this.workerShutdownCallbacks = workerShutdownCallbacks;
       this.workerInterceptors = workerInterceptors;
       this.clientInterceptors = clientInterceptors;
       this.contextPropagators = contextPropagators;
@@ -328,5 +391,24 @@ public void initializeWorker(@Nonnull String taskQueue, @Nonnull Worker worker)
         initializer.accept(taskQueue, worker);
       }
     }
+
+    @Override
+    public void startWorker(
+        @Nonnull String taskQueue, @Nonnull Worker worker, @Nonnull Runnable next)
+        throws Exception {
+      next.run();
+      for (BiConsumer<String, Worker> callback : workerStartCallbacks) {
+        callback.accept(taskQueue, worker);
+      }
+    }
+
+    @Override
+    public void shutdownWorker(
+        @Nonnull String taskQueue, @Nonnull Worker worker, @Nonnull Runnable next) {
+      for (BiConsumer<String, Worker> callback : workerShutdownCallbacks) {
+        callback.accept(taskQueue, worker);
+      }
+      next.run();
+    }
   }
 }

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

@@ -106,7 +106,8 @@ default WorkerOptions.Builder configureWorker(
    * 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}.
+   * created in {@link WorkerFactory#newWorker}. This is the appropriate place for registrations
+   * because it is called before the worker starts polling.
    *
    * <p>Example:
    *
@@ -125,6 +126,68 @@ default void initializeWorker(@Nonnull String taskQueue, @Nonnull Worker worker)
     // Default: no-op
   }
 
+  /**
+   * Allows the plugin to wrap individual worker startup. Called during execution phase in reverse
+   * order (first plugin wraps all others) when {@link WorkerFactory#start()} is invoked.
+   *
+   * <p>This method is called for each worker when the factory starts. Use this for per-worker
+   * resource initialization, logging, or metrics. Note that workflow/activity registration should
+   * be done in {@link #initializeWorker} instead, as this method is called after registrations are
+   * finalized.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * @Override
+   * public void startWorker(String taskQueue, Worker worker, Runnable next) throws Exception {
+   *     logger.info("Starting worker for task queue: {}", taskQueue);
+   *     perWorkerResources.put(taskQueue, new ResourcePool());
+   *     next.run();
+   * }
+   * }</pre>
+   *
+   * @param taskQueue the task queue name for the worker
+   * @param worker the worker being started
+   * @param next runnable that starts the next in chain (eventually starts the actual worker)
+   * @throws Exception if startup fails
+   */
+  default void startWorker(
+      @Nonnull String taskQueue, @Nonnull Worker worker, @Nonnull Runnable next) throws Exception {
+    next.run();
+  }
+
+  /**
+   * Allows the plugin to wrap individual worker shutdown. Called during shutdown phase in reverse
+   * order (first plugin wraps all others) when {@link WorkerFactory#shutdown()} or {@link
+   * WorkerFactory#shutdownNow()} is invoked.
+   *
+   * <p>This method is called for each worker when the factory shuts down. Use this for per-worker
+   * resource cleanup that was initialized in {@link #startWorker} or {@link #initializeWorker}.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * @Override
+   * public void shutdownWorker(String taskQueue, Worker worker, Runnable next) {
+   *     logger.info("Shutting down worker for task queue: {}", taskQueue);
+   *     next.run();
+   *     ResourcePool pool = perWorkerResources.remove(taskQueue);
+   *     if (pool != null) {
+   *         pool.close();
+   *     }
+   * }
+   * }</pre>
+   *
+   * @param taskQueue the task queue name for the worker
+   * @param worker the worker being shut down
+   * @param next runnable that shuts down the next in chain (eventually shuts down the actual
+   *     worker)
+   */
+  default void shutdownWorker(
+      @Nonnull String taskQueue, @Nonnull Worker worker, @Nonnull Runnable next) {
+    next.run();
+  }
+
   /**
    * 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

@@ -263,8 +263,39 @@ public synchronized void start() {
 
   /** Internal method that actually starts the workers. Called from the plugin chain. */
   private void doStart() {
-    for (Worker worker : workers.values()) {
-      worker.start();
+    // Start each worker with plugin hooks
+    for (Map.Entry<String, Worker> entry : workers.entrySet()) {
+      String taskQueue = entry.getKey();
+      Worker worker = entry.getValue();
+
+      // Build plugin chain for this worker (reverse order for proper nesting)
+      Runnable startChain = worker::start;
+      List<Object> reversed = new ArrayList<>(plugins);
+      Collections.reverse(reversed);
+      for (Object plugin : reversed) {
+        if (plugin instanceof Plugin) {
+          final Runnable next = startChain;
+          final Plugin workerPlugin = (Plugin) plugin;
+          startChain =
+              () -> {
+                try {
+                  workerPlugin.startWorker(taskQueue, worker, next);
+                } catch (RuntimeException e) {
+                  throw e;
+                } catch (Exception e) {
+                  throw new RuntimeException(
+                      "Plugin "
+                          + workerPlugin.getName()
+                          + " failed during worker startup for task queue "
+                          + taskQueue,
+                      e);
+                }
+              };
+        }
+      }
+
+      // Execute the chain for this worker
+      startChain.run();
     }
 
     state = State.Started;
@@ -368,10 +399,51 @@ private void shutdownInternal(boolean interruptUserTasks) {
   private void doShutdown(boolean interruptUserTasks) {
     ((WorkflowClientInternal) workflowClient.getInternal()).deregisterWorkerFactory(this);
     ShutdownManager shutdownManager = new ShutdownManager();
-    CompletableFuture.allOf(
-            workers.values().stream()
-                .map(worker -> worker.shutdown(shutdownManager, interruptUserTasks))
-                .toArray(CompletableFuture[]::new))
+
+    // Shutdown each worker with plugin hooks
+    List<CompletableFuture<Void>> shutdownFutures = new ArrayList<>();
+    for (Map.Entry<String, Worker> entry : workers.entrySet()) {
+      String taskQueue = entry.getKey();
+      Worker worker = entry.getValue();
+
+      // Build plugin chain for this worker's shutdown (reverse order for proper nesting)
+      // We use a holder to capture the future from the terminal action
+      @SuppressWarnings("unchecked")
+      CompletableFuture<Void>[] futureHolder = new CompletableFuture[1];
+      Runnable shutdownChain =
+          () -> futureHolder[0] = worker.shutdown(shutdownManager, interruptUserTasks);
+
+      List<Object> reversed = new ArrayList<>(plugins);
+      Collections.reverse(reversed);
+      for (Object plugin : reversed) {
+        if (plugin instanceof Plugin) {
+          final Runnable next = shutdownChain;
+          final Plugin workerPlugin = (Plugin) plugin;
+          shutdownChain =
+              () -> {
+                try {
+                  workerPlugin.shutdownWorker(taskQueue, worker, next);
+                } catch (Exception e) {
+                  log.warn(
+                      "Plugin {} failed during worker shutdown for task queue {}",
+                      workerPlugin.getName(),
+                      taskQueue,
+                      e);
+                  // Still try to continue shutdown
+                  next.run();
+                }
+              };
+        }
+      }
+
+      // Execute the shutdown chain for this worker
+      shutdownChain.run();
+      if (futureHolder[0] != null) {
+        shutdownFutures.add(futureHolder[0]);
+      }
+    }
+
+    CompletableFuture.allOf(shutdownFutures.toArray(new CompletableFuture[0]))
         .thenApply(
             r -> {
               cache.invalidateAll();