Add per-invocation Lambda worker configuration

Add LambdaWorker overloads that accept an invocation configurator so
handlers can adjust worker options using the AWS Lambda Context before
Temporal service stubs, clients, and workers are created.

Copy the base options for each invocation, allow required fields such as
task queue to be supplied dynamically, and keep invocation-local shutdown
hooks from leaking across warm invocations. Run those hooks if the
per-invocation configurator or final option validation fails.

Document the new overload and cover the lifecycle behavior with unit tests.

Author: eamsden

contrib/temporal-aws-lambda/README.md

@@ -33,6 +33,21 @@ public final class Handler implements RequestHandler<Object, Void> {
 
 Connection options are loaded with `temporal-envconfig` when the handler is constructed during Lambda cold start. The Lambda worker checks `TEMPORAL_CONFIG_FILE` first, then readable `$LAMBDA_TASK_ROOT/temporal.toml`, then readable `./temporal.toml`, then falls back to the envconfig defaults and Temporal environment variables. The `configure` callback also runs during handler construction, so non-invocation configuration is prepared once and reused.
 
+Use the per-invocation configuration overload when final options depend on the Lambda `Context` or on resources opened for one invocation. The cold-start callback still runs once; the invocation callback runs before Temporal service stubs, client, and worker are created:
+
+```java
+private static final RequestHandler<Object, Void> WORKER =
+    LambdaWorker.run(
+        new WorkerDeploymentVersion("orders-worker", "2026-06-02"),
+        builder -> builder.registerWorkflowImplementationTypes(OrderWorkflowImpl.class),
+        (builder, context) -> {
+          builder.setTaskQueue(taskQueueFor(context));
+          builder.addShutdownHook(() -> cleanupInvocationResources(context));
+        });
+```
+
+Each invocation receives a fresh copy of the base options, so mutations and shutdown hooks added by the invocation callback do not leak across warm invocations. If the invocation callback throws after adding shutdown hooks, those hooks still run for cleanup. Reusable resources such as AWS SDK clients should usually be created during cold start and reused across warm invocations.
+
 If you need to assemble options outside the `run` callback, call `LambdaWorkerOptions.newBuilderFromEnvironment()`, configure the returned builder, call `build()`, and pass the options to `LambdaWorker.newHandler(...)`.
 
 Dynamic workflow and activity implementations can be registered with `registerDynamicWorkflowImplementationType(...)` and `registerDynamicActivityImplementation(...)`. Java SDK worker rules still apply: only one dynamic workflow implementation type and one dynamic activity implementation can be registered per worker.

contrib/temporal-aws-lambda/src/main/java/io/temporal/aws/lambda/LambdaWorker.java

@@ -5,6 +5,7 @@
 import io.temporal.common.WorkerDeploymentVersion;
 import java.io.IOException;
 import java.time.Duration;
+import java.util.List;
 import java.util.Objects;
 import java.util.function.Consumer;
 import org.slf4j.Logger;
@@ -19,6 +20,19 @@ public final class LambdaWorker {
 
   private LambdaWorker() {}
 
+  /**
+   * Configures options for one Lambda invocation before Temporal service stubs, client, and worker
+   * are created.
+   *
+   * <p>The supplied builder is a fresh copy of the base options for the current invocation.
+   * Shutdown hooks added by this callback run for only that invocation, including when this
+   * callback throws.
+   */
+  @FunctionalInterface
+  public interface InvocationConfigurator {
+    void configure(LambdaWorkerOptions.Builder builder, Context context);
+  }
+
   /**
    * Returns an AWS Lambda Java handler that creates, starts, and shuts down one Temporal worker per
    * invocation.
@@ -40,13 +54,57 @@ public static RequestHandler<Object, Void> run(
     }
   }
 
+  /**
+   * Returns an AWS Lambda Java handler with both cold-start and per-invocation configuration.
+   *
+   * @param version worker deployment version to advertise for this worker.
+   * @param configure callback invoked once while the Lambda handler is constructed.
+   * @param invocationConfigure callback invoked for each Lambda invocation before Temporal service
+   *     stubs, client, and worker are created. Required fields may be supplied by this callback.
+   */
+  public static RequestHandler<Object, Void> run(
+      WorkerDeploymentVersion version,
+      Consumer<LambdaWorkerOptions.Builder> configure,
+      InvocationConfigurator invocationConfigure) {
+    LambdaWorkerOptions.validateVersion(version);
+    Objects.requireNonNull(configure, "configure");
+    Objects.requireNonNull(invocationConfigure, "invocationConfigure");
+    try {
+      LambdaWorkerOptions.Builder builder =
+          LambdaWorkerOptions.newBuilderFromEnvironment(System.getenv());
+      configure.accept(builder);
+      return newHandler(version, builder.build(), invocationConfigure);
+    } catch (IOException e) {
+      throw new RuntimeException("Unable to load Temporal client configuration", e);
+    }
+  }
+
   /** Returns an AWS Lambda Java handler using already-configured Lambda worker options. */
   public static RequestHandler<Object, Void> newHandler(
       WorkerDeploymentVersion version, LambdaWorkerOptions options) {
     return newHandler(
         version, options, new DefaultLambdaWorkerRuntime(), sleep(), systemNanoClock());
   }
 
+  /**
+   * Returns an AWS Lambda Java handler using base options and per-invocation configuration.
+   *
+   * <p>The supplied options are copied for each invocation before {@code invocationConfigure} runs.
+   * Required fields may be supplied by {@code invocationConfigure}.
+   */
+  public static RequestHandler<Object, Void> newHandler(
+      WorkerDeploymentVersion version,
+      LambdaWorkerOptions options,
+      InvocationConfigurator invocationConfigure) {
+    return newHandler(
+        version,
+        options,
+        invocationConfigure,
+        new DefaultLambdaWorkerRuntime(),
+        sleep(),
+        systemNanoClock());
+  }
+
   static RequestHandler<Object, Void> newHandler(
       WorkerDeploymentVersion version,
       LambdaWorkerOptions options,
@@ -56,6 +114,16 @@ static RequestHandler<Object, Void> newHandler(
     return newHandler(version, options, runtime, sleeper, systemNanoClock());
   }
 
+  static RequestHandler<Object, Void> newHandler(
+      WorkerDeploymentVersion version,
+      LambdaWorkerOptions options,
+      InvocationConfigurator invocationConfigure,
+      LambdaWorkerRuntime runtime,
+      Sleeper sleeper) {
+    // This overload exists to let tests inject a fake runtime and sleeper.
+    return newHandler(version, options, invocationConfigure, runtime, sleeper, systemNanoClock());
+  }
+
   static RequestHandler<Object, Void> newHandler(
       WorkerDeploymentVersion version,
       LambdaWorkerOptions options,
@@ -67,7 +135,34 @@ static RequestHandler<Object, Void> newHandler(
         Objects.requireNonNull(options, "options").prepare(version),
         Objects.requireNonNull(runtime, "runtime"),
         Objects.requireNonNull(sleeper, "sleeper"),
-        Objects.requireNonNull(clock, "clock"));
+        Objects.requireNonNull(clock, "clock"),
+        false);
+  }
+
+  static RequestHandler<Object, Void> newHandler(
+      WorkerDeploymentVersion version,
+      LambdaWorkerOptions options,
+      InvocationConfigurator invocationConfigure,
+      LambdaWorkerRuntime runtime,
+      Sleeper sleeper,
+      NanoClock clock) {
+    LambdaWorkerOptions.validateVersion(version);
+    Objects.requireNonNull(options, "options");
+    Objects.requireNonNull(invocationConfigure, "invocationConfigure");
+    return new Handler(
+        context -> {
+          LambdaWorkerOptions.Builder builder = options.toBuilder();
+          try {
+            invocationConfigure.configure(builder, context);
+            return builder.build().prepare(version).materialize(identityFor(context));
+          } catch (RuntimeException e) {
+            throw new InvocationConfigurationException(e, builder);
+          }
+        },
+        Objects.requireNonNull(runtime, "runtime"),
+        Objects.requireNonNull(sleeper, "sleeper"),
+        Objects.requireNonNull(clock, "clock"),
+        true);
   }
 
   private static Sleeper sleep() {
@@ -86,32 +181,73 @@ private static NanoClock systemNanoClock() {
     return System::nanoTime;
   }
 
+  private static String identityFor(Context context) {
+    return emptyToUnknown(context.getAwsRequestId())
+        + "@"
+        + emptyToUnknown(context.getInvokedFunctionArn());
+  }
+
+  private static String emptyToUnknown(String value) {
+    return value == null || value.isEmpty() ? "unknown" : value;
+  }
+
+  private interface OptionsMaterializer {
+    LambdaWorkerOptions.Materialized materialize(Context context);
+  }
+
   private static final class Handler implements RequestHandler<Object, Void> {
-    private final LambdaWorkerOptions.Prepared preparedOptions;
+    private final OptionsMaterializer optionsMaterializer;
     private final LambdaWorkerRuntime runtime;
     private final Sleeper sleeper;
     private final NanoClock clock;
+    private final boolean runShutdownHooksBeforeRuntimeCreation;
 
     private Handler(
         LambdaWorkerOptions.Prepared preparedOptions,
         LambdaWorkerRuntime runtime,
         Sleeper sleeper,
+        NanoClock clock,
+        boolean runShutdownHooksBeforeRuntimeCreation) {
+      this(
+          context -> preparedOptions.materialize(identityFor(context)),
+          runtime,
+          sleeper,
+          clock,
+          runShutdownHooksBeforeRuntimeCreation);
+      Objects.requireNonNull(preparedOptions, "preparedOptions");
+    }
+
+    private Handler(
+        OptionsMaterializer optionsMaterializer,
+        LambdaWorkerRuntime runtime,
+        Sleeper sleeper,
         NanoClock clock) {
-      this.preparedOptions = Objects.requireNonNull(preparedOptions, "preparedOptions");
+      this(optionsMaterializer, runtime, sleeper, clock, false);
+    }
+
+    private Handler(
+        OptionsMaterializer optionsMaterializer,
+        LambdaWorkerRuntime runtime,
+        Sleeper sleeper,
+        NanoClock clock,
+        boolean runShutdownHooksBeforeRuntimeCreation) {
+      this.optionsMaterializer = Objects.requireNonNull(optionsMaterializer, "optionsMaterializer");
       this.runtime = runtime;
       this.sleeper = sleeper;
       this.clock = clock;
+      this.runShutdownHooksBeforeRuntimeCreation = runShutdownHooksBeforeRuntimeCreation;
     }
 
     @Override
     public Void handleRequest(Object input, Context context) {
       Objects.requireNonNull(context, "context");
 
-      LambdaWorkerOptions.Materialized options = preparedOptions.materialize(identityFor(context));
-      validateRemainingTime(context, options.shutdownDeadlineBuffer);
-
+      LambdaWorkerOptions.Materialized options = null;
       LambdaWorkerRuntime.Invocation invocation = null;
       try {
+        options = optionsMaterializer.materialize(context);
+        validateRemainingTime(context, options.shutdownDeadlineBuffer);
+
         invocation =
             runtime.create(
                 options.serviceStubsOptions,
@@ -134,8 +270,26 @@ public Void handleRequest(Object input, Context context) {
 
         sleepUntilShutdownWindow(context, options);
         return null;
+      } catch (InvocationConfigurationException e) {
+        runShutdownHooks(
+            context,
+            e.taskQueue,
+            e.shutdownHooks,
+            cleanupDeadlineNanos(context, e.gracefulShutdownTimeout, e.shutdownDeadlineBuffer));
+        throw e.failure;
+      } catch (RuntimeException e) {
+        if (invocation == null && options != null && runShutdownHooksBeforeRuntimeCreation) {
+          runShutdownHooks(
+              context,
+              options.taskQueue,
+              options.shutdownHooks,
+              cleanupDeadlineNanos(context, options));
+        }
+        throw e;
       } finally {
-        shutdownInvocation(context, invocation, options);
+        if (invocation != null) {
+          shutdownInvocation(context, invocation, options);
+        }
       }
     }
 
@@ -191,17 +345,15 @@ private void shutdownInvocation(
       }
       runShutdownHooks(context, options, cleanupDeadlineNanos.longValue());
 
-      if (invocation != null) {
-        try {
-          invocation.closeStubs(remainingCleanupTime(cleanupDeadlineNanos.longValue()));
-        } catch (RuntimeException e) {
-          log.error(
-              "Temporal Lambda worker service stubs close failed awsRequestId={} invokedFunctionArn={} taskQueue={}",
-              context.getAwsRequestId(),
-              context.getInvokedFunctionArn(),
-              options.taskQueue,
-              e);
-        }
+      try {
+        invocation.closeStubs(remainingCleanupTime(cleanupDeadlineNanos.longValue()));
+      } catch (RuntimeException e) {
+        log.error(
+            "Temporal Lambda worker service stubs close failed awsRequestId={} invokedFunctionArn={} taskQueue={}",
+            context.getAwsRequestId(),
+            context.getInvokedFunctionArn(),
+            options.taskQueue,
+            e);
       }
     }
 
@@ -267,7 +419,15 @@ private boolean isTerminated(
 
     private void runShutdownHooks(
         Context context, LambdaWorkerOptions.Materialized options, long cleanupDeadlineNanos) {
-      for (Runnable hook : options.shutdownHooks) {
+      runShutdownHooks(context, options.taskQueue, options.shutdownHooks, cleanupDeadlineNanos);
+    }
+
+    private void runShutdownHooks(
+        Context context,
+        String taskQueue,
+        List<Runnable> shutdownHooks,
+        long cleanupDeadlineNanos) {
+      for (Runnable hook : shutdownHooks) {
         try {
           if (hook instanceof TimedShutdownHook) {
             ((TimedShutdownHook) hook).run(remainingCleanupTime(cleanupDeadlineNanos));
@@ -279,7 +439,7 @@ private void runShutdownHooks(
               "Temporal Lambda worker shutdown hook failed awsRequestId={} invokedFunctionArn={} taskQueue={}",
               context.getAwsRequestId(),
               context.getInvokedFunctionArn(),
-              options.taskQueue,
+              taskQueue,
               e);
         }
       }
@@ -311,12 +471,25 @@ private void validateRemainingTime(Context context, Duration shutdownDeadlineBuf
     }
 
     private long cleanupDeadlineNanos(Context context, LambdaWorkerOptions.Materialized options) {
-      return clock.nanoTime() + cleanupWindow(context, options).toNanos();
+      return cleanupDeadlineNanos(
+          context, options.gracefulShutdownTimeout, options.shutdownDeadlineBuffer);
+    }
+
+    private long cleanupDeadlineNanos(
+        Context context, Duration gracefulShutdownTimeout, Duration shutdownDeadlineBuffer) {
+      return clock.nanoTime()
+          + cleanupWindow(context, gracefulShutdownTimeout, shutdownDeadlineBuffer).toNanos();
     }
 
     private Duration cleanupWindow(Context context, LambdaWorkerOptions.Materialized options) {
+      return cleanupWindow(
+          context, options.gracefulShutdownTimeout, options.shutdownDeadlineBuffer);
+    }
+
+    private Duration cleanupWindow(
+        Context context, Duration gracefulShutdownTimeout, Duration shutdownDeadlineBuffer) {
       Duration configuredWindow =
-          nonNegative(options.shutdownDeadlineBuffer.minus(options.gracefulShutdownTimeout));
+          nonNegative(shutdownDeadlineBuffer.minus(gracefulShutdownTimeout));
       Duration remaining = remainingInvocationTime(context);
       return configuredWindow.compareTo(remaining) <= 0 ? configuredWindow : remaining;
     }
@@ -332,15 +505,23 @@ private Duration remainingInvocationTime(Context context) {
     private static Duration nonNegative(Duration duration) {
       return duration.isNegative() ? Duration.ZERO : duration;
     }
+  }
 
-    private static String identityFor(Context context) {
-      return emptyToUnknown(context.getAwsRequestId())
-          + "@"
-          + emptyToUnknown(context.getInvokedFunctionArn());
-    }
-
-    private static String emptyToUnknown(String value) {
-      return value == null || value.isEmpty() ? "unknown" : value;
+  private static final class InvocationConfigurationException extends RuntimeException {
+    private final RuntimeException failure;
+    private final String taskQueue;
+    private final Duration gracefulShutdownTimeout;
+    private final Duration shutdownDeadlineBuffer;
+    private final List<Runnable> shutdownHooks;
+
+    private InvocationConfigurationException(
+        RuntimeException failure, LambdaWorkerOptions.Builder builder) {
+      super(failure);
+      this.failure = failure;
+      this.taskQueue = builder.getTaskQueue();
+      this.gracefulShutdownTimeout = builder.getGracefulShutdownTimeout();
+      this.shutdownDeadlineBuffer = builder.getShutdownDeadlineBuffer();
+      this.shutdownHooks = builder.getShutdownHooks();
     }
   }
 }