Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ public static final class Builder {
private boolean usingVirtualThreads;
private WorkerDeploymentOptions deploymentOptions;
private String workerInstanceKey;
private boolean allowActivityHeartbeatDuringShutdown;

private Builder() {}

Expand All @@ -66,6 +67,7 @@ private Builder(SingleWorkerOptions options) {
this.usingVirtualThreads = options.isUsingVirtualThreads();
this.deploymentOptions = options.getDeploymentOptions();
this.workerInstanceKey = options.getWorkerInstanceKey();
this.allowActivityHeartbeatDuringShutdown = options.getAllowActivityHeartbeatDuringShutdown();
}

public Builder setIdentity(String identity) {
Expand Down Expand Up @@ -162,6 +164,12 @@ public Builder setWorkerInstanceKey(String workerInstanceKey) {
return this;
}

public Builder setAllowActivityHeartbeatDuringShutdown(
boolean allowActivityHeartbeatDuringShutdown) {
this.allowActivityHeartbeatDuringShutdown = allowActivityHeartbeatDuringShutdown;
return this;
}

public SingleWorkerOptions build() {
PollerOptions pollerOptions = this.pollerOptions;
if (pollerOptions == null) {
Expand Down Expand Up @@ -201,7 +209,8 @@ public SingleWorkerOptions build() {
drainStickyTaskQueueTimeout,
usingVirtualThreads,
this.deploymentOptions,
this.workerInstanceKey);
this.workerInstanceKey,
this.allowActivityHeartbeatDuringShutdown);
}
}

Expand All @@ -223,6 +232,7 @@ public SingleWorkerOptions build() {
private final boolean usingVirtualThreads;
private final WorkerDeploymentOptions deploymentOptions;
private final String workerInstanceKey;
private final boolean allowActivityHeartbeatDuringShutdown;

private SingleWorkerOptions(
String identity,
Expand All @@ -242,7 +252,8 @@ private SingleWorkerOptions(
Duration drainStickyTaskQueueTimeout,
boolean usingVirtualThreads,
WorkerDeploymentOptions deploymentOptions,
String workerInstanceKey) {
String workerInstanceKey,
boolean allowActivityHeartbeatDuringShutdown) {
this.identity = identity;
this.binaryChecksum = binaryChecksum;
this.buildId = buildId;
Expand All @@ -261,6 +272,7 @@ private SingleWorkerOptions(
this.usingVirtualThreads = usingVirtualThreads;
this.deploymentOptions = deploymentOptions;
this.workerInstanceKey = workerInstanceKey;
this.allowActivityHeartbeatDuringShutdown = allowActivityHeartbeatDuringShutdown;
}

public String getIdentity() {
Expand Down Expand Up @@ -291,6 +303,10 @@ public Duration getDrainStickyTaskQueueTimeout() {
return drainStickyTaskQueueTimeout;
}

public boolean getAllowActivityHeartbeatDuringShutdown() {
return allowActivityHeartbeatDuringShutdown;
}

public DataConverter getDataConverter() {
return dataConverter;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ public class SyncActivityWorker implements SuspendableWorker {
private final ScheduledExecutorService heartbeatExecutor;
private final ActivityTaskHandlerImpl taskHandler;
private final ActivityWorker worker;
private final boolean allowActivityHeartbeatDuringShutdown;

public SyncActivityWorker(
WorkflowClient client,
Expand All @@ -38,6 +39,7 @@ public SyncActivityWorker(
this.identity = options.getIdentity();
this.namespace = namespace;
this.taskQueue = taskQueue;
this.allowActivityHeartbeatDuringShutdown = options.getAllowActivityHeartbeatDuringShutdown();

this.heartbeatExecutor =
Executors.newScheduledThreadPool(
Expand Down Expand Up @@ -89,16 +91,31 @@ public boolean start() {

@Override
public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean interruptTasks) {
return shutdownManager
// we want to shut down heartbeatExecutor before activity worker, so in-flight activities
// could get an ActivityWorkerShutdownException from their heartbeat
.shutdownExecutor(heartbeatExecutor, this + "#heartbeatExecutor", Duration.ofSeconds(5))
.thenCompose(r -> worker.shutdown(shutdownManager, interruptTasks))
.exceptionally(
e -> {
log.error("[BUG] Unexpected exception during shutdown", e);
return null;
});
CompletableFuture<Void> shutdownFuture;
if (allowActivityHeartbeatDuringShutdown && !interruptTasks) {
// we want to shut down heartbeatExecutor only after all outstanding activity tasks have
// finished executing, so in-flight activities can keep heartbeating during the shutdown
shutdownFuture =
worker
.shutdown(shutdownManager, interruptTasks)
.thenCompose(r -> shutdownHeartbeatExecutor(shutdownManager));
} else {
// we want to shut down heartbeatExecutor before activity worker, so in-flight activities
// could get an ActivityWorkerShutdownException from their heartbeat
shutdownFuture =
shutdownHeartbeatExecutor(shutdownManager)
.thenCompose(r -> worker.shutdown(shutdownManager, interruptTasks));
}
return shutdownFuture.exceptionally(
e -> {
log.error("[BUG] Unexpected exception during shutdown", e);
return null;
});
}

private CompletableFuture<Void> shutdownHeartbeatExecutor(ShutdownManager shutdownManager) {
return shutdownManager.shutdownExecutor(
heartbeatExecutor, this + "#heartbeatExecutor", Duration.ofSeconds(5));
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -893,6 +893,7 @@ private static SingleWorkerOptions toActivityOptions(
return toSingleWorkerOptions(
factoryOptions, options, clientOptions, contextPropagators, workerInstanceKey)
.setUsingVirtualThreads(options.isUsingVirtualThreadsOnActivityWorker())
.setAllowActivityHeartbeatDuringShutdown(options.getAllowActivityHeartbeatDuringShutdown())
.setPollerOptions(
PollerOptions.newBuilder()
.setMaximumPollRatePerSecond(options.getMaxWorkerActivitiesPerSecond())
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ public static final class Builder {
private PollerBehavior workflowTaskPollersBehavior;
private PollerBehavior activityTaskPollersBehavior;
private PollerBehavior nexusTaskPollersBehavior;
private boolean allowActivityHeartbeatDuringShutdown;

private Builder() {}

Expand Down Expand Up @@ -112,6 +113,7 @@ private Builder(WorkerOptions o) {
this.workflowTaskPollersBehavior = o.workflowTaskPollersBehavior;
this.activityTaskPollersBehavior = o.activityTaskPollersBehavior;
this.nexusTaskPollersBehavior = o.nexusTaskPollersBehavior;
this.allowActivityHeartbeatDuringShutdown = o.allowActivityHeartbeatDuringShutdown;
}

/**
Expand Down Expand Up @@ -524,6 +526,28 @@ public Builder setNexusTaskPollersBehavior(PollerBehavior pollerBehavior) {
return this;
}

/**
* If true, activities can keep heartbeating during graceful worker shutdown (see {@link
* io.temporal.worker.WorkerFactory#shutdown WorkerFactory.shutdown}). Defaults to false, which
* means that after graceful shutdown is requested, calling {@link
* io.temporal.activity.ActivityExecutionContext#heartbeat ActivityExecutionContext.heartbeat}
* does not send a heartbeat and instead throws {@link
* io.temporal.client.ActivityWorkerShutdownException ActivityWorkerShutdownException}. This
* option is ignored by non-graceful shutdown (see {@link
* io.temporal.worker.WorkerFactory#shutdownNow WorkerFactory.shutdownNow}).
*
* <p>Note that with this option enabled, activities are no longer notified of the worker
* shutdown by the {@link io.temporal.client.ActivityWorkerShutdownException
* ActivityWorkerShutdownException} exception, so they are expected to complete within the
* termination grace period on their own.
*/
@Experimental
public Builder setAllowActivityHeartbeatDuringShutdown(
boolean allowActivityHeartbeatDuringShutdown) {
this.allowActivityHeartbeatDuringShutdown = allowActivityHeartbeatDuringShutdown;
return this;
}

public WorkerOptions build() {
return new WorkerOptions(
maxWorkerActivitiesPerSecond,
Expand Down Expand Up @@ -553,7 +577,8 @@ public WorkerOptions build() {
deploymentOptions,
workflowTaskPollersBehavior,
activityTaskPollersBehavior,
nexusTaskPollersBehavior);
nexusTaskPollersBehavior,
allowActivityHeartbeatDuringShutdown);
}

public WorkerOptions validateAndBuildWithDefaults() {
Expand Down Expand Up @@ -685,7 +710,8 @@ public WorkerOptions validateAndBuildWithDefaults() {
deploymentOptions,
workflowTaskPollersBehavior,
activityTaskPollersBehavior,
nexusTaskPollersBehavior);
nexusTaskPollersBehavior,
allowActivityHeartbeatDuringShutdown);
}
}

Expand Down Expand Up @@ -717,6 +743,7 @@ public WorkerOptions validateAndBuildWithDefaults() {
private final PollerBehavior workflowTaskPollersBehavior;
private final PollerBehavior activityTaskPollersBehavior;
private final PollerBehavior nexusTaskPollersBehavior;
private final boolean allowActivityHeartbeatDuringShutdown;

private WorkerOptions(
double maxWorkerActivitiesPerSecond,
Expand Down Expand Up @@ -746,7 +773,8 @@ private WorkerOptions(
WorkerDeploymentOptions deploymentOptions,
PollerBehavior workflowTaskPollersBehavior,
PollerBehavior activityTaskPollersBehavior,
PollerBehavior nexusTaskPollersBehavior) {
PollerBehavior nexusTaskPollersBehavior,
boolean allowActivityHeartbeatDuringShutdown) {
this.maxWorkerActivitiesPerSecond = maxWorkerActivitiesPerSecond;
this.maxConcurrentActivityExecutionSize = maxConcurrentActivityExecutionSize;
this.maxConcurrentWorkflowTaskExecutionSize = maxConcurrentWorkflowTaskExecutionSize;
Expand Down Expand Up @@ -775,6 +803,7 @@ private WorkerOptions(
this.workflowTaskPollersBehavior = workflowTaskPollersBehavior;
this.activityTaskPollersBehavior = activityTaskPollersBehavior;
this.nexusTaskPollersBehavior = nexusTaskPollersBehavior;
this.allowActivityHeartbeatDuringShutdown = allowActivityHeartbeatDuringShutdown;
}

public double getMaxWorkerActivitiesPerSecond() {
Expand Down Expand Up @@ -912,6 +941,11 @@ public PollerBehavior getNexusTaskPollersBehavior() {
return nexusTaskPollersBehavior;
}

@Experimental
public boolean getAllowActivityHeartbeatDuringShutdown() {
return allowActivityHeartbeatDuringShutdown;
}

@Override
public boolean equals(Object o) {
if (this == o) return true;
Expand Down Expand Up @@ -944,7 +978,8 @@ && compare(maxTaskQueueActivitiesPerSecond, that.maxTaskQueueActivitiesPerSecond
&& Objects.equals(deploymentOptions, that.deploymentOptions)
&& Objects.equals(workflowTaskPollersBehavior, that.workflowTaskPollersBehavior)
&& Objects.equals(activityTaskPollersBehavior, that.activityTaskPollersBehavior)
&& Objects.equals(nexusTaskPollersBehavior, that.nexusTaskPollersBehavior);
&& Objects.equals(nexusTaskPollersBehavior, that.nexusTaskPollersBehavior)
&& allowActivityHeartbeatDuringShutdown == that.allowActivityHeartbeatDuringShutdown;
}

@Override
Expand Down Expand Up @@ -977,7 +1012,8 @@ public int hashCode() {
deploymentOptions,
workflowTaskPollersBehavior,
activityTaskPollersBehavior,
nexusTaskPollersBehavior);
nexusTaskPollersBehavior,
allowActivityHeartbeatDuringShutdown);
}

@Override
Expand Down Expand Up @@ -1040,6 +1076,8 @@ public String toString() {
+ activityTaskPollersBehavior
+ ", nexusTaskPollersBehavior="
+ nexusTaskPollersBehavior
+ ", allowActivityHeartbeatDuringShutdown="
+ allowActivityHeartbeatDuringShutdown
+ '}';
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ public void verifyNewBuilderFromExistingWorkerOptions() {
.setBuildId("build-id")
.setStickyTaskQueueDrainTimeout(Duration.ofSeconds(15))
.setIdentity("worker-identity")
.setAllowActivityHeartbeatDuringShutdown(true)
.build();

WorkerOptions w2 = WorkerOptions.newBuilder(w1).build();
Expand Down Expand Up @@ -89,6 +90,8 @@ public void verifyNewBuilderFromExistingWorkerOptions() {
assertEquals(w1.getBuildId(), w2.getBuildId());
assertEquals(w1.getStickyTaskQueueDrainTimeout(), w2.getStickyTaskQueueDrainTimeout());
assertEquals(w1.getIdentity(), w2.getIdentity());
assertEquals(
w1.getAllowActivityHeartbeatDuringShutdown(), w2.getAllowActivityHeartbeatDuringShutdown());
}

@Test
Expand Down
Loading