diff --git a/temporal-sdk/src/main/java/io/temporal/client/ClientPlugin.java b/temporal-sdk/src/main/java/io/temporal/client/ClientPlugin.java new file mode 100644 index 000000000..3f2b6f259 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/client/ClientPlugin.java @@ -0,0 +1,126 @@ +/* + * Copyright (C) 2022 Temporal Technologies, Inc. All Rights Reserved. + * + * Copyright (C) 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Modifications copyright (C) 2017 Uber Technologies, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this material except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.temporal.client; + +import io.temporal.common.Experimental; +import io.temporal.common.SimplePlugin; +import io.temporal.serviceclient.WorkflowServiceStubs; +import io.temporal.serviceclient.WorkflowServiceStubs.ClientPluginCallback; +import io.temporal.serviceclient.WorkflowServiceStubsOptions; +import java.util.function.Supplier; +import javax.annotation.Nonnull; + +/** + * Plugin interface for customizing Temporal client configuration and lifecycle. + * + *

Plugins participate in two phases: + * + *

+ * + *

Example implementation: + * + *

{@code
+ * public class LoggingPlugin extends SimplePlugin {
+ *     public LoggingPlugin() {
+ *         super("my-org.logging");
+ *     }
+ *
+ *     @Override
+ *     public void configureClient(WorkflowClientOptions.Builder builder) {
+ *         // Add custom interceptor
+ *         builder.setInterceptors(new LoggingInterceptor());
+ *     }
+ *
+ *     @Override
+ *     public WorkflowServiceStubs connectServiceClient(
+ *             WorkflowServiceStubsOptions options,
+ *             Supplier<WorkflowServiceStubs> next) {
+ *         logger.info("Connecting to Temporal at {}", options.getTarget());
+ *         WorkflowServiceStubs stubs = next.get();
+ *         logger.info("Connected successfully");
+ *         return stubs;
+ *     }
+ * }
+ * }
+ * + * @see io.temporal.worker.WorkerPlugin + * @see SimplePlugin + */ +@Experimental +public interface ClientPlugin extends ClientPluginCallback { + + /** + * Returns a unique name for this plugin. Used for logging and duplicate detection. Recommended + * format: "organization.plugin-name" (e.g., "io.temporal.tracing") + * + * @return fully qualified plugin name + */ + @Nonnull + String getName(); + + /** + * Allows the plugin to modify service stubs options before the service stubs are created. Called + * during configuration phase in forward (registration) order. + * + * @param builder the options builder to modify + */ + @Override + void configureServiceStubs(@Nonnull WorkflowServiceStubsOptions.Builder builder); + + /** + * Allows the plugin to modify workflow client options before the client is created. Called during + * configuration phase in forward (registration) order. + * + * @param builder the options builder to modify + */ + void configureClient(@Nonnull WorkflowClientOptions.Builder builder); + + /** + * Allows the plugin to wrap service client connection. Called during connection phase in reverse + * order (first plugin wraps all others). + * + *

Example: + * + *

{@code
+   * @Override
+   * public WorkflowServiceStubs connectServiceClient(
+   *         WorkflowServiceStubsOptions options,
+   *         Supplier next) {
+   *     logger.info("Connecting to Temporal...");
+   *     WorkflowServiceStubs stubs = next.get();
+   *     logger.info("Connected successfully");
+   *     return stubs;
+   * }
+   * }
+ * + * @param options the final options being used for connection + * @param next supplier that creates the service stubs (calls next plugin or actual connection) + * @return the service stubs (possibly wrapped or decorated) + */ + @Override + @Nonnull + WorkflowServiceStubs connectServiceClient( + @Nonnull WorkflowServiceStubsOptions options, @Nonnull Supplier next); +} diff --git a/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientInternalImpl.java b/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientInternalImpl.java index 9b9a36897..6f475520f 100644 --- a/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientInternalImpl.java +++ b/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientInternalImpl.java @@ -65,7 +65,10 @@ public static WorkflowClient newInstance( WorkflowClientInternalImpl( WorkflowServiceStubs workflowServiceStubs, WorkflowClientOptions options) { - options = WorkflowClientOptions.newBuilder(options).validateAndBuildWithDefaults(); + // Apply plugin configuration phase (forward order), then validate + WorkflowClientOptions.Builder builder = WorkflowClientOptions.newBuilder(options); + applyClientPluginConfiguration(builder, options.getPlugins()); + options = builder.validateAndBuildWithDefaults(); workflowServiceStubs = new NamespaceInjectWorkflowServiceStubs(workflowServiceStubs, options.getNamespace()); this.options = options; @@ -771,4 +774,18 @@ public NexusStartWorkflowResponse startNexus( WorkflowInvocationHandler.closeAsyncInvocation(); } } + + /** + * Applies client plugin configuration phase. Plugins are called in forward (registration) order + * to modify the client options. + */ + private static void applyClientPluginConfiguration( + WorkflowClientOptions.Builder builder, ClientPlugin[] plugins) { + if (plugins == null || plugins.length == 0) { + return; + } + for (ClientPlugin plugin : plugins) { + plugin.configureClient(builder); + } + } } diff --git a/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientOptions.java b/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientOptions.java index 944f395a4..6e4a7a067 100644 --- a/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientOptions.java +++ b/temporal-sdk/src/main/java/io/temporal/client/WorkflowClientOptions.java @@ -1,6 +1,7 @@ package io.temporal.client; import io.temporal.api.enums.v1.QueryRejectCondition; +import io.temporal.common.Experimental; import io.temporal.common.context.ContextPropagator; import io.temporal.common.converter.DataConverter; import io.temporal.common.converter.GlobalDataConverter; @@ -47,6 +48,7 @@ public static final class Builder { private String binaryChecksum; private List contextPropagators; private QueryRejectCondition queryRejectCondition; + private ClientPlugin[] plugins; private Builder() {} @@ -61,6 +63,7 @@ private Builder(WorkflowClientOptions options) { binaryChecksum = options.binaryChecksum; contextPropagators = options.contextPropagators; queryRejectCondition = options.queryRejectCondition; + plugins = options.plugins; } public Builder setNamespace(String namespace) { @@ -132,6 +135,24 @@ public Builder setQueryRejectCondition(QueryRejectCondition queryRejectCondition return this; } + /** + * Sets the plugins to use with this client. Plugins can modify client configuration, intercept + * connection, and wrap execution lifecycle. + * + *

Plugins that also implement {@link io.temporal.worker.WorkerPlugin} are automatically + * propagated to workers created from this client. + * + * @param plugins the client plugins to use + * @return this builder for chaining + * @see io.temporal.client.ClientPlugin + * @see io.temporal.worker.WorkerPlugin + */ + @Experimental + public Builder setPlugins(ClientPlugin... plugins) { + this.plugins = Objects.requireNonNull(plugins); + return this; + } + public WorkflowClientOptions build() { return new WorkflowClientOptions( namespace, @@ -140,7 +161,8 @@ public WorkflowClientOptions build() { identity, binaryChecksum, contextPropagators, - queryRejectCondition); + queryRejectCondition, + plugins); } public WorkflowClientOptions validateAndBuildWithDefaults() { @@ -154,7 +176,8 @@ public WorkflowClientOptions validateAndBuildWithDefaults() { contextPropagators == null ? EMPTY_CONTEXT_PROPAGATORS : contextPropagators, queryRejectCondition == null ? QueryRejectCondition.QUERY_REJECT_CONDITION_UNSPECIFIED - : queryRejectCondition); + : queryRejectCondition, + plugins == null ? EMPTY_PLUGINS : plugins); } } @@ -163,6 +186,8 @@ public WorkflowClientOptions validateAndBuildWithDefaults() { private static final List EMPTY_CONTEXT_PROPAGATORS = Collections.emptyList(); + private static final ClientPlugin[] EMPTY_PLUGINS = new ClientPlugin[0]; + private final String namespace; private final DataConverter dataConverter; @@ -177,6 +202,8 @@ public WorkflowClientOptions validateAndBuildWithDefaults() { private final QueryRejectCondition queryRejectCondition; + private final ClientPlugin[] plugins; + private WorkflowClientOptions( String namespace, DataConverter dataConverter, @@ -184,7 +211,8 @@ private WorkflowClientOptions( String identity, String binaryChecksum, List contextPropagators, - QueryRejectCondition queryRejectCondition) { + QueryRejectCondition queryRejectCondition, + ClientPlugin[] plugins) { this.namespace = namespace; this.dataConverter = dataConverter; this.interceptors = interceptors; @@ -192,6 +220,7 @@ private WorkflowClientOptions( this.binaryChecksum = binaryChecksum; this.contextPropagators = contextPropagators; this.queryRejectCondition = queryRejectCondition; + this.plugins = plugins; } /** @@ -236,6 +265,19 @@ public QueryRejectCondition getQueryRejectCondition() { return queryRejectCondition; } + /** + * Returns the client plugins configured for this client. + * + *

Plugins that also implement {@link io.temporal.worker.WorkerPlugin} are automatically + * propagated to workers created from this client. + * + * @return the array of client plugins, never null + */ + @Experimental + public ClientPlugin[] getPlugins() { + return plugins; + } + @Override public String toString() { return "WorkflowClientOptions{" @@ -256,6 +298,8 @@ public String toString() { + contextPropagators + ", queryRejectCondition=" + queryRejectCondition + + ", plugins=" + + Arrays.toString(plugins) + '}'; } @@ -270,7 +314,8 @@ public boolean equals(Object o) { && com.google.common.base.Objects.equal(identity, that.identity) && com.google.common.base.Objects.equal(binaryChecksum, that.binaryChecksum) && com.google.common.base.Objects.equal(contextPropagators, that.contextPropagators) - && queryRejectCondition == that.queryRejectCondition; + && queryRejectCondition == that.queryRejectCondition + && Arrays.equals(plugins, that.plugins); } @Override @@ -282,6 +327,7 @@ public int hashCode() { identity, binaryChecksum, contextPropagators, - queryRejectCondition); + queryRejectCondition, + Arrays.hashCode(plugins)); } } diff --git a/temporal-sdk/src/main/java/io/temporal/common/SimplePlugin.java b/temporal-sdk/src/main/java/io/temporal/common/SimplePlugin.java new file mode 100644 index 000000000..1be916f01 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/common/SimplePlugin.java @@ -0,0 +1,535 @@ +/* + * Copyright (C) 2022 Temporal Technologies, Inc. All Rights Reserved. + * + * Copyright (C) 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Modifications copyright (C) 2017 Uber Technologies, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this material except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.temporal.common; + +import io.temporal.client.ClientPlugin; +import io.temporal.client.WorkflowClientOptions; +import io.temporal.common.context.ContextPropagator; +import io.temporal.common.interceptors.WorkerInterceptor; +import io.temporal.common.interceptors.WorkflowClientInterceptor; +import io.temporal.serviceclient.WorkflowServiceStubs; +import io.temporal.serviceclient.WorkflowServiceStubsOptions; +import io.temporal.worker.Worker; +import io.temporal.worker.WorkerFactory; +import io.temporal.worker.WorkerFactoryOptions; +import io.temporal.worker.WorkerOptions; +import io.temporal.worker.WorkerPlugin; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import java.util.function.BiConsumer; +import java.util.function.Consumer; +import java.util.function.Supplier; +import javax.annotation.Nonnull; + +/** + * A plugin that implements both {@link ClientPlugin} and {@link WorkerPlugin}. This class can be + * used in two ways: + * + *

    + *
  1. Builder pattern: Use {@link #newBuilder(String)} to declaratively configure a plugin + * with callbacks + *
  2. Subclassing: Extend this class and override specific methods for custom behavior + *
+ * + *

Builder Pattern Example

+ * + *
{@code
+ * SimplePlugin myPlugin = SimplePlugin.newBuilder("my-plugin")
+ *     .addWorkerInterceptors(new TracingInterceptor())
+ *     .addClientInterceptors(new LoggingInterceptor())
+ *     .customizeClient(b -> b.setIdentity("custom-identity"))
+ *     .build();
+ * }
+ * + *

Subclassing Example

+ * + *
{@code
+ * public class TracingPlugin extends SimplePlugin {
+ *     private final Tracer tracer;
+ *
+ *     public TracingPlugin(Tracer tracer) {
+ *         super("io.temporal.tracing");
+ *         this.tracer = tracer;
+ *     }
+ *
+ *     @Override
+ *     public void configureClient(WorkflowClientOptions.Builder builder) {
+ *         builder.setInterceptors(new TracingClientInterceptor(tracer));
+ *     }
+ * }
+ * }
+ * + *

Hybrid Example (Builder + Override)

+ * + *
{@code
+ * public class HybridPlugin extends SimplePlugin {
+ *     public HybridPlugin() {
+ *         super(SimplePlugin.newBuilder("hybrid")
+ *             .addClientInterceptors(new LoggingInterceptor()));
+ *     }
+ *
+ *     @Override
+ *     public void initializeWorker(String taskQueue, Worker worker) {
+ *         worker.registerWorkflowImplementationTypes(MyWorkflow.class);
+ *     }
+ * }
+ * }
+ * + * @see ClientPlugin + * @see WorkerPlugin + */ +@Experimental +public class SimplePlugin implements ClientPlugin, WorkerPlugin { + + private final String name; + private final List> stubsCustomizers; + private final List> clientCustomizers; + private final List> factoryCustomizers; + private final List> workerCustomizers; + private final List> workerInitializers; + private final List> workerStartCallbacks; + private final List> workerShutdownCallbacks; + private final List> workerFactoryStartCallbacks; + private final List> workerFactoryShutdownCallbacks; + private final List workerInterceptors; + private final List clientInterceptors; + private final List contextPropagators; + + /** + * Creates a new plugin with the specified name. Use this constructor when subclassing to override + * specific methods. + * + * @param name a unique name for this plugin, used for logging and duplicate detection. + * Recommended format: "organization.plugin-name" (e.g., "io.temporal.tracing") + * @throws NullPointerException if name is null + */ + protected SimplePlugin(@Nonnull String name) { + this.name = Objects.requireNonNull(name, "Plugin name cannot be null"); + this.stubsCustomizers = Collections.emptyList(); + this.clientCustomizers = Collections.emptyList(); + this.factoryCustomizers = Collections.emptyList(); + this.workerCustomizers = Collections.emptyList(); + this.workerInitializers = Collections.emptyList(); + this.workerStartCallbacks = Collections.emptyList(); + this.workerShutdownCallbacks = Collections.emptyList(); + this.workerFactoryStartCallbacks = Collections.emptyList(); + this.workerFactoryShutdownCallbacks = Collections.emptyList(); + this.workerInterceptors = Collections.emptyList(); + this.clientInterceptors = Collections.emptyList(); + this.contextPropagators = Collections.emptyList(); + } + + /** + * Creates a new plugin from a builder. Use this constructor when subclassing to combine builder + * configuration with method overrides. + * + * @param builder the builder with configuration + * @throws NullPointerException if builder is null + */ + protected SimplePlugin(@Nonnull Builder builder) { + Objects.requireNonNull(builder, "Builder cannot be null"); + this.name = builder.name; + this.stubsCustomizers = new ArrayList<>(builder.stubsCustomizers); + this.clientCustomizers = new ArrayList<>(builder.clientCustomizers); + this.factoryCustomizers = new ArrayList<>(builder.factoryCustomizers); + this.workerCustomizers = new ArrayList<>(builder.workerCustomizers); + this.workerInitializers = new ArrayList<>(builder.workerInitializers); + this.workerStartCallbacks = new ArrayList<>(builder.workerStartCallbacks); + this.workerShutdownCallbacks = new ArrayList<>(builder.workerShutdownCallbacks); + this.workerFactoryStartCallbacks = new ArrayList<>(builder.workerFactoryStartCallbacks); + this.workerFactoryShutdownCallbacks = new ArrayList<>(builder.workerFactoryShutdownCallbacks); + this.workerInterceptors = new ArrayList<>(builder.workerInterceptors); + this.clientInterceptors = new ArrayList<>(builder.clientInterceptors); + this.contextPropagators = new ArrayList<>(builder.contextPropagators); + } + + /** + * Creates a new builder with the specified plugin name. + * + * @param name a unique name for the plugin, used for logging and duplicate detection. Recommended + * format: "organization.plugin-name" (e.g., "my-org.tracing") + * @return a new builder instance + */ + public static Builder newBuilder(@Nonnull String name) { + return new Builder(name); + } + + @Override + @Nonnull + public String getName() { + return name; + } + + @Override + public void configureServiceStubs(@Nonnull WorkflowServiceStubsOptions.Builder builder) { + for (Consumer customizer : stubsCustomizers) { + customizer.accept(builder); + } + } + + @Override + public void configureClient(@Nonnull WorkflowClientOptions.Builder builder) { + // Apply customizers + for (Consumer customizer : clientCustomizers) { + customizer.accept(builder); + } + + // Add client interceptors + if (!clientInterceptors.isEmpty()) { + WorkflowClientInterceptor[] existing = builder.build().getInterceptors(); + List combined = + new ArrayList<>(existing != null ? Arrays.asList(existing) : new ArrayList<>()); + combined.addAll(clientInterceptors); + builder.setInterceptors(combined.toArray(new WorkflowClientInterceptor[0])); + } + + // Add context propagators + if (!contextPropagators.isEmpty()) { + List existing = builder.build().getContextPropagators(); + List combined = + new ArrayList<>(existing != null ? existing : new ArrayList<>()); + combined.addAll(contextPropagators); + builder.setContextPropagators(combined); + } + } + + @Override + public void configureWorkerFactory(@Nonnull WorkerFactoryOptions.Builder builder) { + // Apply customizers + for (Consumer customizer : factoryCustomizers) { + customizer.accept(builder); + } + + // Add worker interceptors + if (!workerInterceptors.isEmpty()) { + WorkerInterceptor[] existing = builder.build().getWorkerInterceptors(); + List combined = + new ArrayList<>(existing != null ? Arrays.asList(existing) : new ArrayList<>()); + combined.addAll(workerInterceptors); + builder.setWorkerInterceptors(combined.toArray(new WorkerInterceptor[0])); + } + } + + @Override + public void configureWorker(@Nonnull String taskQueue, @Nonnull WorkerOptions.Builder builder) { + for (Consumer customizer : workerCustomizers) { + customizer.accept(builder); + } + } + + @Override + public void initializeWorker(@Nonnull String taskQueue, @Nonnull Worker worker) { + for (BiConsumer initializer : workerInitializers) { + initializer.accept(taskQueue, worker); + } + } + + @Override + public void startWorker(@Nonnull String taskQueue, @Nonnull Worker worker, @Nonnull Runnable next) + throws Exception { + next.run(); + for (BiConsumer callback : workerStartCallbacks) { + callback.accept(taskQueue, worker); + } + } + + @Override + public void shutdownWorker( + @Nonnull String taskQueue, @Nonnull Worker worker, @Nonnull Runnable next) { + for (BiConsumer callback : workerShutdownCallbacks) { + callback.accept(taskQueue, worker); + } + next.run(); + } + + @Override + public WorkflowServiceStubs connectServiceClient( + WorkflowServiceStubsOptions options, Supplier next) { + return next.get(); + } + + @Override + public void startWorkerFactory(WorkerFactory factory, Runnable next) throws Exception { + next.run(); + for (Consumer callback : workerFactoryStartCallbacks) { + callback.accept(factory); + } + } + + @Override + public void shutdownWorkerFactory(WorkerFactory factory, Runnable next) throws Exception { + for (Consumer callback : workerFactoryShutdownCallbacks) { + callback.accept(factory); + } + next.run(); + } + + @Override + public String toString() { + return getClass().getSimpleName() + "{name='" + name + "'}"; + } + + /** Builder for creating {@link SimplePlugin} instances with declarative configuration. */ + public static final class Builder { + + private final String name; + private final List> stubsCustomizers = + new ArrayList<>(); + private final List> clientCustomizers = + new ArrayList<>(); + private final List> factoryCustomizers = + new ArrayList<>(); + private final List> workerCustomizers = new ArrayList<>(); + private final List> workerInitializers = new ArrayList<>(); + private final List> workerStartCallbacks = new ArrayList<>(); + private final List> workerShutdownCallbacks = new ArrayList<>(); + private final List> workerFactoryStartCallbacks = new ArrayList<>(); + private final List> workerFactoryShutdownCallbacks = new ArrayList<>(); + private final List workerInterceptors = new ArrayList<>(); + private final List clientInterceptors = new ArrayList<>(); + private final List contextPropagators = new ArrayList<>(); + + private Builder(@Nonnull String name) { + this.name = Objects.requireNonNull(name, "Plugin name cannot be null"); + } + + /** + * Adds a customizer for {@link WorkflowServiceStubsOptions}. Multiple customizers are applied + * in the order they are added. + * + * @param customizer a consumer that modifies the options builder + * @return this builder for chaining + */ + public Builder customizeServiceStubs( + @Nonnull Consumer customizer) { + stubsCustomizers.add(Objects.requireNonNull(customizer)); + return this; + } + + /** + * Adds a customizer for {@link WorkflowClientOptions}. Multiple customizers are applied in the + * order they are added. + * + * @param customizer a consumer that modifies the options builder + * @return this builder for chaining + */ + public Builder customizeClient(@Nonnull Consumer customizer) { + clientCustomizers.add(Objects.requireNonNull(customizer)); + return this; + } + + /** + * Adds a customizer for {@link WorkerFactoryOptions}. Multiple customizers are applied in the + * order they are added. + * + * @param customizer a consumer that modifies the options builder + * @return this builder for chaining + */ + public Builder customizeWorkerFactory( + @Nonnull Consumer customizer) { + factoryCustomizers.add(Objects.requireNonNull(customizer)); + return this; + } + + /** + * Adds a customizer for {@link WorkerOptions}. Multiple customizers are applied in the order + * they are added. The customizer is applied to all workers created by the factory. + * + * @param customizer a consumer that modifies the options builder + * @return this builder for chaining + */ + public Builder customizeWorker(@Nonnull Consumer customizer) { + workerCustomizers.add(Objects.requireNonNull(customizer)); + 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. + * + *

Example: + * + *

{@code
+     * SimplePlugin.newBuilder("my-plugin")
+     *     .initializeWorker((taskQueue, worker) -> {
+     *         worker.registerWorkflowImplementationTypes(MyWorkflow.class);
+     *         worker.registerActivitiesImplementations(new MyActivityImpl());
+     *     })
+     *     .build();
+     * }
+ * + * @param initializer a consumer that receives the task queue name and worker + * @return this builder for chaining + */ + public Builder initializeWorker(@Nonnull BiConsumer initializer) { + workerInitializers.add(Objects.requireNonNull(initializer)); + return this; + } + + /** + * Adds a callback that is invoked when a worker starts. This can be used to start per-worker + * resources or record metrics. + * + *

Note: For registering workflows and activities, use {@link #initializeWorker} instead, as + * registrations must happen before the worker starts polling. + * + *

Example: + * + *

{@code
+     * SimplePlugin.newBuilder("my-plugin")
+     *     .onWorkerStart((taskQueue, worker) -> {
+     *         logger.info("Worker started for task queue: {}", taskQueue);
+     *         perWorkerResources.put(taskQueue, new ResourcePool());
+     *     })
+     *     .build();
+     * }
+ * + * @param callback a consumer that receives the task queue name and worker when the worker + * starts + * @return this builder for chaining + */ + public Builder onWorkerStart(@Nonnull BiConsumer 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}. + * + *

Example: + * + *

{@code
+     * SimplePlugin.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();
+     * }
+ * + * @param callback a consumer that receives the task queue name and worker when the worker shuts + * down + * @return this builder for chaining + */ + public Builder onWorkerShutdown(@Nonnull BiConsumer callback) { + workerShutdownCallbacks.add(Objects.requireNonNull(callback)); + return this; + } + + /** + * Adds a callback that is invoked when the worker factory starts. This can be used to + * initialize factory-level resources or record metrics. + * + *

Example: + * + *

{@code
+     * SimplePlugin.newBuilder("my-plugin")
+     *     .onWorkerFactoryStart(factory -> {
+     *         logger.info("Worker factory started");
+     *         globalResources.initialize();
+     *     })
+     *     .build();
+     * }
+ * + * @param callback a consumer that receives the worker factory when it starts + * @return this builder for chaining + */ + public Builder onWorkerFactoryStart(@Nonnull Consumer callback) { + workerFactoryStartCallbacks.add(Objects.requireNonNull(callback)); + return this; + } + + /** + * Adds a callback that is invoked when the worker factory shuts down. This can be used to clean + * up factory-level resources. + * + *

Example: + * + *

{@code
+     * SimplePlugin.newBuilder("my-plugin")
+     *     .onWorkerFactoryShutdown(factory -> {
+     *         logger.info("Worker factory shutting down");
+     *         globalResources.cleanup();
+     *     })
+     *     .build();
+     * }
+ * + * @param callback a consumer that receives the worker factory when it shuts down + * @return this builder for chaining + */ + public Builder onWorkerFactoryShutdown(@Nonnull Consumer callback) { + workerFactoryShutdownCallbacks.add(Objects.requireNonNull(callback)); + return this; + } + + /** + * Adds worker interceptors. Interceptors are appended to any existing interceptors in the + * configuration. + * + * @param interceptors the interceptors to add + * @return this builder for chaining + */ + public Builder addWorkerInterceptors(WorkerInterceptor... interceptors) { + workerInterceptors.addAll(Arrays.asList(interceptors)); + return this; + } + + /** + * Adds client interceptors. Interceptors are appended to any existing interceptors in the + * configuration. + * + * @param interceptors the interceptors to add + * @return this builder for chaining + */ + public Builder addClientInterceptors(WorkflowClientInterceptor... interceptors) { + clientInterceptors.addAll(Arrays.asList(interceptors)); + return this; + } + + /** + * Adds context propagators. Propagators are appended to any existing propagators in the + * configuration. + * + * @param propagators the propagators to add + * @return this builder for chaining + */ + public Builder addContextPropagators(ContextPropagator... propagators) { + contextPropagators.addAll(Arrays.asList(propagators)); + return this; + } + + /** + * Builds the plugin with the configured settings. + * + * @return a new plugin instance + */ + public SimplePlugin build() { + return new SimplePlugin(this); + } + } +} diff --git a/temporal-sdk/src/main/java/io/temporal/worker/WorkerFactory.java b/temporal-sdk/src/main/java/io/temporal/worker/WorkerFactory.java index 20540f4a9..c9ae8b82a 100644 --- a/temporal-sdk/src/main/java/io/temporal/worker/WorkerFactory.java +++ b/temporal-sdk/src/main/java/io/temporal/worker/WorkerFactory.java @@ -15,7 +15,10 @@ import io.temporal.internal.worker.WorkflowExecutorCache; import io.temporal.internal.worker.WorkflowRunLockManager; import io.temporal.serviceclient.MetricsTag; +import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.Objects; import java.util.concurrent.CompletableFuture; @@ -46,6 +49,9 @@ public final class WorkerFactory { private final @Nonnull WorkflowExecutorCache cache; + /** Plugins propagated from the client and applied to this factory. */ + private final List plugins; + private State state = State.Initial; private final String statusErrorMessage = @@ -72,6 +78,12 @@ private WorkerFactory(WorkflowClient workflowClient, WorkerFactoryOptions factor WorkflowClientOptions workflowClientOptions = workflowClient.getOptions(); String namespace = workflowClientOptions.getNamespace(); + // Extract worker plugins from client (auto-propagation) + this.plugins = extractWorkerPlugins(workflowClientOptions.getPlugins()); + + // Apply plugin configuration to factory options (forward order) + factoryOptions = applyPluginConfiguration(factoryOptions, this.plugins); + this.factoryOptions = WorkerFactoryOptions.newBuilder(factoryOptions).validateAndBuildWithDefaults(); @@ -137,6 +149,9 @@ public synchronized Worker newWorker(String taskQueue, WorkerOptions options) { state == State.Initial, String.format(statusErrorMessage, "create new worker", state.name(), State.Initial.name())); + // Apply plugin configuration to worker options (forward order) + options = applyWorkerPluginConfiguration(taskQueue, options, this.plugins); + // Only one worker can exist for a task queue Worker existingWorker = workers.get(taskQueue); if (existingWorker == null) { @@ -153,6 +168,13 @@ 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 (WorkerPlugin plugin : plugins) { + plugin.initializeWorker(taskQueue, worker); + } + return worker; } else { log.warn( @@ -211,8 +233,59 @@ public synchronized void start() { .setNamespace(workflowClient.getOptions().getNamespace()) .build()); - for (Worker worker : workers.values()) { - worker.start(); + // Build plugin execution chain (reverse order for proper nesting) + Runnable startChain = this::doStart; + for (int i = plugins.size() - 1; i >= 0; i--) { + final Runnable next = startChain; + final WorkerPlugin workerPlugin = plugins.get(i); + startChain = + () -> { + try { + workerPlugin.startWorkerFactory(this, next); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new RuntimeException( + "Plugin " + workerPlugin.getName() + " failed during startup", e); + } + }; + } + + // Execute the chain + startChain.run(); + } + + /** Internal method that actually starts the workers. Called from the plugin chain. */ + private void doStart() { + // Start each worker with plugin hooks + for (Map.Entry 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; + for (int i = plugins.size() - 1; i >= 0; i--) { + final Runnable next = startChain; + final WorkerPlugin workerPlugin = plugins.get(i); + 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; @@ -286,12 +359,73 @@ public synchronized void shutdownNow() { private void shutdownInternal(boolean interruptUserTasks) { state = State.Shutdown; + + // Build plugin shutdown chain (reverse order for proper nesting) + Runnable shutdownChain = () -> doShutdown(interruptUserTasks); + for (int i = plugins.size() - 1; i >= 0; i--) { + final Runnable next = shutdownChain; + final WorkerPlugin workerPlugin = plugins.get(i); + shutdownChain = + () -> { + try { + workerPlugin.shutdownWorkerFactory(this, next); + } catch (Exception e) { + log.warn("Plugin {} failed during shutdown", workerPlugin.getName(), e); + // Still try to continue shutdown + next.run(); + } + }; + } + + // Execute the chain + shutdownChain.run(); + } + + /** Internal method that actually shuts down workers. Called from the plugin chain. */ + 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> shutdownFutures = new ArrayList<>(); + for (Map.Entry 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[] futureHolder = new CompletableFuture[1]; + Runnable shutdownChain = + () -> futureHolder[0] = worker.shutdown(shutdownManager, interruptUserTasks); + + for (int i = plugins.size() - 1; i >= 0; i--) { + final Runnable next = shutdownChain; + final WorkerPlugin workerPlugin = plugins.get(i); + 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(); @@ -359,6 +493,65 @@ public String toString() { return String.format("WorkerFactory{identity=%s}", workflowClient.getOptions().getIdentity()); } + /** + * Extracts worker plugins from the client plugins array. Only plugins that also implement {@link + * WorkerPlugin} are included. + */ + private static List extractWorkerPlugins( + io.temporal.client.ClientPlugin[] clientPlugins) { + if (clientPlugins == null || clientPlugins.length == 0) { + return Collections.emptyList(); + } + + List workerPlugins = new ArrayList<>(); + for (io.temporal.client.ClientPlugin plugin : clientPlugins) { + if (plugin instanceof WorkerPlugin) { + workerPlugins.add((WorkerPlugin) plugin); + } + } + return Collections.unmodifiableList(workerPlugins); + } + + /** + * Applies plugin configuration to worker factory options. Plugins are called in forward + * (registration) order. + */ + private static WorkerFactoryOptions applyPluginConfiguration( + WorkerFactoryOptions options, List plugins) { + if (plugins == null || plugins.isEmpty()) { + return options; + } + + WorkerFactoryOptions.Builder builder = + options == null + ? WorkerFactoryOptions.newBuilder() + : WorkerFactoryOptions.newBuilder(options); + + for (WorkerPlugin plugin : plugins) { + plugin.configureWorkerFactory(builder); + } + return builder.build(); + } + + /** + * Applies plugin configuration to worker options. Plugins are called in forward (registration) + * order. + */ + private static WorkerOptions applyWorkerPluginConfiguration( + String taskQueue, WorkerOptions options, List plugins) { + if (plugins == null || plugins.isEmpty()) { + return options; + } + + WorkerOptions.Builder builder = + options == null ? WorkerOptions.newBuilder() : WorkerOptions.newBuilder(options); + + for (WorkerPlugin plugin : plugins) { + plugin.configureWorker(taskQueue, builder); + } + return builder.build(); + } + enum State { Initial, Started, diff --git a/temporal-sdk/src/main/java/io/temporal/worker/WorkerFactoryOptions.java b/temporal-sdk/src/main/java/io/temporal/worker/WorkerFactoryOptions.java index c50da81cd..8fda76d20 100644 --- a/temporal-sdk/src/main/java/io/temporal/worker/WorkerFactoryOptions.java +++ b/temporal-sdk/src/main/java/io/temporal/worker/WorkerFactoryOptions.java @@ -37,6 +37,7 @@ public static class Builder { private int workflowCacheSize; private int maxWorkflowThreadCount; private WorkerInterceptor[] workerInterceptors; + private WorkerPlugin[] plugins; private boolean enableLoggingInReplay; private boolean usingVirtualWorkflowThreads; private ExecutorService overrideLocalActivityTaskExecutor; @@ -52,6 +53,7 @@ private Builder(WorkerFactoryOptions options) { this.workflowCacheSize = options.workflowCacheSize; this.maxWorkflowThreadCount = options.maxWorkflowThreadCount; this.workerInterceptors = options.workerInterceptors; + this.plugins = options.plugins; this.enableLoggingInReplay = options.enableLoggingInReplay; this.usingVirtualWorkflowThreads = options.usingVirtualWorkflowThreads; this.overrideLocalActivityTaskExecutor = options.overrideLocalActivityTaskExecutor; @@ -101,6 +103,24 @@ public Builder setWorkerInterceptors(WorkerInterceptor... workerInterceptors) { return this; } + /** + * Sets the worker plugins to use with workers created by this factory. Plugins can modify worker + * configuration and wrap worker lifecycle. + * + *

Note: Plugins that implement both {@link io.temporal.client.ClientPlugin} and {@link + * WorkerPlugin} are automatically propagated from the client. Use this method for worker-only + * plugins that don't need client-side configuration. + * + * @param plugins the worker plugins to use + * @return this builder for chaining + * @see WorkerPlugin + */ + @Experimental + public Builder setPlugins(WorkerPlugin... plugins) { + this.plugins = plugins; + return this; + } + public Builder setEnableLoggingInReplay(boolean enableLoggingInReplay) { this.enableLoggingInReplay = enableLoggingInReplay; return this; @@ -141,6 +161,7 @@ public WorkerFactoryOptions build() { maxWorkflowThreadCount, workflowHostLocalTaskQueueScheduleToStartTimeout, workerInterceptors, + plugins, enableLoggingInReplay, usingVirtualWorkflowThreads, overrideLocalActivityTaskExecutor, @@ -153,6 +174,7 @@ public WorkerFactoryOptions validateAndBuildWithDefaults() { maxWorkflowThreadCount, workflowHostLocalTaskQueueScheduleToStartTimeout, workerInterceptors == null ? new WorkerInterceptor[0] : workerInterceptors, + plugins == null ? new WorkerPlugin[0] : plugins, enableLoggingInReplay, usingVirtualWorkflowThreads, overrideLocalActivityTaskExecutor, @@ -164,6 +186,7 @@ public WorkerFactoryOptions validateAndBuildWithDefaults() { private final int maxWorkflowThreadCount; private final @Nullable Duration workflowHostLocalTaskQueueScheduleToStartTimeout; private final WorkerInterceptor[] workerInterceptors; + private final WorkerPlugin[] plugins; private final boolean enableLoggingInReplay; private final boolean usingVirtualWorkflowThreads; private final ExecutorService overrideLocalActivityTaskExecutor; @@ -173,6 +196,7 @@ private WorkerFactoryOptions( int maxWorkflowThreadCount, @Nullable Duration workflowHostLocalTaskQueueScheduleToStartTimeout, WorkerInterceptor[] workerInterceptors, + WorkerPlugin[] plugins, boolean enableLoggingInReplay, boolean usingVirtualWorkflowThreads, ExecutorService overrideLocalActivityTaskExecutor, @@ -195,12 +219,16 @@ private WorkerFactoryOptions( if (workerInterceptors == null) { workerInterceptors = new WorkerInterceptor[0]; } + if (plugins == null) { + plugins = new WorkerPlugin[0]; + } } this.workflowCacheSize = workflowCacheSize; this.maxWorkflowThreadCount = maxWorkflowThreadCount; this.workflowHostLocalTaskQueueScheduleToStartTimeout = workflowHostLocalTaskQueueScheduleToStartTimeout; this.workerInterceptors = workerInterceptors; + this.plugins = plugins; this.enableLoggingInReplay = enableLoggingInReplay; this.usingVirtualWorkflowThreads = usingVirtualWorkflowThreads; this.overrideLocalActivityTaskExecutor = overrideLocalActivityTaskExecutor; @@ -223,6 +251,16 @@ public WorkerInterceptor[] getWorkerInterceptors() { return workerInterceptors; } + /** + * Returns the worker plugins configured for this factory. + * + * @return the array of worker plugins, never null + */ + @Experimental + public WorkerPlugin[] getPlugins() { + return plugins; + } + public boolean isEnableLoggingInReplay() { return enableLoggingInReplay; } diff --git a/temporal-sdk/src/main/java/io/temporal/worker/WorkerPlugin.java b/temporal-sdk/src/main/java/io/temporal/worker/WorkerPlugin.java new file mode 100644 index 000000000..5446645f0 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/worker/WorkerPlugin.java @@ -0,0 +1,221 @@ +/* + * Copyright (C) 2022 Temporal Technologies, Inc. All Rights Reserved. + * + * Copyright (C) 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Modifications copyright (C) 2017 Uber Technologies, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this material except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.temporal.worker; + +import io.temporal.common.Experimental; +import io.temporal.common.SimplePlugin; +import javax.annotation.Nonnull; + +/** + * Plugin interface for customizing Temporal worker configuration and lifecycle. + * + *

Plugins that also implement {@link io.temporal.client.ClientPlugin} are automatically + * propagated from the client to workers created from that client. + * + *

Example implementation: + * + *

{@code
+ * public class MetricsPlugin extends SimplePlugin {
+ *     private final MetricsRegistry registry;
+ *
+ *     public MetricsPlugin(MetricsRegistry registry) {
+ *         super("my-org.metrics");
+ *         this.registry = registry;
+ *     }
+ *
+ *     @Override
+ *     public void configureWorkerFactory(WorkerFactoryOptions.Builder builder) {
+ *         builder.setWorkerInterceptors(new MetricsWorkerInterceptor(registry));
+ *     }
+ *
+ *     @Override
+ *     public void startWorkerFactory(WorkerFactory factory, Runnable next) throws Exception {
+ *         registry.recordWorkerStart();
+ *         try {
+ *             next.run();
+ *         } finally {
+ *             registry.recordWorkerStop();
+ *         }
+ *     }
+ * }
+ * }
+ * + * @see io.temporal.client.ClientPlugin + * @see SimplePlugin + */ +@Experimental +public interface WorkerPlugin { + + /** + * Returns a unique name for this plugin. Used for logging and duplicate detection. Recommended + * format: "organization.plugin-name" (e.g., "io.temporal.tracing") + * + * @return fully qualified plugin name + */ + @Nonnull + String getName(); + + /** + * Allows the plugin to modify worker factory options before the factory is created. Called during + * configuration phase in forward (registration) order. + * + * @param builder the options builder to modify + */ + void configureWorkerFactory(@Nonnull WorkerFactoryOptions.Builder builder); + + /** + * Allows the plugin to modify worker options before a worker is created. Called during + * configuration phase in forward (registration) order. + * + * @param taskQueue the task queue name for the worker being created + * @param builder the options builder to modify + */ + void configureWorker(@Nonnull String taskQueue, @Nonnull WorkerOptions.Builder builder); + + /** + * Called after a worker is created, allowing plugins to register workflows, activities, Nexus + * services, and other components on the worker. + * + *

This method is called in forward (registration) order immediately after the worker is + * created in {@link WorkerFactory#newWorker}. This is the appropriate place for registrations + * because it is called before the worker starts polling. + * + *

Example: + * + *

{@code
+   * @Override
+   * public void initializeWorker(String taskQueue, Worker worker) {
+   *     worker.registerWorkflowImplementationTypes(MyWorkflow.class);
+   *     worker.registerActivitiesImplementations(new MyActivityImpl());
+   * }
+   * }
+ * + * @param taskQueue the task queue name for the worker + * @param worker the newly created worker + */ + void initializeWorker(@Nonnull String taskQueue, @Nonnull Worker worker); + + /** + * 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. + * + *

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. + * + *

Example: + * + *

{@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();
+   * }
+   * }
+ * + * @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 + */ + void startWorker(@Nonnull String taskQueue, @Nonnull Worker worker, @Nonnull Runnable next) + throws Exception; + + /** + * 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. + * + *

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}. + * + *

Example: + * + *

{@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();
+   *     }
+   * }
+   * }
+ * + * @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) + */ + void shutdownWorker(@Nonnull String taskQueue, @Nonnull Worker worker, @Nonnull Runnable next); + + /** + * Allows the plugin to wrap worker factory startup. Called during execution phase in reverse + * order (first plugin wraps all others). + * + *

This method is called when {@link WorkerFactory#start()} is invoked. The plugin can perform + * setup before starting and cleanup logic. + * + *

Example: + * + *

{@code
+   * @Override
+   * public void startWorkerFactory(WorkerFactory factory, Runnable next) throws Exception {
+   *     logger.info("Starting workers...");
+   *     next.run();
+   *     logger.info("Workers started");
+   * }
+   * }
+ * + * @param factory the worker factory being started + * @param next runnable that starts the next in chain (eventually starts actual workers) + * @throws Exception if startup fails + */ + void startWorkerFactory(@Nonnull WorkerFactory factory, @Nonnull Runnable next) throws Exception; + + /** + * Allows the plugin to wrap worker factory shutdown. Called during shutdown phase in reverse + * order (first plugin wraps all others). + * + *

This method is called when {@link WorkerFactory#shutdown()} or {@link + * WorkerFactory#shutdownNow()} is invoked. The plugin can perform actions before and after the + * actual shutdown occurs. + * + *

Example: + * + *

{@code
+   * @Override
+   * public void shutdownWorkerFactory(WorkerFactory factory, Runnable next) {
+   *     logger.info("Shutting down workers...");
+   *     next.run();
+   *     logger.info("Workers shut down");
+   * }
+   * }
+ * + * @param factory the worker factory being shut down + * @param next runnable that shuts down the next in chain (eventually shuts down actual workers) + */ + void shutdownWorkerFactory(@Nonnull WorkerFactory factory, @Nonnull Runnable next) + throws Exception; +} diff --git a/temporal-sdk/src/test/java/io/temporal/client/WorkflowClientOptionsPluginTest.java b/temporal-sdk/src/test/java/io/temporal/client/WorkflowClientOptionsPluginTest.java new file mode 100644 index 000000000..af971355f --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/client/WorkflowClientOptionsPluginTest.java @@ -0,0 +1,100 @@ +/* + * Copyright (C) 2022 Temporal Technologies, Inc. All Rights Reserved. + * + * Copyright (C) 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Modifications copyright (C) 2017 Uber Technologies, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this material except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.temporal.client; + +import static org.junit.Assert.*; + +import io.temporal.common.SimplePlugin; +import org.junit.Test; + +public class WorkflowClientOptionsPluginTest { + + @Test + public void testDefaultPluginsEmpty() { + WorkflowClientOptions options = WorkflowClientOptions.newBuilder().build(); + assertEquals("Default plugins should be empty", 0, options.getPlugins().length); + } + + @Test + public void testSetPlugins() { + SimplePlugin plugin1 = new TestPlugin("plugin1"); + SimplePlugin plugin2 = new TestPlugin("plugin2"); + + WorkflowClientOptions options = + WorkflowClientOptions.newBuilder().setPlugins(plugin1, plugin2).build(); + + ClientPlugin[] plugins = options.getPlugins(); + assertEquals(2, plugins.length); + assertEquals("plugin1", plugins[0].getName()); + assertEquals("plugin2", plugins[1].getName()); + } + + @Test + public void testToBuilder() { + SimplePlugin plugin = new TestPlugin("plugin"); + + WorkflowClientOptions original = WorkflowClientOptions.newBuilder().setPlugins(plugin).build(); + + WorkflowClientOptions copy = original.toBuilder().build(); + + assertEquals(1, copy.getPlugins().length); + assertEquals("plugin", copy.getPlugins()[0].getName()); + } + + @Test + public void testValidateAndBuildWithDefaults() { + SimplePlugin plugin = new TestPlugin("plugin"); + + WorkflowClientOptions options = + WorkflowClientOptions.newBuilder().setPlugins(plugin).validateAndBuildWithDefaults(); + + assertEquals(1, options.getPlugins().length); + assertEquals("plugin", options.getPlugins()[0].getName()); + } + + @Test + public void testEqualsWithPlugins() { + SimplePlugin plugin = new TestPlugin("plugin"); + + WorkflowClientOptions options1 = WorkflowClientOptions.newBuilder().setPlugins(plugin).build(); + + WorkflowClientOptions options2 = WorkflowClientOptions.newBuilder().setPlugins(plugin).build(); + + assertEquals(options1, options2); + assertEquals(options1.hashCode(), options2.hashCode()); + } + + @Test + public void testToStringWithPlugins() { + SimplePlugin plugin = new TestPlugin("my-plugin"); + + WorkflowClientOptions options = WorkflowClientOptions.newBuilder().setPlugins(plugin).build(); + + String str = options.toString(); + assertTrue("toString should contain plugins", str.contains("plugins")); + } + + private static class TestPlugin extends SimplePlugin { + TestPlugin(String name) { + super(name); + } + } +} diff --git a/temporal-sdk/src/test/java/io/temporal/common/PluginTest.java b/temporal-sdk/src/test/java/io/temporal/common/PluginTest.java new file mode 100644 index 000000000..ec41acbb1 --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/common/PluginTest.java @@ -0,0 +1,301 @@ +/* + * Copyright (C) 2022 Temporal Technologies, Inc. All Rights Reserved. + * + * Copyright (C) 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Modifications copyright (C) 2017 Uber Technologies, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this material except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.temporal.common; + +import static org.junit.Assert.*; + +import io.temporal.client.WorkflowClientOptions; +import io.temporal.serviceclient.WorkflowServiceStubsOptions; +import io.temporal.worker.WorkerFactoryOptions; +import io.temporal.worker.WorkerOptions; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.junit.Test; + +public class PluginTest { + + @Test + public void testSimplePluginName() { + SimplePlugin plugin = new SimplePlugin("test-plugin") {}; + assertEquals("test-plugin", plugin.getName()); + } + + @Test + public void testSimplePluginToString() { + SimplePlugin plugin = new SimplePlugin("my-plugin") {}; + assertTrue(plugin.toString().contains("my-plugin")); + } + + @Test(expected = NullPointerException.class) + public void testSimplePluginNullName() { + new SimplePlugin((String) null) {}; + } + + @Test + public void testSimplePluginDefaultBehavior() throws Exception { + SimplePlugin plugin = new SimplePlugin("test") {}; + + // Test configureServiceStubs doesn't throw (no customizers) + WorkflowServiceStubsOptions.Builder stubsBuilder = WorkflowServiceStubsOptions.newBuilder(); + plugin.configureServiceStubs(stubsBuilder); + + // Test configureClient doesn't throw (no customizers) + WorkflowClientOptions.Builder clientBuilder = WorkflowClientOptions.newBuilder(); + plugin.configureClient(clientBuilder); + + // Test configureWorkerFactory doesn't throw (no customizers) + WorkerFactoryOptions.Builder factoryBuilder = WorkerFactoryOptions.newBuilder(); + plugin.configureWorkerFactory(factoryBuilder); + + // Test configureWorker doesn't throw (no customizers) + WorkerOptions.Builder workerBuilder = WorkerOptions.newBuilder(); + plugin.configureWorker("test-queue", workerBuilder); + + // Test startWorkerFactory calls next + final boolean[] called = {false}; + plugin.startWorkerFactory(null, () -> called[0] = true); + assertTrue("startWorkerFactory should call next", called[0]); + + // Test startWorker calls next + called[0] = false; + plugin.startWorker("test-queue", null, () -> called[0] = true); + assertTrue("startWorker should call next", called[0]); + + // Test shutdownWorker calls next + called[0] = false; + plugin.shutdownWorker("test-queue", null, () -> called[0] = true); + assertTrue("shutdownWorker should call next", called[0]); + + // Test shutdownWorkerFactory calls next + called[0] = false; + plugin.shutdownWorkerFactory(null, () -> called[0] = true); + assertTrue("shutdownWorkerFactory should call next", called[0]); + + // Test initializeWorker is a no-op (doesn't throw) + plugin.initializeWorker("test-queue", null); + } + + @Test + public void testConfigurationPhaseOrder() { + List order = new ArrayList<>(); + + SimplePlugin pluginA = createTrackingPlugin("A", order); + SimplePlugin pluginB = createTrackingPlugin("B", order); + SimplePlugin pluginC = createTrackingPlugin("C", order); + + List plugins = Arrays.asList(pluginA, pluginB, pluginC); + + // Simulate configuration phase (forward order) + WorkflowClientOptions.Builder builder = WorkflowClientOptions.newBuilder(); + for (Object plugin : plugins) { + if (plugin instanceof io.temporal.client.ClientPlugin) { + ((io.temporal.client.ClientPlugin) plugin).configureClient(builder); + } + } + + // Configuration should be in forward order + assertEquals(Arrays.asList("A-config", "B-config", "C-config"), order); + } + + @Test + public void testExecutionPhaseReverseOrder() throws Exception { + List order = new ArrayList<>(); + + SimplePlugin pluginA = createExecutionTrackingPlugin("A", order); + SimplePlugin pluginB = createExecutionTrackingPlugin("B", order); + SimplePlugin pluginC = createExecutionTrackingPlugin("C", order); + + List plugins = Arrays.asList(pluginA, pluginB, pluginC); + + // Build chain in reverse (like WorkerFactory does) + Runnable chain = + () -> { + order.add("terminal"); + }; + + List reversed = new ArrayList<>(plugins); + java.util.Collections.reverse(reversed); + for (Object plugin : reversed) { + if (plugin instanceof io.temporal.worker.WorkerPlugin) { + final Runnable next = chain; + final io.temporal.worker.WorkerPlugin workerPlugin = + (io.temporal.worker.WorkerPlugin) plugin; + chain = + () -> { + order.add(workerPlugin.getName() + "-before"); + try { + workerPlugin.startWorkerFactory(null, next); + } catch (Exception e) { + throw new RuntimeException(e); + } + order.add(workerPlugin.getName() + "-after"); + }; + } + } + + // Execute the chain + chain.run(); + + // First plugin should wrap all others + assertEquals( + Arrays.asList( + "A-before", "B-before", "C-before", "terminal", "C-after", "B-after", "A-after"), + order); + } + + @Test + public void testStartWorkerReverseOrder() throws Exception { + List order = new ArrayList<>(); + + SimplePlugin pluginA = createWorkerLifecycleTrackingPlugin("A", order); + SimplePlugin pluginB = createWorkerLifecycleTrackingPlugin("B", order); + SimplePlugin pluginC = createWorkerLifecycleTrackingPlugin("C", order); + + List plugins = Arrays.asList(pluginA, pluginB, pluginC); + + // Build chain in reverse (like WorkerFactory does) + Runnable chain = () -> order.add("worker-start"); + + List reversed = new ArrayList<>(plugins); + java.util.Collections.reverse(reversed); + for (Object plugin : reversed) { + if (plugin instanceof io.temporal.worker.WorkerPlugin) { + final Runnable next = chain; + final io.temporal.worker.WorkerPlugin workerPlugin = + (io.temporal.worker.WorkerPlugin) plugin; + chain = + () -> { + order.add(workerPlugin.getName() + "-startWorker-before"); + try { + workerPlugin.startWorker("test-queue", null, next); + } catch (Exception e) { + throw new RuntimeException(e); + } + order.add(workerPlugin.getName() + "-startWorker-after"); + }; + } + } + + chain.run(); + + // First plugin should wrap all others + assertEquals( + Arrays.asList( + "A-startWorker-before", + "B-startWorker-before", + "C-startWorker-before", + "worker-start", + "C-startWorker-after", + "B-startWorker-after", + "A-startWorker-after"), + order); + } + + @Test + public void testShutdownWorkerReverseOrder() { + List order = new ArrayList<>(); + + SimplePlugin pluginA = createWorkerLifecycleTrackingPlugin("A", order); + SimplePlugin pluginB = createWorkerLifecycleTrackingPlugin("B", order); + SimplePlugin pluginC = createWorkerLifecycleTrackingPlugin("C", order); + + List plugins = Arrays.asList(pluginA, pluginB, pluginC); + + // Build chain in reverse (like WorkerFactory does) + Runnable chain = () -> order.add("worker-shutdown"); + + List reversed = new ArrayList<>(plugins); + java.util.Collections.reverse(reversed); + for (Object plugin : reversed) { + if (plugin instanceof io.temporal.worker.WorkerPlugin) { + final Runnable next = chain; + final io.temporal.worker.WorkerPlugin workerPlugin = + (io.temporal.worker.WorkerPlugin) plugin; + chain = + () -> { + order.add(workerPlugin.getName() + "-shutdownWorker-before"); + workerPlugin.shutdownWorker("test-queue", null, next); + order.add(workerPlugin.getName() + "-shutdownWorker-after"); + }; + } + } + + chain.run(); + + // First plugin should wrap all others + assertEquals( + Arrays.asList( + "A-shutdownWorker-before", + "B-shutdownWorker-before", + "C-shutdownWorker-before", + "worker-shutdown", + "C-shutdownWorker-after", + "B-shutdownWorker-after", + "A-shutdownWorker-after"), + order); + } + + @Test + public void testSimplePluginImplementsBothInterfaces() { + SimplePlugin plugin = new SimplePlugin("dual-plugin") {}; + + assertTrue( + "SimplePlugin should implement ClientPlugin", + plugin instanceof io.temporal.client.ClientPlugin); + assertTrue( + "SimplePlugin should implement WorkerPlugin", + plugin instanceof io.temporal.worker.WorkerPlugin); + } + + private SimplePlugin createTrackingPlugin(String name, List order) { + return new SimplePlugin(name) { + @Override + public void configureClient(WorkflowClientOptions.Builder builder) { + order.add(name + "-config"); + } + }; + } + + private SimplePlugin createExecutionTrackingPlugin(String name, List order) { + return new SimplePlugin(name) { + @Override + public void startWorkerFactory(io.temporal.worker.WorkerFactory factory, Runnable next) { + next.run(); + } + }; + } + + private SimplePlugin createWorkerLifecycleTrackingPlugin(String name, List order) { + return new SimplePlugin(name) { + @Override + public void startWorker(String taskQueue, io.temporal.worker.Worker worker, Runnable next) { + next.run(); + } + + @Override + public void shutdownWorker( + String taskQueue, io.temporal.worker.Worker worker, Runnable next) { + next.run(); + } + }; + } +} diff --git a/temporal-sdk/src/test/java/io/temporal/common/SimplePluginBuilderTest.java b/temporal-sdk/src/test/java/io/temporal/common/SimplePluginBuilderTest.java new file mode 100644 index 000000000..b7e7c78b9 --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/common/SimplePluginBuilderTest.java @@ -0,0 +1,411 @@ +/* + * Copyright (C) 2022 Temporal Technologies, Inc. All Rights Reserved. + * + * Copyright (C) 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Modifications copyright (C) 2017 Uber Technologies, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this material except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.temporal.common; + +import static org.junit.Assert.*; + +import io.temporal.client.WorkflowClientOptions; +import io.temporal.common.interceptors.WorkerInterceptor; +import io.temporal.common.interceptors.WorkerInterceptorBase; +import io.temporal.common.interceptors.WorkflowClientInterceptor; +import io.temporal.common.interceptors.WorkflowClientInterceptorBase; +import io.temporal.serviceclient.WorkflowServiceStubsOptions; +import io.temporal.worker.WorkerFactoryOptions; +import io.temporal.worker.WorkerOptions; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; + +public class SimplePluginBuilderTest { + + @Test + public void testSimplePluginName() { + SimplePlugin plugin = SimplePlugin.newBuilder("test-plugin").build(); + assertEquals("test-plugin", plugin.getName()); + } + + @Test + public void testSimplePluginImplementsBothInterfaces() { + SimplePlugin plugin = SimplePlugin.newBuilder("test").build(); + assertTrue( + "Should implement io.temporal.client.ClientPlugin", + plugin instanceof io.temporal.client.ClientPlugin); + assertTrue( + "Should implement io.temporal.worker.WorkerPlugin", + plugin instanceof io.temporal.worker.WorkerPlugin); + } + + @Test + public void testCustomizeServiceStubs() { + AtomicBoolean customized = new AtomicBoolean(false); + + SimplePlugin plugin = + SimplePlugin.newBuilder("test") + .customizeServiceStubs( + builder -> { + customized.set(true); + }) + .build(); + + WorkflowServiceStubsOptions.Builder builder = WorkflowServiceStubsOptions.newBuilder(); + ((io.temporal.client.ClientPlugin) plugin).configureServiceStubs(builder); + + assertTrue("Customizer should have been called", customized.get()); + } + + @Test + public void testCustomizeClient() { + AtomicBoolean customized = new AtomicBoolean(false); + + SimplePlugin plugin = + SimplePlugin.newBuilder("test") + .customizeClient( + builder -> { + customized.set(true); + builder.setIdentity("custom-identity"); + }) + .build(); + + WorkflowClientOptions.Builder builder = WorkflowClientOptions.newBuilder(); + ((io.temporal.client.ClientPlugin) plugin).configureClient(builder); + + assertTrue("Customizer should have been called", customized.get()); + assertEquals("custom-identity", builder.build().getIdentity()); + } + + @Test + public void testCustomizeWorkerFactory() { + AtomicBoolean customized = new AtomicBoolean(false); + + SimplePlugin plugin = + SimplePlugin.newBuilder("test") + .customizeWorkerFactory( + builder -> { + customized.set(true); + builder.setWorkflowCacheSize(100); + }) + .build(); + + WorkerFactoryOptions.Builder builder = WorkerFactoryOptions.newBuilder(); + ((io.temporal.worker.WorkerPlugin) plugin).configureWorkerFactory(builder); + + assertTrue("Customizer should have been called", customized.get()); + assertEquals(100, builder.build().getWorkflowCacheSize()); + } + + @Test + public void testCustomizeWorker() { + AtomicBoolean customized = new AtomicBoolean(false); + + SimplePlugin plugin = + SimplePlugin.newBuilder("test") + .customizeWorker( + builder -> { + customized.set(true); + builder.setMaxConcurrentActivityExecutionSize(50); + }) + .build(); + + WorkerOptions.Builder builder = WorkerOptions.newBuilder(); + ((io.temporal.worker.WorkerPlugin) plugin).configureWorker("test-queue", builder); + + assertTrue("Customizer should have been called", customized.get()); + assertEquals(50, builder.build().getMaxConcurrentActivityExecutionSize()); + } + + @Test + public void testMultipleCustomizers() { + AtomicInteger callCount = new AtomicInteger(0); + + SimplePlugin plugin = + SimplePlugin.newBuilder("test") + .customizeClient(builder -> callCount.incrementAndGet()) + .customizeClient(builder -> callCount.incrementAndGet()) + .customizeClient(builder -> callCount.incrementAndGet()) + .build(); + + WorkflowClientOptions.Builder builder = WorkflowClientOptions.newBuilder(); + ((io.temporal.client.ClientPlugin) plugin).configureClient(builder); + + assertEquals("All customizers should be called", 3, callCount.get()); + } + + @Test + public void testAddWorkerInterceptors() { + WorkerInterceptor interceptor = new WorkerInterceptorBase() {}; + + SimplePlugin plugin = + SimplePlugin.newBuilder("test").addWorkerInterceptors(interceptor).build(); + + WorkerFactoryOptions.Builder builder = WorkerFactoryOptions.newBuilder(); + ((io.temporal.worker.WorkerPlugin) plugin).configureWorkerFactory(builder); + + WorkerInterceptor[] interceptors = builder.build().getWorkerInterceptors(); + assertEquals(1, interceptors.length); + assertSame(interceptor, interceptors[0]); + } + + @Test + public void testAddClientInterceptors() { + WorkflowClientInterceptor interceptor = new WorkflowClientInterceptorBase() {}; + + SimplePlugin plugin = + SimplePlugin.newBuilder("test").addClientInterceptors(interceptor).build(); + + WorkflowClientOptions.Builder builder = WorkflowClientOptions.newBuilder(); + ((io.temporal.client.ClientPlugin) plugin).configureClient(builder); + + WorkflowClientInterceptor[] interceptors = builder.build().getInterceptors(); + assertEquals(1, interceptors.length); + assertSame(interceptor, interceptors[0]); + } + + @Test + public void testInterceptorsAppendToExisting() { + WorkerInterceptor existingInterceptor = new WorkerInterceptorBase() {}; + WorkerInterceptor newInterceptor = new WorkerInterceptorBase() {}; + + SimplePlugin plugin = + SimplePlugin.newBuilder("test").addWorkerInterceptors(newInterceptor).build(); + + WorkerFactoryOptions.Builder builder = + WorkerFactoryOptions.newBuilder().setWorkerInterceptors(existingInterceptor); + ((io.temporal.worker.WorkerPlugin) plugin).configureWorkerFactory(builder); + + WorkerInterceptor[] interceptors = builder.build().getWorkerInterceptors(); + assertEquals(2, interceptors.length); + assertSame(existingInterceptor, interceptors[0]); + assertSame(newInterceptor, interceptors[1]); + } + + @Test + public void testInitializeWorker() { + AtomicBoolean initialized = new AtomicBoolean(false); + String[] capturedTaskQueue = {null}; + + SimplePlugin plugin = + SimplePlugin.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) + ((io.temporal.worker.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); + + SimplePlugin plugin = + SimplePlugin.newBuilder("test") + .initializeWorker((taskQueue, worker) -> callCount.incrementAndGet()) + .initializeWorker((taskQueue, worker) -> callCount.incrementAndGet()) + .initializeWorker((taskQueue, worker) -> callCount.incrementAndGet()) + .build(); + + ((io.temporal.worker.WorkerPlugin) plugin).initializeWorker("test-queue", null); + + assertEquals("All initializers should be called", 3, callCount.get()); + } + + @Test(expected = NullPointerException.class) + public void testNullInitializeWorker() { + SimplePlugin.newBuilder("test").initializeWorker(null); + } + + @Test + public void testOnWorkerStart() throws Exception { + AtomicBoolean started = new AtomicBoolean(false); + String[] capturedTaskQueue = {null}; + + SimplePlugin plugin = + SimplePlugin.newBuilder("test") + .onWorkerStart( + (taskQueue, worker) -> { + started.set(true); + capturedTaskQueue[0] = taskQueue; + }) + .build(); + + AtomicBoolean nextCalled = new AtomicBoolean(false); + ((io.temporal.worker.WorkerPlugin) plugin) + .startWorker("my-task-queue", null, () -> nextCalled.set(true)); + + assertTrue("next should be called", nextCalled.get()); + assertTrue("Callback should have been called", started.get()); + assertEquals("my-task-queue", capturedTaskQueue[0]); + } + + @Test + public void testOnWorkerShutdown() { + AtomicBoolean shutdown = new AtomicBoolean(false); + String[] capturedTaskQueue = {null}; + + SimplePlugin plugin = + SimplePlugin.newBuilder("test") + .onWorkerShutdown( + (taskQueue, worker) -> { + shutdown.set(true); + capturedTaskQueue[0] = taskQueue; + }) + .build(); + + AtomicBoolean nextCalled = new AtomicBoolean(false); + ((io.temporal.worker.WorkerPlugin) plugin) + .shutdownWorker("my-task-queue", null, () -> nextCalled.set(true)); + + assertTrue("next should be called", nextCalled.get()); + assertTrue("Callback should have been called", shutdown.get()); + assertEquals("my-task-queue", capturedTaskQueue[0]); + } + + @Test + public void testMultipleOnWorkerStartCallbacks() throws Exception { + AtomicInteger callCount = new AtomicInteger(0); + + SimplePlugin plugin = + SimplePlugin.newBuilder("test") + .onWorkerStart((taskQueue, worker) -> callCount.incrementAndGet()) + .onWorkerStart((taskQueue, worker) -> callCount.incrementAndGet()) + .onWorkerStart((taskQueue, worker) -> callCount.incrementAndGet()) + .build(); + + ((io.temporal.worker.WorkerPlugin) plugin).startWorker("test-queue", null, () -> {}); + + assertEquals("All callbacks should be called", 3, callCount.get()); + } + + @Test + public void testMultipleOnWorkerShutdownCallbacks() { + AtomicInteger callCount = new AtomicInteger(0); + + SimplePlugin plugin = + SimplePlugin.newBuilder("test") + .onWorkerShutdown((taskQueue, worker) -> callCount.incrementAndGet()) + .onWorkerShutdown((taskQueue, worker) -> callCount.incrementAndGet()) + .onWorkerShutdown((taskQueue, worker) -> callCount.incrementAndGet()) + .build(); + + ((io.temporal.worker.WorkerPlugin) plugin).shutdownWorker("test-queue", null, () -> {}); + + assertEquals("All callbacks should be called", 3, callCount.get()); + } + + @Test + public void testOnWorkerFactoryStart() throws Exception { + AtomicBoolean started = new AtomicBoolean(false); + + SimplePlugin plugin = + SimplePlugin.newBuilder("test").onWorkerFactoryStart(factory -> started.set(true)).build(); + + AtomicBoolean nextCalled = new AtomicBoolean(false); + ((io.temporal.worker.WorkerPlugin) plugin).startWorkerFactory(null, () -> nextCalled.set(true)); + + assertTrue("next should be called", nextCalled.get()); + assertTrue("Callback should have been called", started.get()); + } + + @Test + public void testOnWorkerFactoryShutdown() throws Exception { + AtomicBoolean shutdown = new AtomicBoolean(false); + + SimplePlugin plugin = + SimplePlugin.newBuilder("test") + .onWorkerFactoryShutdown(factory -> shutdown.set(true)) + .build(); + + AtomicBoolean nextCalled = new AtomicBoolean(false); + ((io.temporal.worker.WorkerPlugin) plugin) + .shutdownWorkerFactory(null, () -> nextCalled.set(true)); + + assertTrue("next should be called", nextCalled.get()); + assertTrue("Callback should have been called", shutdown.get()); + } + + @Test + public void testMultipleOnWorkerFactoryStartCallbacks() throws Exception { + AtomicInteger callCount = new AtomicInteger(0); + + SimplePlugin plugin = + SimplePlugin.newBuilder("test") + .onWorkerFactoryStart(factory -> callCount.incrementAndGet()) + .onWorkerFactoryStart(factory -> callCount.incrementAndGet()) + .onWorkerFactoryStart(factory -> callCount.incrementAndGet()) + .build(); + + ((io.temporal.worker.WorkerPlugin) plugin).startWorkerFactory(null, () -> {}); + + assertEquals("All callbacks should be called", 3, callCount.get()); + } + + @Test + public void testMultipleOnWorkerFactoryShutdownCallbacks() throws Exception { + AtomicInteger callCount = new AtomicInteger(0); + + SimplePlugin plugin = + SimplePlugin.newBuilder("test") + .onWorkerFactoryShutdown(factory -> callCount.incrementAndGet()) + .onWorkerFactoryShutdown(factory -> callCount.incrementAndGet()) + .onWorkerFactoryShutdown(factory -> callCount.incrementAndGet()) + .build(); + + ((io.temporal.worker.WorkerPlugin) plugin).shutdownWorkerFactory(null, () -> {}); + + assertEquals("All callbacks should be called", 3, callCount.get()); + } + + @Test(expected = NullPointerException.class) + public void testNullOnWorkerStart() { + SimplePlugin.newBuilder("test").onWorkerStart(null); + } + + @Test(expected = NullPointerException.class) + public void testNullOnWorkerShutdown() { + SimplePlugin.newBuilder("test").onWorkerShutdown(null); + } + + @Test(expected = NullPointerException.class) + public void testNullOnWorkerFactoryStart() { + SimplePlugin.newBuilder("test").onWorkerFactoryStart(null); + } + + @Test(expected = NullPointerException.class) + public void testNullOnWorkerFactoryShutdown() { + SimplePlugin.newBuilder("test").onWorkerFactoryShutdown(null); + } + + @Test(expected = NullPointerException.class) + public void testNullName() { + SimplePlugin.newBuilder(null); + } + + @Test(expected = NullPointerException.class) + public void testNullCustomizer() { + SimplePlugin.newBuilder("test").customizeClient(null); + } +} diff --git a/temporal-serviceclient/src/main/java/io/temporal/serviceclient/WorkflowServiceStubs.java b/temporal-serviceclient/src/main/java/io/temporal/serviceclient/WorkflowServiceStubs.java index 2ce76512a..7ae3077c7 100644 --- a/temporal-serviceclient/src/main/java/io/temporal/serviceclient/WorkflowServiceStubs.java +++ b/temporal-serviceclient/src/main/java/io/temporal/serviceclient/WorkflowServiceStubs.java @@ -6,6 +6,8 @@ import io.temporal.internal.WorkflowThreadMarker; import io.temporal.internal.testservice.InProcessGRPCServer; import java.time.Duration; +import java.util.function.Supplier; +import javax.annotation.Nonnull; import javax.annotation.Nullable; /** Initializes and holds gRPC blocking and future stubs. */ @@ -122,5 +124,74 @@ static WorkflowServiceStubs newInstance( new WorkflowServiceStubsImpl(service, options), WorkflowServiceStubs.class); } + /** + * Creates WorkflowService gRPC stubs with plugin support. + * + *

This method applies plugins in two phases: + * + *

    + *
  1. Configuration phase: Each plugin's {@code configureServiceStubs} method is called + * in forward (registration) order to modify the options builder + *
  2. Connection phase: Each plugin's {@code connectServiceClient} method is called in + * reverse order to wrap the connection (first plugin wraps all others) + *
+ * + *

This method creates stubs with "lazy" connectivity. The connection is not performed during + * the creation time and happens on the first request. + * + * @param options stub options to use + * @param plugins array of client plugins to apply + * @return the workflow service stubs + */ + static WorkflowServiceStubs newServiceStubs( + @Nonnull WorkflowServiceStubsOptions options, @Nonnull ClientPluginCallback[] plugins) { + enforceNonWorkflowThread(); + + // Apply plugin configuration phase (forward order) + WorkflowServiceStubsOptions.Builder builder = WorkflowServiceStubsOptions.newBuilder(options); + for (ClientPluginCallback plugin : plugins) { + plugin.configureServiceStubs(builder); + } + WorkflowServiceStubsOptions finalOptions = builder.validateAndBuildWithDefaults(); + + // Build connection chain (reverse order for proper nesting) + Supplier connectionChain = + () -> + WorkflowThreadMarker.protectFromWorkflowThread( + new WorkflowServiceStubsImpl(null, finalOptions), WorkflowServiceStubs.class); + + for (int i = plugins.length - 1; i >= 0; i--) { + final Supplier next = connectionChain; + final ClientPluginCallback callback = plugins[i]; + connectionChain = () -> callback.connectServiceClient(finalOptions, next); + } + + return connectionChain.get(); + } + + /** + * Callback interface for client plugins to participate in service stubs creation. This interface + * is implemented by {@code io.temporal.client.ClientPlugin} in the temporal-sdk module. + */ + interface ClientPluginCallback { + /** + * Allows the plugin to modify service stubs options before the service stubs are created. + * + * @param builder the options builder to modify + */ + void configureServiceStubs(@Nonnull WorkflowServiceStubsOptions.Builder builder); + + /** + * Allows the plugin to wrap service client connection. + * + * @param options the final options being used for connection + * @param next supplier that creates the service stubs (calls next plugin or actual connection) + * @return the service stubs (possibly wrapped or decorated) + */ + @Nonnull + WorkflowServiceStubs connectServiceClient( + @Nonnull WorkflowServiceStubsOptions options, @Nonnull Supplier next); + } + WorkflowServiceStubsOptions getOptions(); }