From 609161fe7964133edef4260f790d99416485c38b Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Wed, 19 Aug 2026 00:12:26 -0400 Subject: [PATCH] feat(extstore): integrate into nexus pipeline. --- .../temporal/internal/worker/NexusWorker.java | 90 +++++++++++++------ 1 file changed, 63 insertions(+), 27 deletions(-) diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/NexusWorker.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/NexusWorker.java index 33416a807b..a9bc2e89bb 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/worker/NexusWorker.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/worker/NexusWorker.java @@ -4,6 +4,7 @@ import static io.temporal.serviceclient.MetricsTag.TASK_FAILURE_TYPE; import com.google.protobuf.ByteString; +import com.google.protobuf.Message; import com.uber.m3.tally.Scope; import com.uber.m3.tally.Stopwatch; import com.uber.m3.util.Duration; @@ -17,6 +18,7 @@ import io.temporal.internal.common.NexusUtil; import io.temporal.internal.common.ProtobufTimeUtils; import io.temporal.internal.logging.LoggerTag; +import io.temporal.internal.payload.storage.ExternalStorage; import io.temporal.internal.retryer.GrpcRetryer; import io.temporal.serviceclient.MetricsTag; import io.temporal.serviceclient.WorkflowServiceStubs; @@ -274,6 +276,16 @@ public String toString() { options.getIdentity(), namespace, taskQueue); } + /** + * Marks a failure raised by the external storage. Reported to the server as a retryable handler + * error. + */ + private static final class ExternalStorageTaskFailure extends RuntimeException { + ExternalStorageTaskFailure(String message, Throwable cause) { + super(message, cause); + } + } + private class TaskHandlerImpl implements PollTaskExecutor.TaskHandler { final NexusTaskHandler handler; @@ -304,28 +316,30 @@ private String getNexusTaskOperation(PollNexusTaskQueueResponseOrBuilder pollRes @Override public void handle(NexusTask task) { - PollNexusTaskQueueResponseOrBuilder pollResponse = task.getResponse(); - // Extract service and operation from the request and set them as MDC and metrics - // scope tags. If the request does not have a service or operation, do not set the tags. - // If we don't know how to handle the task, we will fail the task further down the line. - Scope metricsScope = workerMetricsScope; - String service = getNexusTaskService(pollResponse); - if (!service.isEmpty()) { - MDC.put(LoggerTag.NEXUS_SERVICE, service); - metricsScope = metricsScope.tagged(ImmutableMap.of(MetricsTag.NEXUS_SERVICE, service)); - } - String operation = getNexusTaskOperation(pollResponse); - if (!operation.isEmpty()) { - MDC.put(LoggerTag.NEXUS_OPERATION, operation); - metricsScope = metricsScope.tagged(ImmutableMap.of(MetricsTag.NEXUS_OPERATION, operation)); - } - slotSupplier.markSlotUsed( - new NexusSlotInfo( - service, operation, taskQueue, options.getIdentity(), options.getBuildId()), - task.getPermit()); - boolean taskFailed = false; try { + task = retrieveInboundPayloads(task); + PollNexusTaskQueueResponseOrBuilder pollResponse = task.getResponse(); + // Extract service and operation from the request and set them as MDC and metrics + // scope tags. If the request does not have a service or operation, do not set the tags. + // If we don't know how to handle the task, we will fail the task further down the line. + Scope metricsScope = workerMetricsScope; + String service = getNexusTaskService(pollResponse); + if (!service.isEmpty()) { + MDC.put(LoggerTag.NEXUS_SERVICE, service); + metricsScope = metricsScope.tagged(ImmutableMap.of(MetricsTag.NEXUS_SERVICE, service)); + } + String operation = getNexusTaskOperation(pollResponse); + if (!operation.isEmpty()) { + MDC.put(LoggerTag.NEXUS_OPERATION, operation); + metricsScope = + metricsScope.tagged(ImmutableMap.of(MetricsTag.NEXUS_OPERATION, operation)); + } + slotSupplier.markSlotUsed( + new NexusSlotInfo( + service, operation, taskQueue, options.getIdentity(), options.getBuildId()), + task.getPermit()); + taskFailed = handleNexusTask(task, metricsScope); } catch (Throwable e) { taskFailed = true; @@ -485,12 +499,13 @@ private void sendReply( taskResponse = getResponseForOldServer(taskResponse); } RespondNexusTaskCompletedRequest request = - RespondNexusTaskCompletedRequest.newBuilder() - .setTaskToken(taskToken) - .setIdentity(options.getIdentity()) - .setNamespace(namespace) - .setResponse(taskResponse) - .build(); + storeOutbound( + RespondNexusTaskCompletedRequest.newBuilder() + .setTaskToken(taskToken) + .setIdentity(options.getIdentity()) + .setNamespace(namespace) + .setResponse(taskResponse) + .build()); grpcRetryer.retry( () -> @@ -512,17 +527,38 @@ private void sendReply( } else { request.setError(NexusUtil.handlerErrorToNexusError(handlerException, dataConverter)); } + RespondNexusTaskFailedRequest failedRequest = storeOutbound(request.build()); grpcRetryer.retry( () -> service .blockingStub() .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, metricsScope) - .respondNexusTaskFailed(request.build()), + .respondNexusTaskFailed(failedRequest), replyGrpcRetryerOptions); } else { throw new IllegalArgumentException("[BUG] Either response or failure must be set"); } } } + + private NexusTask retrieveInboundPayloads(NexusTask task) { + ExternalStorage externalStorage = options.getExternalStorage(); + PollNexusTaskQueueResponseOrBuilder response = task.getResponse(); + PollNexusTaskQueueResponse built = + response instanceof PollNexusTaskQueueResponse + ? (PollNexusTaskQueueResponse) response + : ((PollNexusTaskQueueResponse.Builder) response).build(); + if (externalStorage == null) { + ExternalStorage.throwIfContainsReference(built); + return task; + } + return new NexusTask( + externalStorage.retrieveBlocking(built), task.getPermit(), task.getCompletionCallback()); + } + + private T storeOutbound(T request) { + ExternalStorage externalStorage = options.getExternalStorage(); + return externalStorage == null ? request : externalStorage.storeBlocking(request, null); + } } }