Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Call shutdown RPC on worker shutdown #2264

Merged
merged 3 commits into from
Oct 10, 2024
Merged

Call shutdown RPC on worker shutdown #2264

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
83b73ff
Update API to v1.39.0
Quinn-With-Two-Ns Oct 9, 2024
425a49b
Call shutdownWorker on worker shutdown
Quinn-With-Two-Ns Oct 9, 2024
02ffa77
fix stub
Quinn-With-Two-Ns Oct 9, 2024
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
31 changes: 31 additions & 0 deletions temporal-sdk/src/main/java/io/temporal/internal/worker/ShutdownManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,10 @@

package io.temporal.internal.worker;

import com.google.common.util.concurrent.ListenableFuture;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import io.temporal.api.workflowservice.v1.ShutdownWorkerResponse;
import java.io.Closeable;
import java.time.Duration;
import java.util.concurrent.*;
Expand Down Expand Up @@ -119,6 +123,33 @@ private CompletableFuture<Void> limitedWait(
return future;
}

/**
* Wait for {@code shutdownRequest} to finish. shutdownRequest is considered best effort, so we do
* not fail the shutdown if it fails.
*/
public CompletableFuture<Void> waitOnWorkerShutdownRequest(
ListenableFuture<ShutdownWorkerResponse> shutdownRequest) {
CompletableFuture<Void> future = new CompletableFuture<>();
shutdownRequest.addListener(
() -> {
try {
shutdownRequest.get();
} catch (StatusRuntimeException e) {
// If the server does not support shutdown, ignore the exception
if (Status.Code.UNIMPLEMENTED.equals(e.getStatus().getCode())) {
return;
}
log.warn("failed to call shutdown worker", e);
} catch (Exception e) {
log.warn("failed to call shutdown worker", e);
} finally {
future.complete(null);
}
},
scheduledExecutorService);
return future;
}

@Override
public void close() {
scheduledExecutorService.shutdownNow();
Expand Down
68 changes: 44 additions & 24 deletions temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@
import org.slf4j.MDC;

final class WorkflowWorker implements SuspendableWorker {
private static final String GRACEFUL_SHUTDOWN_MESSAGE = "graceful shutdown";
private static final Logger log = LoggerFactory.getLogger(WorkflowWorker.class);

private final WorkflowRunLockManager runLocks;
Expand Down Expand Up @@ -162,30 +163,49 @@ public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean
&& stickyTaskQueueName != null
&& stickyQueueBalancer != null;

return CompletableFuture.completedFuture(null)
.thenCompose(
ignore ->
stickyQueueBalancerDrainEnabled
? shutdownManager.waitForStickyQueueBalancer(
stickyQueueBalancer, options.getDrainStickyTaskQueueTimeout())
: CompletableFuture.completedFuture(null))
.thenCompose(ignore -> poller.shutdown(shutdownManager, interruptTasks))
.thenCompose(
ignore ->
!interruptTasks
? shutdownManager.waitForSupplierPermitsReleasedUnlimited(
slotSupplier, supplierName)
: CompletableFuture.completedFuture(null))
.thenCompose(
ignore ->
pollTaskExecutor != null
? pollTaskExecutor.shutdown(shutdownManager, interruptTasks)
: CompletableFuture.completedFuture(null))
.exceptionally(
e -> {
log.error("Unexpected exception during shutdown", e);
return null;
});
CompletableFuture<Void> pollerShutdown =
CompletableFuture.completedFuture(null)
.thenCompose(
ignore ->
stickyQueueBalancerDrainEnabled
? shutdownManager.waitForStickyQueueBalancer(
stickyQueueBalancer, options.getDrainStickyTaskQueueTimeout())
: CompletableFuture.completedFuture(null))
.thenCompose(ignore -> poller.shutdown(shutdownManager, interruptTasks));
return CompletableFuture.allOf(
pollerShutdown.thenCompose(
ignore -> {
if (!interruptTasks && stickyTaskQueueName != null) {
return shutdownManager.waitOnWorkerShutdownRequest(
service
.futureStub()
.shutdownWorker(
ShutdownWorkerRequest.newBuilder()
.setIdentity(options.getIdentity())
.setNamespace(namespace)
.setStickyTaskQueue(stickyTaskQueueName)
.setReason(GRACEFUL_SHUTDOWN_MESSAGE)
.build()));
}
return CompletableFuture.completedFuture(null);
}),
pollerShutdown
.thenCompose(
ignore ->
!interruptTasks
? shutdownManager.waitForSupplierPermitsReleasedUnlimited(
slotSupplier, supplierName)
: CompletableFuture.completedFuture(null))
.thenCompose(
ignore ->
pollTaskExecutor != null
? pollTaskExecutor.shutdown(shutdownManager, interruptTasks)
: CompletableFuture.completedFuture(null))
.exceptionally(
e -> {
log.error("Unexpected exception during shutdown", e);
return null;
}));
}

@Override
Expand Down
13 changes: 13 additions & 0 deletions temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;

import com.google.common.util.concurrent.Futures;
import com.google.protobuf.ByteString;
import com.uber.m3.tally.NoopScope;
import com.uber.m3.tally.RootScopeBuilder;
Expand Down Expand Up @@ -100,9 +101,15 @@ public void concurrentPollRequestLockTest() throws Exception {
eagerActivityDispatcher,
slotSupplier);

WorkflowServiceGrpc.WorkflowServiceFutureStub futureStub =
mock(WorkflowServiceGrpc.WorkflowServiceFutureStub.class);
when(futureStub.shutdownWorker(any(ShutdownWorkerRequest.class)))
.thenReturn(Futures.immediateFuture(ShutdownWorkerResponse.newBuilder().build()));

WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub =
mock(WorkflowServiceGrpc.WorkflowServiceBlockingStub.class);
when(client.blockingStub()).thenReturn(blockingStub);
when(client.futureStub()).thenReturn(futureStub);
when(blockingStub.withOption(any(), any())).thenReturn(blockingStub);

PollWorkflowTaskQueueResponse pollResponse =
Expand Down Expand Up @@ -259,9 +266,15 @@ public void respondWorkflowTaskFailureMetricTest() throws Exception {
eagerActivityDispatcher,
slotSupplier);

WorkflowServiceGrpc.WorkflowServiceFutureStub futureStub =
mock(WorkflowServiceGrpc.WorkflowServiceFutureStub.class);
when(futureStub.shutdownWorker(any(ShutdownWorkerRequest.class)))
.thenReturn(Futures.immediateFuture(ShutdownWorkerResponse.newBuilder().build()));

WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub =
mock(WorkflowServiceGrpc.WorkflowServiceBlockingStub.class);
when(client.blockingStub()).thenReturn(blockingStub);
when(client.futureStub()).thenReturn(futureStub);
when(blockingStub.withOption(any(), any())).thenReturn(blockingStub);

PollWorkflowTaskQueueResponse pollResponse =
Expand Down
2 changes: 1 addition & 1 deletion temporal-serviceclient/src/main/proto
Open in desktop
Submodule proto updated 5 files
+11 −8 openapi/openapiv2.json
+49 −14 openapi/openapiv3.yaml
+46 −12 temporal/api/taskqueue/v1/message.proto
+10 −0 temporal/api/workflowservice/v1/request_response.proto
+20 −0 temporal/api/workflowservice/v1/service.proto
Loading