From 4ebb95c476bf2406ef2941ce689922a82773d4a1 Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Wed, 7 Oct 2026 18:08:22 +0530 Subject: [PATCH 01/16] Block segment publish action for a streaming task until prior task groups are done --- .../common/actions/LocalTaskActionClient.java | 26 +++- .../common/actions/SegmentAllocateAction.java | 3 +- .../actions/SegmentAllocationQueue.java | 22 +-- .../SegmentTransactionalAppendAction.java | 26 +++- .../SegmentTransactionalInsertAction.java | 26 +++- .../indexing/common/actions/TaskAction.java | 8 +- .../common/actions/TaskActionClient.java | 15 ++ .../common/actions/TaskActionToolbox.java | 22 ++- .../overlord/http/OverlordResource.java | 86 ++++++++--- .../supervisor/SupervisorManager.java | 40 ++--- .../supervisor/SeekableStreamSupervisor.java | 138 ++++++++++++++++-- .../overlord/http/OverlordResourceTest.java | 5 +- .../indexing/overlord/http/OverlordTest.java | 4 +- .../supervisor/SupervisorManagerTest.java | 44 +++--- .../overlord/SegmentPublishResult.java | 41 +++++- .../IndexerSQLMetadataStorageCoordinator.java | 2 +- 16 files changed, 380 insertions(+), 128 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/LocalTaskActionClient.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/LocalTaskActionClient.java index 1d0059335ed1..e7ff750ae133 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/LocalTaskActionClient.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/LocalTaskActionClient.java @@ -20,8 +20,12 @@ package org.apache.druid.indexing.common.actions; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.MoreExecutors; import org.apache.druid.indexing.common.task.IndexTaskUtils; import org.apache.druid.indexing.common.task.Task; +import org.apache.druid.java.util.common.Stopwatch; import org.apache.druid.java.util.common.jackson.JacksonUtils; import org.apache.druid.java.util.emitter.EmittingLogger; import org.apache.druid.java.util.emitter.service.ServiceMetricEvent; @@ -34,6 +38,8 @@ public class LocalTaskActionClient implements TaskActionClient { private static final EmittingLogger log = new EmittingLogger(LocalTaskActionClient.class); + private static final String RUN_TIME_METRIC = "task/action/run/time"; + private final Task task; private final TaskActionToolbox toolbox; @@ -52,10 +58,28 @@ public RetType submit(TaskAction taskAction) log.debug("Performing action for task[%s]: %s", task.getId(), taskAction); final long performStartTime = System.currentTimeMillis(); final RetType result = performAction(taskAction); - emitTimerMetric("task/action/run/time", taskAction, System.currentTimeMillis() - performStartTime); + emitTimerMetric(RUN_TIME_METRIC, taskAction, System.currentTimeMillis() - performStartTime); return result; } + @Override + public ListenableFuture submitAsync(TaskAction taskAction) + { + if (taskAction.canPerformAsync(task, toolbox)) { + final Stopwatch actionRunTime = Stopwatch.createStarted(); + return Futures.transform( + taskAction.performAsync(task, toolbox), + v -> { + emitTimerMetric(RUN_TIME_METRIC, taskAction, actionRunTime.millisElapsed()); + return v; + }, + MoreExecutors.directExecutor() + ); + } else { + return Futures.immediateFuture(submit(taskAction)); + } + } + private R performAction(TaskAction taskAction) { try { diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocateAction.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocateAction.java index 3f240e0aa966..944dc7fb43a9 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocateAction.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocateAction.java @@ -23,6 +23,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.core.type.TypeReference; import com.google.common.base.Preconditions; +import com.google.common.util.concurrent.ListenableFuture; import org.apache.druid.error.DruidException; import org.apache.druid.indexing.common.LockGranularity; import org.apache.druid.indexing.common.TaskLockType; @@ -194,7 +195,7 @@ public boolean canPerformAsync(Task task, TaskActionToolbox toolbox) } @Override - public Future performAsync(Task task, TaskActionToolbox toolbox) + public ListenableFuture performAsync(Task task, TaskActionToolbox toolbox) { if (!toolbox.canBatchSegmentAllocation()) { throw new ISE("Batched segment allocation is disabled"); diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocationQueue.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocationQueue.java index 7b504a336b40..7606bde2ff99 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocationQueue.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocationQueue.java @@ -19,6 +19,8 @@ package org.apache.druid.indexing.common.actions; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; import com.google.inject.Inject; import org.apache.druid.guice.ManageLifecycle; import org.apache.druid.indexing.common.LockGranularity; @@ -54,9 +56,7 @@ import java.util.Set; import java.util.TreeSet; import java.util.concurrent.BlockingDeque; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingDeque; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; @@ -211,7 +211,7 @@ public int size() * Queues a SegmentAllocateRequest. The returned future may complete successfully * with a non-null value or with a non-null value. */ - public Future add(SegmentAllocateRequest request) + public ListenableFuture add(SegmentAllocateRequest request) { if (!isLeader.get()) { throw new ISE("Cannot allocate segment if not leader."); @@ -220,7 +220,7 @@ public Future add(SegmentAllocateRequest request) } final AllocateRequestKey requestKey = new AllocateRequestKey(request); - final AtomicReference> futureReference = new AtomicReference<>(); + final AtomicReference> futureReference = new AtomicReference<>(); // Possible race condition: // t1 -> new batch is added to queue or batch already exists in queue @@ -644,7 +644,7 @@ private class AllocateRequestBatch * Map from allocate requests (represents a single SegmentAllocateAction) * to the future of allocated segment id. */ - private final Map> + private final Map> requestToFuture = new HashMap<>(); AllocateRequestBatch(AllocateRequestKey key) @@ -692,10 +692,10 @@ boolean isFull() return size() >= MAX_BATCH_SIZE; } - Future add(SegmentAllocateRequest request) + ListenableFuture add(SegmentAllocateRequest request) { log.debug("Adding request to batch [%s]: %s", key, request.getAction()); - return requestToFuture.computeIfAbsent(request, req -> new CompletableFuture<>()); + return requestToFuture.computeIfAbsent(request, req -> SettableFuture.create()); } void transferRequestsFrom(AllocateRequestBatch batch) @@ -718,7 +718,7 @@ void failPendingRequests(Throwable cause) { if (!requestToFuture.isEmpty()) { log.warn("Failing [%d] requests in batch[%s], reason[%s].", size(), key, cause.getMessage()); - requestToFuture.values().forEach(future -> future.completeExceptionally(cause)); + requestToFuture.values().forEach(future -> future.setException(cause)); requestToFuture.keySet().forEach( request -> emitTaskMetric("task/action/failed/count", 1L, request) ); @@ -732,7 +732,7 @@ void completePendingRequestsWithNull() return; } - requestToFuture.values().forEach(future -> future.complete(null)); + requestToFuture.values().forEach(future -> future.set(null)); requestToFuture.keySet().forEach( request -> emitTaskMetric("task/action/failed/count", 1L, request) ); @@ -746,7 +746,7 @@ void handleResult(SegmentAllocateResult result, SegmentAllocateRequest request) if (result.isSuccess()) { emitTaskMetric("task/action/success/count", 1L, request); IndexTaskUtils.emitSegmentAllocateMetric(result.getSegmentId(), request.getTask(), emitter); - requestToFuture.remove(request).complete(result.getSegmentId()); + requestToFuture.remove(request).set(result.getSegmentId()); } else if (request.canRetry()) { log.debug( "Allocation failed on attempt [%d] due to error[%s]. Can still retry action[%s].", @@ -759,7 +759,7 @@ void handleResult(SegmentAllocateResult result, SegmentAllocateRequest request) + " Completing action[%s] with a null value.", request.getAttempts(), result.getErrorMessage(), request.getAction() ); - requestToFuture.remove(request).complete(null); + requestToFuture.remove(request).set(null); } } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalAppendAction.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalAppendAction.java index 15001fcd4a21..5b0e6a442ffe 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalAppendAction.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalAppendAction.java @@ -22,6 +22,9 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.core.type.TypeReference; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.MoreExecutors; import org.apache.druid.error.DruidException; import org.apache.druid.error.InvalidInput; import org.apache.druid.indexing.common.TaskLock; @@ -219,12 +222,25 @@ public SegmentPublishResult perform(Task task, TaskActionToolbox toolbox) } IndexTaskUtils.emitSegmentPublishMetrics(retVal, task, toolbox); + return retVal; + } - if (toolbox.shouldFailSegmentPublishImmediately(retVal, task, supervisorId, startMetadata)) { - return SegmentPublishResult.fail(retVal.getErrorMsg()); - } else { - return retVal; - } + @Override + public ListenableFuture performAsync(Task task, TaskActionToolbox toolbox) + { + final ListenableFuture taskReadyToPublishFuture + = toolbox.isTaskReadyToPublish(task, supervisorId, startMetadata); + return Futures.transform( + taskReadyToPublishFuture, + taskIsReadyToPublish -> { + if (Boolean.TRUE.equals(taskIsReadyToPublish)) { + return SegmentPublishResult.doNotRetryOffsetMismatchFailure(perform(task, toolbox)); + } else { + return SegmentPublishResult.retryableFailure("Task is not ready to publish yet"); + } + }, + MoreExecutors.directExecutor() + ); } @Override diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertAction.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertAction.java index bf1218125730..465aee7bb843 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertAction.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertAction.java @@ -24,6 +24,9 @@ import com.fasterxml.jackson.core.type.TypeReference; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableSet; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.MoreExecutors; import org.apache.druid.common.config.Configs; import org.apache.druid.indexing.common.LockGranularity; import org.apache.druid.indexing.common.TaskLock; @@ -259,12 +262,25 @@ public SegmentPublishResult perform(Task task, TaskActionToolbox toolbox) } IndexTaskUtils.emitSegmentPublishMetrics(retVal, task, toolbox); + return retVal; + } - if (toolbox.shouldFailSegmentPublishImmediately(retVal, task, supervisorId, startMetadata)) { - return SegmentPublishResult.fail(retVal.getErrorMsg()); - } else { - return retVal; - } + @Override + public ListenableFuture performAsync(Task task, TaskActionToolbox toolbox) + { + final ListenableFuture taskReadyToPublishFuture + = toolbox.isTaskReadyToPublish(task, supervisorId, startMetadata); + return Futures.transform( + taskReadyToPublishFuture, + taskIsReadyToPublish -> { + if (Boolean.TRUE.equals(taskIsReadyToPublish)) { + return SegmentPublishResult.doNotRetryOffsetMismatchFailure(perform(task, toolbox)); + } else { + return SegmentPublishResult.retryableFailure("Task is not ready to publish yet"); + } + }, + MoreExecutors.directExecutor() + ); } private void checkWithSegmentLock() diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskAction.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskAction.java index 1deab619f75f..1f9b62c991dc 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskAction.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskAction.java @@ -22,10 +22,12 @@ import com.fasterxml.jackson.annotation.JsonSubTypes; import com.fasterxml.jackson.annotation.JsonTypeInfo; import com.fasterxml.jackson.core.type.TypeReference; +import com.google.common.util.concurrent.ListenableFuture; import org.apache.druid.indexing.common.task.Task; -import java.util.concurrent.Future; - +/** + * An action performed on behalf of a Task by the Overlord. + */ @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = TaskAction.TYPE_FIELD) @JsonSubTypes(value = { @JsonSubTypes.Type(name = "lockAcquire", value = TimeChunkLockAcquireAction.class), @@ -66,7 +68,7 @@ default boolean canPerformAsync(Task task, TaskActionToolbox toolbox) return false; } - default Future performAsync(Task task, TaskActionToolbox toolbox) + default ListenableFuture performAsync(Task task, TaskActionToolbox toolbox) { throw new UnsupportedOperationException(); } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionClient.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionClient.java index fb0e8c514000..b0439d70f890 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionClient.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionClient.java @@ -19,9 +19,24 @@ package org.apache.druid.indexing.common.actions; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; + import java.io.IOException; +/** + * Client to submit a {@link TaskAction} on behalf of a Task to the Overlord. + */ public interface TaskActionClient { RetType submit(TaskAction taskAction) throws IOException; + + default ListenableFuture submitAsync(TaskAction taskAction) + { + try { + return Futures.immediateFuture(submit(taskAction)); + } catch (IOException e) { + return Futures.immediateFailedFuture(e); + } + } } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java index 6294710f1677..4fd9b557345b 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java @@ -21,13 +21,14 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.base.Optional; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; import com.google.inject.Inject; import org.apache.druid.guice.annotations.Json; import org.apache.druid.indexing.common.task.Task; import org.apache.druid.indexing.overlord.DataSourceMetadata; import org.apache.druid.indexing.overlord.GlobalTaskLockbox; import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator; -import org.apache.druid.indexing.overlord.SegmentPublishResult; import org.apache.druid.indexing.overlord.TaskRunner; import org.apache.druid.indexing.overlord.TaskRunnerFactory; import org.apache.druid.indexing.overlord.TaskStorage; @@ -140,25 +141,20 @@ public boolean canBatchSegmentAllocation() } /** - * Checks if the given publish action should be failed without allowing any - * more retries. A failed publish action should be retried only if there is - * another task waiting to publish offsets for an overlapping set of partitions. + * Returns a future that completes when the given task is ready to publish its + * segments. A streaming task must wait for previously created tasks that are + * yet to publish offsets for an overlapping set of partitions. */ - public boolean shouldFailSegmentPublishImmediately( - SegmentPublishResult result, + public ListenableFuture isTaskReadyToPublish( Task task, String supervisorId, DataSourceMetadata startMetadata ) { - if (result.isSuccess() || !result.isRetryable() || startMetadata == null) { - return false; + if (startMetadata == null) { + return Futures.immediateFuture(true); } - return !getSupervisorManager().isAnotherTaskGroupPublishingToPartitions( - supervisorId, - task.getId(), - startMetadata - ); + return getSupervisorManager().isTaskReadyToPublishSegments(supervisorId, task.getId(), startMetadata); } } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java index a12a9ea22feb..a3747f7b4773 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java @@ -19,6 +19,7 @@ package org.apache.druid.indexing.overlord.http; +import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Function; import com.google.common.base.Optional; @@ -26,6 +27,10 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Lists; +import com.google.common.util.concurrent.FutureCallback; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.MoreExecutors; import com.google.inject.Inject; import com.sun.jersey.spi.container.ResourceFilters; import org.apache.druid.audit.AuditEntry; @@ -79,7 +84,9 @@ import org.joda.time.Interval; import javax.annotation.Nullable; +import javax.servlet.AsyncContext; import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletResponse; import javax.ws.rs.Consumes; import javax.ws.rs.DELETE; import javax.ws.rs.DefaultValue; @@ -96,7 +103,6 @@ import javax.ws.rs.core.Response.Status; import java.io.InputStream; import java.util.Collections; -import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -121,6 +127,7 @@ public class OverlordResource private final WorkerTaskRunnerQueryAdapter workerTaskRunnerQueryAdapter; private final AuthConfig authConfig; + private final ObjectMapper jsonMapper; private static final List API_TASK_STATES = ImmutableList.of("pending", "waiting", "running", "complete"); private static final Set AUDITED_TASK_TYPES @@ -137,7 +144,8 @@ public OverlordResource( AuditManager auditManager, AuthorizerMapper authorizerMapper, WorkerTaskRunnerQueryAdapter workerTaskRunnerQueryAdapter, - AuthConfig authConfig + AuthConfig authConfig, + ObjectMapper jsonMapper ) { this.overlord = overlord; @@ -150,6 +158,7 @@ public OverlordResource( this.authorizerMapper = authorizerMapper; this.workerTaskRunnerQueryAdapter = workerTaskRunnerQueryAdapter; this.authConfig = authConfig; + this.jsonMapper = jsonMapper; } /** @@ -506,34 +515,51 @@ public Response getWorkerConfigHistory( @Path("/action") @Produces(MediaType.APPLICATION_JSON) @ResourceFilters(StateResourceFilter.class) - public Response doAction(final TaskActionHolder holder) + public void doAction( + final TaskActionHolder holder, + @Context HttpServletRequest request + ) { - return asLeaderWith( - taskMaster.getTaskActionClient(holder.getTask()), - new Function<>() + final AsyncContext asyncContext = request.startAsync(); + final Optional taskActionClient = taskMaster.getTaskActionClient(holder.getTask()); + if (!taskActionClient.isPresent()) { + // Encourage client to try again soon, when we'll likely have a redirect set up + completeRequest(asyncContext, Status.SERVICE_UNAVAILABLE.getStatusCode(), null); + return; + } + + final ListenableFuture future = taskActionClient.get().submitAsync(holder.getAction()); + asyncContext.addListener( + ServletResourceUtils.createAsyncTimeoutListener(event -> { + // HTTP 204 NO_CONTENT is sent to the client. + // TODO: send back an error response that allows the client to retry + future.cancel(true); + event.getAsyncContext().complete(); + }) + ); + asyncContext.setTimeout(5 * 60_000); + + Futures.addCallback( + future, + new FutureCallback() { @Override - public Response apply(TaskActionClient taskActionClient) + public void onSuccess(Object result) { - final Map retMap; - - // It would be great to verify that this worker is actually supposed to be running the task before - // actually doing the action. Some ideas for how that could be done would be using some sort of attempt_id - // or token that gets passed around. + completeRequest(asyncContext, Status.OK.getStatusCode(), Map.of("result", result)); + } - try { - final Object ret = taskActionClient.submit(holder.getAction()); - retMap = new HashMap<>(); - retMap.put("result", ret); - } - catch (Exception e) { - log.warn(e, "Failed to perform task action"); - return Response.serverError().entity(ImmutableMap.of("error", e.getMessage())).build(); + @Override + public void onFailure(Throwable t) + { + if (t instanceof DruidException druidException) { + completeRequest(asyncContext, druidException.getStatusCode(), druidException.toErrorResponse()); + } else { + completeRequest(asyncContext, Status.INTERNAL_SERVER_ERROR.getStatusCode(), null); } - - return Response.ok().entity(retMap).build(); } - } + }, + MoreExecutors.directExecutor() ); } @@ -884,4 +910,18 @@ private List securedTaskStatusPlus( ) ); } + + private void completeRequest(AsyncContext context, int statusCode, Object result) + { + try { + final HttpServletResponse response = (HttpServletResponse) context.getResponse(); + response.setStatus(statusCode); + if (result != null) { + response.setContentType(MediaType.APPLICATION_JSON); + jsonMapper.writeValue(response.getOutputStream(), result); + } + } catch (Exception e) { + log.noStackTrace().warn(e, "Request timed out or is closed already"); + } + } } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java index 3510ae7cfb6d..fabae3a48247 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java @@ -24,6 +24,7 @@ import com.google.common.base.Optional; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import com.google.inject.Inject; import org.apache.druid.common.config.Configs; @@ -597,14 +598,13 @@ public OptionalInt registerUpgradedPendingSegmentOnSupervisor( } /** - * Checks if there is a Task distinct from the given {@code taskId} or its replicas - * that is currently waiting to publish offsets for the given partitions. + * Checks if the given {@code supervisorId} represents a {@link SeekableStreamSupervisor} + * and if the {@code taskId} is ready to publish its segments. * - * @return true only if the given {@param supervisorId} represents a - * {@link SeekableStreamSupervisor} and the supervisor has other tasks that - * are currently publishing offsets to an overlapping set of partitions. + * @return A future that completes successfully with true when the given + * {@code taskId} can proceed with publishing its segments. */ - public boolean isAnotherTaskGroupPublishingToPartitions( + public ListenableFuture isTaskReadyToPublishSegments( String supervisorId, String taskId, DataSourceMetadata startMetadata @@ -616,7 +616,7 @@ public boolean isAnotherTaskGroupPublishingToPartitions( throw NotFound.exception("Could not find supervisor[%s]", supervisorId); } if (!(supervisor.lhs instanceof SeekableStreamSupervisor)) { - return false; + return Futures.immediateFuture(true); } if (!(startMetadata instanceof SeekableStreamDataSourceMetadata)) { @@ -626,24 +626,14 @@ public boolean isAnotherTaskGroupPublishingToPartitions( ); } - try { - final Set partitionIds = Set.copyOf( - ((SeekableStreamDataSourceMetadata) startMetadata) - .getSeekableStreamSequenceNumbers() - .getPartitionSequenceNumberMap() - .keySet() - ); - return ((SeekableStreamSupervisor) supervisor.lhs) - .isAnotherTaskGroupPublishingToPartitions(taskId, partitionIds); - } - catch (Exception e) { - log.error( - e, - "Failed to check if a publish is pending for supervisor[%s], metadata[%s]", - supervisorId, startMetadata - ); - return false; - } + final Set partitionIds = Set.copyOf( + ((SeekableStreamDataSourceMetadata) startMetadata) + .getSeekableStreamSequenceNumbers() + .getPartitionSequenceNumberMap() + .keySet() + ); + return ((SeekableStreamSupervisor) supervisor.lhs) + .isTaskReadyToPublishSegments(taskId, partitionIds); } /** diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java index 3377090a29c4..950338afcecc 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java @@ -38,6 +38,7 @@ import com.google.common.util.concurrent.ListeningExecutorService; import com.google.common.util.concurrent.ListeningScheduledExecutorService; import com.google.common.util.concurrent.MoreExecutors; +import com.google.common.util.concurrent.SettableFuture; import it.unimi.dsi.fastutil.ints.Int2ObjectLinkedOpenHashMap; import it.unimi.dsi.fastutil.ints.Int2ObjectMap; import jakarta.validation.constraints.NotNull; @@ -129,6 +130,7 @@ import java.util.concurrent.Callable; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.ScheduledExecutorService; @@ -225,6 +227,9 @@ public class TaskGroup boolean handoffEarly = false; // set by SupervisorManager.stopTaskGroupEarly + final AtomicBoolean completed = new AtomicBoolean(false); + final List completionListeners = new ArrayList<>(); + TaskGroup( int groupId, ImmutableMap startingSequences, @@ -358,6 +363,33 @@ public String getBaseSequenceName() return baseSequenceName; } + /** + * Called when this TaskGroup is removed from {@link #pendingCompletionTaskGroups}. + */ + private void onCompleted() + { + if (completed.compareAndSet(false, true)) { + runCompletionListeners(); + } + } + + private synchronized void addCompletionListener(Runnable runnable) + { + if (completed.get()) { + runnable.run(); + } else { + completionListeners.add(runnable); + } + } + + private synchronized void runCompletionListeners() + { + for (Runnable runnable : completionListeners) { + runnable.run(); + } + completionListeners.clear(); + } + @Override public String toString() { @@ -385,6 +417,37 @@ public String toString() } } + /** + * Represents a task that waits until all prior task groups that were publishing + * to overlapping partitions have completed. + */ + private class TaskWaitingToPublish + { + final CountDownLatch latch; + final SettableFuture readyToPublish = SettableFuture.create(); + final Map> blockingTaskGroups; + + TaskWaitingToPublish(String taskId, Map> blockingTaskGroups) + { + this.latch = new CountDownLatch(blockingTaskGroups.size()); + this.blockingTaskGroups = new HashMap<>(blockingTaskGroups); + + tasksWaitingToPublish.put(taskId, this); + + for (TaskGroup group : blockingTaskGroups.keySet()) { + group.addCompletionListener(() -> { + latch.countDown(); + blockingTaskGroups.remove(group); + + if (latch.getCount() == 0) { + readyToPublish.set(true); + tasksWaitingToPublish.remove(taskId); + } + }); + } + } + } + /** * Notice is used to queue tasks that are internal to the supervisor */ @@ -1009,6 +1072,13 @@ public String getType() // Map<{group id}, List<{pending completion task groups}>> private final ConcurrentHashMap> pendingCompletionTaskGroups = new ConcurrentHashMap<>(); + /** + * Tasks waiting on other task groups to finish publishing to overlapping partitions. + * Map from taskId to {@link TaskWaitingToPublish}. This map is maintained only + * for reporting purposes. + */ + private final ConcurrentHashMap tasksWaitingToPublish = new ConcurrentHashMap<>(); + // We keep two separate maps for tracking the current state of partition->task group mappings [partitionGroups] and partition->offset // mappings [partitionOffsets]. The starting offset for a new partition in [partitionOffsets] is initially set to getNotSetMarker(). When a new task group // is created and is assigned partitions, if the offset for an assigned partition in [partitionOffsets] is getNotSetMarker() it will take the starting @@ -2926,11 +2996,17 @@ private void verifyAndMergeCheckpoints( /** * Checks if there is a Task distinct from the given {@code taskId} or its replicas - * publishing to any of the given partitions. If this method returns true, it + * publishing to any of the given partitions. If this method returns a non-empty map,it * indicates that the current task would need to wait for the other tasks to * finish publishing before it can publish its own offsets. + * + * @return Map containing the conflicting TaskGroups that are publishing to + * overlapping partitions. */ - public boolean isAnotherTaskGroupPublishingToPartitions(String taskId, Set partitions) + private Map> getOtherTaskGroupsPublishingToPartitions( + String taskId, + Set partitions + ) { // Identify all the partitions that are being published by other taskGroups final Map> partitionIdToPublishingGroups = new HashMap<>(); @@ -2945,27 +3021,62 @@ public boolean isAnotherTaskGroupPublishingToPartitions(String taskId, Set> blockingTaskGroups = new HashMap<>(); for (Object partition : partitions) { @SuppressWarnings("unchecked") final PartitionIdType partitionId = (PartitionIdType) partition; if (partitionIdToPublishingGroups.containsKey(partitionId)) { - log.info( - "Task[%s] needs to wait before publishing as other taskGroups[%s] are currently publishing to partition[%s].", - taskId, - partitionIdToPublishingGroups.get(partitionId), - partitionId - ); - return true; + for (TaskGroup taskGroup : partitionIdToPublishingGroups.get(partitionId)) { + blockingTaskGroups.computeIfAbsent(taskGroup, _ -> new HashSet<>()) + .add(partitionId); + } } } - return false; + return blockingTaskGroups; } + private boolean isAnotherTaskGroupPublishingToPartitions( + String taskId, + Set partitions + ) + { + return !getOtherTaskGroupsPublishingToPartitions(taskId, partitions).isEmpty(); + } + + /** + * Checks if there is a Task distinct from the given {@code taskId} or its replicas + * publishing to any of the given partitions. + * + * @return A future that completes when the current task is ready to publish + * its segments since all previous tasks that were reading from an overlapping + * set of partitions have already published their offsets. + */ + public ListenableFuture isTaskReadyToPublishSegments(String taskId, Set partitions) + { + final Map> blockingTaskGroups = + getOtherTaskGroupsPublishingToPartitions(taskId, partitions); + + // No tasks writing to overlapping partitions, task may proceed with publish + if (blockingTaskGroups.isEmpty()) { + return Futures.immediateFuture(true); + } else { + final TaskWaitingToPublish waitingTask = new TaskWaitingToPublish(taskId, blockingTaskGroups); + + log.info( + "Task[%s] needs to wait before publishing as [%d] other taskGroups are" + + " currently publishing to partitions[%s].", + taskId, + blockingTaskGroups.size(), + blockingTaskGroups.values().stream().flatMap(Collection::stream).collect(Collectors.toSet()) + ); + return waitingTask.readyToPublish; + } + } @VisibleForTesting protected void addDiscoveredTaskToPendingCompletionTaskGroups( @@ -4189,7 +4300,7 @@ private void checkPendingCompletionTasks() // kill all the tasks in this pending completion group killTasksInGroup( group, - "No task in pending completion taskGroup[%d] succeeded before completion timeout elapsed", + "No successful task in taskGroup within completion timeout", groupId ); // set a flag so the other pending completion groups for this set of partitions will also stop @@ -4198,7 +4309,7 @@ private void checkPendingCompletionTasks() // kill all the tasks in the currently reading task group and remove the bad task group killTasksInGroup( activelyReadingTaskGroups.remove(groupId), - "No task in the corresponding pending completion taskGroup[%d] succeeded before completion timeout elapsed", + "No successful task in prior taskGroup within completion timeout", groupId ); toRemove.add(group); @@ -4206,6 +4317,7 @@ private void checkPendingCompletionTasks() } taskGroupList.removeAll(toRemove); + toRemove.forEach(TaskGroup::onCompleted); } // Ignore return value; just await. diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordResourceTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordResourceTest.java index b65465a6e9f2..74ab86de2d83 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordResourceTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordResourceTest.java @@ -184,7 +184,8 @@ public Access authorize(AuthenticationResult authenticationResult, Resource reso auditManager, authMapper, workerTaskRunnerQueryAdapter, - authConfig + authConfig, + TestHelper.JSON_MAPPER ); } @@ -1463,7 +1464,7 @@ public void testGetTaskSegmentsReturns404() { replayAll(); OverlordResource overlordResource = - new OverlordResource(null, null, null, null, null, null, null, null, null, null); + new OverlordResource(null, null, null, null, null, null, null, null, null, null, null); final Response response = overlordResource.getTaskSegments("taskId"); Assertions.assertEquals(404, response.getStatus()); Assertions.assertEquals( diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordTest.java index 1bca81b8759a..822362cf59ed 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordTest.java @@ -75,6 +75,7 @@ import org.apache.druid.java.util.emitter.EmittingLogger; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.metadata.segment.cache.SegmentMetadataCache; +import org.apache.druid.segment.TestHelper; import org.apache.druid.server.DruidNode; import org.apache.druid.server.coordinator.CoordinatorOverlordServiceConfig; import org.apache.druid.server.metrics.NoopServiceEmitter; @@ -296,7 +297,8 @@ public void testOverlordRun() throws Exception auditManager, AuthTestUtils.TEST_AUTHORIZER_MAPPER, workerTaskRunnerQueryAdapter, - new AuthConfig() + new AuthConfig(), + TestHelper.JSON_MAPPER ); Response response = overlordResource.getLeader(); Assertions.assertEquals(druidNode.getHostAndPort(), response.getEntity()); diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java index a93686972a17..a9688d79dfc2 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java @@ -27,6 +27,8 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.SettableFuture; import org.apache.druid.data.input.impl.ByteEntity; import org.apache.druid.data.input.impl.DimensionsSpec; @@ -1150,7 +1152,7 @@ public void test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifSupe { final DruidException exception = Assertions.assertThrows( DruidException.class, - () -> manager.isAnotherTaskGroupPublishingToPartitions(null, "task1", null) + () -> manager.isTaskReadyToPublishSegments(null, "task1", null) ); Assertions.assertEquals(DruidException.Persona.USER, exception.getTargetPersona()); Assertions.assertEquals(DruidException.Category.INVALID_INPUT, exception.getCategory()); @@ -1159,11 +1161,11 @@ public void test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifSupe } @Test - public void test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifSupervisorNotFound() + public void test_isTaskReadyToPublish_Segments_throwsException_ifSupervisorNotFound() { final DruidException exception = Assertions.assertThrows( DruidException.class, - () -> manager.isAnotherTaskGroupPublishingToPartitions("supervisor1", "task1", null) + () -> manager.isTaskReadyToPublishSegments("supervisor1", "task1", null) ); Assertions.assertEquals(DruidException.Persona.USER, exception.getTargetPersona()); Assertions.assertEquals(DruidException.Category.NOT_FOUND, exception.getCategory()); @@ -1172,7 +1174,8 @@ public void test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifSupe } @Test - public void test_isAnotherTaskGroupPublishingToPartitions_returnsFalse_forNonStreamingSupervisor() + public void test_isTaskReadyToPublishSegments_returnsCompleteFuture_forNonStreamingSupervisor() + throws Exception { final String supervisorId = "supervisor1"; EasyMock.expect(metadataSupervisorManager.getLatest()).andReturn( @@ -1185,8 +1188,8 @@ public void test_isAnotherTaskGroupPublishingToPartitions_returnsFalse_forNonStr manager.start(); - Assertions.assertFalse( - manager.isAnotherTaskGroupPublishingToPartitions(supervisorId, "task1", null) + Assertions.assertTrue( + manager.isTaskReadyToPublishSegments(supervisorId, "task1", null).get() ); } @@ -1209,7 +1212,7 @@ public void test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifMeta final DruidException exception = Assertions.assertThrows( DruidException.class, - () -> manager.isAnotherTaskGroupPublishingToPartitions(supervisorId, "task1", null) + () -> manager.isTaskReadyToPublishSegments(supervisorId, "task1", null) ); Assertions.assertEquals(DruidException.Persona.USER, exception.getTargetPersona()); Assertions.assertEquals(DruidException.Category.INVALID_INPUT, exception.getCategory()); @@ -1239,7 +1242,7 @@ public void test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifMeta final DruidException exception = Assertions.assertThrows( DruidException.class, - () -> manager.isAnotherTaskGroupPublishingToPartitions(supervisorId, "task1", new ObjectMetadata("abc")) + () -> manager.isTaskReadyToPublishSegments(supervisorId, "task1", new ObjectMetadata("abc")) ); Assertions.assertEquals(DruidException.Persona.USER, exception.getTargetPersona()); Assertions.assertEquals(DruidException.Category.INVALID_INPUT, exception.getCategory()); @@ -1253,7 +1256,7 @@ public void test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifMeta } @Test - public void test_isAnotherTaskGroupPublishingToPartitions() + public void test_isTaskReadyToPublishSegments() throws Exception { final String supervisorId = "supervisor1"; final SeekableStreamSupervisor seekableStreamSupervisor = @@ -1268,20 +1271,20 @@ public void test_isAnotherTaskGroupPublishingToPartitions() // Expect a readyTaskId for which no other group is currently publishing final String readyTaskId = "task1"; EasyMock.expect( - seekableStreamSupervisor.isAnotherTaskGroupPublishingToPartitions( + seekableStreamSupervisor.isTaskReadyToPublishSegments( EasyMock.eq(readyTaskId), EasyMock.anyObject() ) - ).andReturn(false).atLeastOnce(); + ).andReturn(Futures.immediateFuture(true)).atLeastOnce(); // Expect a conflictingTaskId for which another group is currently publishing final String conflictingTaskId = "task2"; EasyMock.expect( - seekableStreamSupervisor.isAnotherTaskGroupPublishingToPartitions( + seekableStreamSupervisor.isTaskReadyToPublishSegments( EasyMock.eq(conflictingTaskId), EasyMock.anyObject() ) - ).andReturn(true).atLeastOnce(); + ).andReturn(SettableFuture.create()).atLeastOnce(); replayAll(); EasyMock.replay(seekableStreamSupervisor); @@ -1291,12 +1294,15 @@ public void test_isAnotherTaskGroupPublishingToPartitions() "topic", Map.of("0", "10") ); - Assertions.assertTrue( - manager.isAnotherTaskGroupPublishingToPartitions(supervisorId, conflictingTaskId, startMetadata) - ); - Assertions.assertFalse( - manager.isAnotherTaskGroupPublishingToPartitions(supervisorId, readyTaskId, startMetadata) - ); + + final ListenableFuture conflictingTaskFuture = + manager.isTaskReadyToPublishSegments(supervisorId, conflictingTaskId, startMetadata); + Assertions.assertFalse(conflictingTaskFuture.isDone()); + + final ListenableFuture readyTaskFuture = + manager.isTaskReadyToPublishSegments(supervisorId, readyTaskId, startMetadata); + Assertions.assertTrue(readyTaskFuture.isDone()); + Assertions.assertTrue(readyTaskFuture.get()); } @Test diff --git a/server/src/main/java/org/apache/druid/indexing/overlord/SegmentPublishResult.java b/server/src/main/java/org/apache/druid/indexing/overlord/SegmentPublishResult.java index 04b745c812e5..e19c9dc5f259 100644 --- a/server/src/main/java/org/apache/druid/indexing/overlord/SegmentPublishResult.java +++ b/server/src/main/java/org/apache/druid/indexing/overlord/SegmentPublishResult.java @@ -41,6 +41,7 @@ public class SegmentPublishResult private final Set segments; private final boolean success; private final boolean retryable; + private final boolean offsetMismatch; private final String errorMsg; private final List upgradedPendingSegments; @@ -51,17 +52,36 @@ public static SegmentPublishResult ok(Set segments) public static SegmentPublishResult ok(Set segments, List upgradedPendingSegments) { - return new SegmentPublishResult(segments, true, false, null, upgradedPendingSegments); + return new SegmentPublishResult(segments, true, false, null, upgradedPendingSegments, false); } public static SegmentPublishResult fail(String errorMsg, Object... args) { - return new SegmentPublishResult(Set.of(), false, false, StringUtils.format(errorMsg, args), null); + return new SegmentPublishResult(Set.of(), false, false, StringUtils.format(errorMsg, args), null, false); } public static SegmentPublishResult retryableFailure(String errorMsg, Object... args) { - return new SegmentPublishResult(Set.of(), false, true, StringUtils.format(errorMsg, args), null); + return new SegmentPublishResult(Set.of(), false, true, StringUtils.format(errorMsg, args), null, false); + } + + public static SegmentPublishResult retryableOffsetMismatchFailure(String errorMsg, Object... args) + { + return new SegmentPublishResult(Set.of(), false, true, StringUtils.format(errorMsg, args), null, true); + } + + /** + * Converts the given publish result to a non-retryable failure if the publish + * had failed due to offset mismatch. If the task has already waited for older + * tasks before trying to publish its segments, retrying the publish would not make it pass. + */ + public static SegmentPublishResult doNotRetryOffsetMismatchFailure(SegmentPublishResult result) + { + if (result.isSuccess() || !result.isOffsetMismatch() || !result.isRetryable()) { + return result; + } else { + return SegmentPublishResult.fail(result.getErrorMsg()); + } } @JsonCreator @@ -72,7 +92,7 @@ private SegmentPublishResult( @JsonProperty("errorMsg") @Nullable String errorMsg ) { - this(segments, success, retryable, errorMsg, null); + this(segments, success, retryable, errorMsg, null, false); } private SegmentPublishResult( @@ -80,7 +100,8 @@ private SegmentPublishResult( boolean success, boolean retryable, @Nullable String errorMsg, - List upgradedPendingSegments + List upgradedPendingSegments, + boolean offsetMismatch ) { this.segments = Preconditions.checkNotNull(segments, "segments"); @@ -88,6 +109,7 @@ private SegmentPublishResult( this.errorMsg = errorMsg; this.retryable = retryable; this.upgradedPendingSegments = upgradedPendingSegments; + this.offsetMismatch = offsetMismatch; if (!success) { Preconditions.checkArgument(segments.isEmpty(), "segments must be empty for unsuccessful publishes"); @@ -129,6 +151,15 @@ public boolean isRetryable() return retryable; } + /** + * @return true if and only if the segment publish failed due to a mismatch + * between provided start offset and existing end offset in DB. + */ + public boolean isOffsetMismatch() + { + return offsetMismatch; + } + @Nullable public List getUpgradedPendingSegments() { diff --git a/server/src/main/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinator.java b/server/src/main/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinator.java index f469b0e7ebba..ee8933cf32cd 100644 --- a/server/src/main/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinator.java +++ b/server/src/main/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinator.java @@ -2401,7 +2401,7 @@ protected SegmentPublishResult updateDataSourceMetadataInTransaction( // Offsets stored in startMetadata is greater than the last committed metadata. // This can happen because the previous task is still publishing its segments and can resolve once // the previous task finishes publishing. - return SegmentPublishResult.retryableFailure( + return SegmentPublishResult.retryableOffsetMismatchFailure( "The new start metadata state[%s] is ahead of the last committed" + " end state[%s]. Try resetting the supervisor.", startMetadata, oldCommitMetadataFromDb From 4472883c27fc6be1bcde9064d3e31dd2f3e0f17b Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Wed, 7 Oct 2026 18:23:04 +0530 Subject: [PATCH 02/16] Add missing method --- .../common/actions/SegmentTransactionalAppendAction.java | 6 ++++++ .../common/actions/SegmentTransactionalInsertAction.java | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalAppendAction.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalAppendAction.java index 5b0e6a442ffe..f2c4e45b0511 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalAppendAction.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalAppendAction.java @@ -225,6 +225,12 @@ public SegmentPublishResult perform(Task task, TaskActionToolbox toolbox) return retVal; } + @Override + public boolean canPerformAsync(Task task, TaskActionToolbox toolbox) + { + return supervisorId != null && startMetadata != null; + } + @Override public ListenableFuture performAsync(Task task, TaskActionToolbox toolbox) { diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertAction.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertAction.java index 465aee7bb843..33b2a452b0ab 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertAction.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertAction.java @@ -265,6 +265,12 @@ public SegmentPublishResult perform(Task task, TaskActionToolbox toolbox) return retVal; } + @Override + public boolean canPerformAsync(Task task, TaskActionToolbox toolbox) + { + return supervisorId != null && startMetadata != null; + } + @Override public ListenableFuture performAsync(Task task, TaskActionToolbox toolbox) { From 0330152aa2d4f5c876ceee86c28baf715d59455e Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Wed, 7 Oct 2026 18:29:31 +0530 Subject: [PATCH 03/16] Fix test and stuff --- .../druid/indexing/common/actions/TaskActionToolbox.java | 8 ++++++-- .../actions/SegmentTransactionalInsertActionTest.java | 3 ++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java index 4fd9b557345b..8083cc132a84 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java @@ -151,10 +151,14 @@ public ListenableFuture isTaskReadyToPublish( DataSourceMetadata startMetadata ) { - if (startMetadata == null) { + if (startMetadata == null || supervisorId == null) { return Futures.immediateFuture(true); } - return getSupervisorManager().isTaskReadyToPublishSegments(supervisorId, task.getId(), startMetadata); + return supervisorManager.isTaskReadyToPublishSegments( + supervisorId, + task.getId(), + startMetadata + ); } } diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertActionTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertActionTest.java index 4e9ae27f8875..eb09e017acbd 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertActionTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertActionTest.java @@ -207,8 +207,9 @@ public void test_fail_transactionalUpdateDataSourceMetadata() throws Exception actionTestKit.getTaskActionToolbox() ); + Assertions.assertTrue(result.isOffsetMismatch()); Assertions.assertEquals( - SegmentPublishResult.fail( + SegmentPublishResult.retryableOffsetMismatchFailure( "The new start metadata state[ObjectMetadata{theObject=[1]}] is" + " ahead of the last committed end state[null]. Try resetting the supervisor." ), From a4d6f530e159b8d444ad1bf04d3ae542fdc2d6e6 Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Wed, 7 Oct 2026 23:48:24 +0530 Subject: [PATCH 04/16] Fix checkstyle --- .../druid/indexing/common/actions/SegmentAllocateAction.java | 1 - .../druid/indexing/common/actions/TaskActionClient.java | 3 ++- .../apache/druid/indexing/overlord/http/OverlordResource.java | 3 ++- .../seekablestream/supervisor/SeekableStreamSupervisor.java | 4 ++-- 4 files changed, 6 insertions(+), 5 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocateAction.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocateAction.java index 944dc7fb43a9..859d60bdd16c 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocateAction.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentAllocateAction.java @@ -51,7 +51,6 @@ import java.util.HashSet; import java.util.List; import java.util.Set; -import java.util.concurrent.Future; import java.util.concurrent.ThreadLocalRandom; /** diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionClient.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionClient.java index b0439d70f890..a7a8a545e93c 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionClient.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionClient.java @@ -35,7 +35,8 @@ default ListenableFuture submitAsync(TaskAction task { try { return Futures.immediateFuture(submit(taskAction)); - } catch (IOException e) { + } + catch (IOException e) { return Futures.immediateFailedFuture(e); } } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java index a3747f7b4773..fc24d7788294 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java @@ -920,7 +920,8 @@ private void completeRequest(AsyncContext context, int statusCode, Object result response.setContentType(MediaType.APPLICATION_JSON); jsonMapper.writeValue(response.getOutputStream(), result); } - } catch (Exception e) { + } + catch (Exception e) { log.noStackTrace().warn(e, "Request timed out or is closed already"); } } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java index 950338afcecc..5090d9fa0894 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java @@ -228,7 +228,7 @@ public class TaskGroup boolean handoffEarly = false; // set by SupervisorManager.stopTaskGroupEarly final AtomicBoolean completed = new AtomicBoolean(false); - final List completionListeners = new ArrayList<>(); + final List completionListeners = new ArrayList<>(); TaskGroup( int groupId, @@ -3030,7 +3030,7 @@ private Map> getOtherTaskGroupsPublishingToParti @SuppressWarnings("unchecked") final PartitionIdType partitionId = (PartitionIdType) partition; if (partitionIdToPublishingGroups.containsKey(partitionId)) { - for (TaskGroup taskGroup : partitionIdToPublishingGroups.get(partitionId)) { + for (TaskGroup taskGroup : partitionIdToPublishingGroups.get(partitionId)) { blockingTaskGroups.computeIfAbsent(taskGroup, _ -> new HashSet<>()) .add(partitionId); } From 97b64ae7eeec1a46a62195ac72cac80df2604c9c Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Wed, 7 Oct 2026 23:53:58 +0530 Subject: [PATCH 05/16] Remove TODO --- .../apache/druid/indexing/overlord/http/OverlordResource.java | 1 - 1 file changed, 1 deletion(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java index fc24d7788294..6d51c59812e9 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java @@ -532,7 +532,6 @@ public void doAction( asyncContext.addListener( ServletResourceUtils.createAsyncTimeoutListener(event -> { // HTTP 204 NO_CONTENT is sent to the client. - // TODO: send back an error response that allows the client to retry future.cancel(true); event.getAsyncContext().complete(); }) From 7ae8a082653393b274b5d30dff5b5d11e0d68743 Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Fri, 9 Oct 2026 12:52:36 +0530 Subject: [PATCH 06/16] Fixes and feedback --- .../SegmentTransactionalAppendAction.java | 16 +----- .../SegmentTransactionalInsertAction.java | 16 +----- .../common/actions/TaskActionToolbox.java | 53 ++++++++++++++++--- .../overlord/http/OverlordResource.java | 12 ++++- .../supervisor/SeekableStreamSupervisor.java | 6 +-- .../common/actions/TaskActionTestKit.java | 3 +- .../overlord/SegmentPublishResult.java | 14 ----- 7 files changed, 62 insertions(+), 58 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalAppendAction.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalAppendAction.java index f2c4e45b0511..77a62cbd3e99 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalAppendAction.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalAppendAction.java @@ -22,9 +22,7 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.core.type.TypeReference; -import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; -import com.google.common.util.concurrent.MoreExecutors; import org.apache.druid.error.DruidException; import org.apache.druid.error.InvalidInput; import org.apache.druid.indexing.common.TaskLock; @@ -234,19 +232,7 @@ public boolean canPerformAsync(Task task, TaskActionToolbox toolbox) @Override public ListenableFuture performAsync(Task task, TaskActionToolbox toolbox) { - final ListenableFuture taskReadyToPublishFuture - = toolbox.isTaskReadyToPublish(task, supervisorId, startMetadata); - return Futures.transform( - taskReadyToPublishFuture, - taskIsReadyToPublish -> { - if (Boolean.TRUE.equals(taskIsReadyToPublish)) { - return SegmentPublishResult.doNotRetryOffsetMismatchFailure(perform(task, toolbox)); - } else { - return SegmentPublishResult.retryableFailure("Task is not ready to publish yet"); - } - }, - MoreExecutors.directExecutor() - ); + return toolbox.publishSegmentsWhenReady(task, supervisorId, startMetadata, this::perform); } @Override diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertAction.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertAction.java index 33b2a452b0ab..2035ff5a7967 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertAction.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/SegmentTransactionalInsertAction.java @@ -24,9 +24,7 @@ import com.fasterxml.jackson.core.type.TypeReference; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableSet; -import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; -import com.google.common.util.concurrent.MoreExecutors; import org.apache.druid.common.config.Configs; import org.apache.druid.indexing.common.LockGranularity; import org.apache.druid.indexing.common.TaskLock; @@ -274,19 +272,7 @@ public boolean canPerformAsync(Task task, TaskActionToolbox toolbox) @Override public ListenableFuture performAsync(Task task, TaskActionToolbox toolbox) { - final ListenableFuture taskReadyToPublishFuture - = toolbox.isTaskReadyToPublish(task, supervisorId, startMetadata); - return Futures.transform( - taskReadyToPublishFuture, - taskIsReadyToPublish -> { - if (Boolean.TRUE.equals(taskIsReadyToPublish)) { - return SegmentPublishResult.doNotRetryOffsetMismatchFailure(perform(task, toolbox)); - } else { - return SegmentPublishResult.retryableFailure("Task is not ready to publish yet"); - } - }, - MoreExecutors.directExecutor() - ); + return toolbox.publishSegmentsWhenReady(task, supervisorId, startMetadata, this::perform); } private void checkWithSegmentLock() diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java index 8083cc132a84..78e3b154e7ae 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java @@ -29,12 +29,18 @@ import org.apache.druid.indexing.overlord.DataSourceMetadata; import org.apache.druid.indexing.overlord.GlobalTaskLockbox; import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator; +import org.apache.druid.indexing.overlord.SegmentPublishResult; import org.apache.druid.indexing.overlord.TaskRunner; import org.apache.druid.indexing.overlord.TaskRunnerFactory; import org.apache.druid.indexing.overlord.TaskStorage; import org.apache.druid.indexing.overlord.supervisor.SupervisorManager; +import org.apache.druid.java.util.common.concurrent.ScheduledExecutorFactory; +import org.apache.druid.java.util.common.concurrent.ScheduledExecutors; import org.apache.druid.java.util.emitter.service.ServiceEmitter; +import java.util.concurrent.ScheduledExecutorService; +import java.util.function.BiFunction; + public class TaskActionToolbox { private final GlobalTaskLockbox taskLockbox; @@ -44,6 +50,7 @@ public class TaskActionToolbox private final ServiceEmitter emitter; private final SupervisorManager supervisorManager; private final ObjectMapper jsonMapper; + private final ScheduledExecutorService actionExec; private Optional factory = Optional.absent(); @Inject @@ -54,7 +61,8 @@ public TaskActionToolbox( SegmentAllocationQueue segmentAllocationQueue, ServiceEmitter emitter, SupervisorManager supervisorManager, - @Json ObjectMapper jsonMapper + @Json ObjectMapper jsonMapper, + ScheduledExecutorFactory scheduledExecutorFactory ) { this.taskLockbox = taskLockbox; @@ -64,6 +72,7 @@ public TaskActionToolbox( this.supervisorManager = supervisorManager; this.jsonMapper = jsonMapper; this.segmentAllocationQueue = segmentAllocationQueue; + this.actionExec = scheduledExecutorFactory.create(4, "TaskActionToolbox-%s"); } public TaskActionToolbox( @@ -82,7 +91,8 @@ public TaskActionToolbox( null, emitter, supervisorManager, - jsonMapper + jsonMapper, + ScheduledExecutors::fixed ); } @@ -141,24 +151,51 @@ public boolean canBatchSegmentAllocation() } /** - * Returns a future that completes when the given task is ready to publish its - * segments. A streaming task must wait for previously created tasks that are + * Performs the segment publish action when the given task is unblocked for publish. + * A streaming task must wait for previously created tasks that are * yet to publish offsets for an overlapping set of partitions. */ - public ListenableFuture isTaskReadyToPublish( + public ListenableFuture publishSegmentsWhenReady( Task task, String supervisorId, - DataSourceMetadata startMetadata + DataSourceMetadata startMetadata, + BiFunction publishAction ) { if (startMetadata == null || supervisorId == null) { - return Futures.immediateFuture(true); + return Futures.immediateFuture(publishAction.apply(task, this)); } - return supervisorManager.isTaskReadyToPublishSegments( + final ListenableFuture taskReadyToPublishFuture = supervisorManager.isTaskReadyToPublishSegments( supervisorId, task.getId(), startMetadata ); + + return Futures.transform( + taskReadyToPublishFuture, + readyToPublish -> { + if (Boolean.TRUE.equals(readyToPublish)) { + // Task is already unblocked for publish, retrying will not fix offset mismatch + return doNotRetryOffsetMismatchFailure(publishAction.apply(task, this)); + } else { + return SegmentPublishResult.retryableFailure("Task is not ready to publish yet"); + } + }, + actionExec + ); + } + + /** + * Converts the given publish result to a non-retryable failure if the publish + * had failed due to offset mismatch. + */ + private static SegmentPublishResult doNotRetryOffsetMismatchFailure(SegmentPublishResult result) + { + if (result.isSuccess() || !result.isOffsetMismatch() || !result.isRetryable()) { + return result; + } else { + return SegmentPublishResult.fail(result.getErrorMsg()); + } } } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java index 6d51c59812e9..9a4a09b1961a 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java @@ -517,6 +517,7 @@ public Response getWorkerConfigHistory( @ResourceFilters(StateResourceFilter.class) public void doAction( final TaskActionHolder holder, + @Nullable @PathParam("timeout") final Long timeoutMillis, @Context HttpServletRequest request ) { @@ -536,7 +537,9 @@ public void doAction( event.getAsyncContext().complete(); }) ); - asyncContext.setTimeout(5 * 60_000); + + // Use a default timeout of 5 minutes + asyncContext.setTimeout(Configs.valueOrDefault(timeoutMillis, 5 * 60_000)); Futures.addCallback( future, @@ -545,7 +548,11 @@ public void doAction( @Override public void onSuccess(Object result) { - completeRequest(asyncContext, Status.OK.getStatusCode(), Map.of("result", result)); + completeRequest( + asyncContext, + Status.OK.getStatusCode(), + result == null ? Map.of() : Map.of("result", result) + ); } @Override @@ -919,6 +926,7 @@ private void completeRequest(AsyncContext context, int statusCode, Object result response.setContentType(MediaType.APPLICATION_JSON); jsonMapper.writeValue(response.getOutputStream(), result); } + context.complete(); } catch (Exception e) { log.noStackTrace().warn(e, "Request timed out or is closed already"); diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java index 5090d9fa0894..88d5beb439fa 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java @@ -425,19 +425,19 @@ private class TaskWaitingToPublish { final CountDownLatch latch; final SettableFuture readyToPublish = SettableFuture.create(); - final Map> blockingTaskGroups; + final ConcurrentHashMap> blockingTaskGroups; TaskWaitingToPublish(String taskId, Map> blockingTaskGroups) { this.latch = new CountDownLatch(blockingTaskGroups.size()); - this.blockingTaskGroups = new HashMap<>(blockingTaskGroups); + this.blockingTaskGroups = new ConcurrentHashMap<>(blockingTaskGroups); tasksWaitingToPublish.put(taskId, this); for (TaskGroup group : blockingTaskGroups.keySet()) { group.addCompletionListener(() -> { latch.countDown(); - blockingTaskGroups.remove(group); + this.blockingTaskGroups.remove(group); if (latch.getCount() == 0) { readyToPublish.set(true); diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/common/actions/TaskActionTestKit.java b/indexing-service/src/test/java/org/apache/druid/indexing/common/actions/TaskActionTestKit.java index 6b9c88727ae3..73900ea2c2e3 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/common/actions/TaskActionTestKit.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/common/actions/TaskActionTestKit.java @@ -268,7 +268,8 @@ public boolean isBatchAllocationReduceMetadataIO() segmentAllocationQueue, emitter, supervisorManager, - objectMapper + objectMapper, + ScheduledExecutors::fixed ); taskActionDelegate = new HashMap<>(); testDerbyConnector.createDataSourceTable(); diff --git a/server/src/main/java/org/apache/druid/indexing/overlord/SegmentPublishResult.java b/server/src/main/java/org/apache/druid/indexing/overlord/SegmentPublishResult.java index e19c9dc5f259..c34553a617d5 100644 --- a/server/src/main/java/org/apache/druid/indexing/overlord/SegmentPublishResult.java +++ b/server/src/main/java/org/apache/druid/indexing/overlord/SegmentPublishResult.java @@ -70,20 +70,6 @@ public static SegmentPublishResult retryableOffsetMismatchFailure(String errorMs return new SegmentPublishResult(Set.of(), false, true, StringUtils.format(errorMsg, args), null, true); } - /** - * Converts the given publish result to a non-retryable failure if the publish - * had failed due to offset mismatch. If the task has already waited for older - * tasks before trying to publish its segments, retrying the publish would not make it pass. - */ - public static SegmentPublishResult doNotRetryOffsetMismatchFailure(SegmentPublishResult result) - { - if (result.isSuccess() || !result.isOffsetMismatch() || !result.isRetryable()) { - return result; - } else { - return SegmentPublishResult.fail(result.getErrorMsg()); - } - } - @JsonCreator private SegmentPublishResult( @JsonProperty("segments") Set segments, From d460766126d73793d4016a972ed734f6a53b6d55 Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Fri, 9 Oct 2026 14:22:29 +0530 Subject: [PATCH 07/16] Add first publish attempt --- .../indexing/common/actions/TaskActionToolbox.java | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java index 78e3b154e7ae..73d5650345a5 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java @@ -162,10 +162,20 @@ public ListenableFuture publishSegmentsWhenReady( BiFunction publishAction ) { - if (startMetadata == null || supervisorId == null) { - return Futures.immediateFuture(publishAction.apply(task, this)); + // First try publishing synchronously + try { + final SegmentPublishResult firstAttemptResult = publishAction.apply(task, this); + if (firstAttemptResult.isSuccess() + || !firstAttemptResult.isOffsetMismatch() + || !firstAttemptResult.isRetryable()) { + return Futures.immediateFuture(firstAttemptResult); + } + } + catch (Exception e) { + return Futures.immediateFailedFuture(e); } + // Try publishing later if the failure was due to offset mismatch final ListenableFuture taskReadyToPublishFuture = supervisorManager.isTaskReadyToPublishSegments( supervisorId, task.getId(), From c3bbf37805cc4d36040949db5265c8ac67bf4fe1 Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Fri, 9 Oct 2026 18:48:54 +0530 Subject: [PATCH 08/16] Fixes --- .../common/actions/LocalTaskActionClient.java | 29 +++++++++++-------- .../common/actions/TaskActionClient.java | 2 +- .../common/actions/TaskActionToolbox.java | 26 +++++++---------- .../overlord/http/OverlordResource.java | 28 +++++++++--------- 4 files changed, 43 insertions(+), 42 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/LocalTaskActionClient.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/LocalTaskActionClient.java index e7ff750ae133..49c29fdde12d 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/LocalTaskActionClient.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/LocalTaskActionClient.java @@ -65,18 +65,23 @@ public RetType submit(TaskAction taskAction) @Override public ListenableFuture submitAsync(TaskAction taskAction) { - if (taskAction.canPerformAsync(task, toolbox)) { - final Stopwatch actionRunTime = Stopwatch.createStarted(); - return Futures.transform( - taskAction.performAsync(task, toolbox), - v -> { - emitTimerMetric(RUN_TIME_METRIC, taskAction, actionRunTime.millisElapsed()); - return v; - }, - MoreExecutors.directExecutor() - ); - } else { - return Futures.immediateFuture(submit(taskAction)); + try { + if (taskAction.canPerformAsync(task, toolbox)) { + final Stopwatch actionRunTime = Stopwatch.createStarted(); + return Futures.transform( + taskAction.performAsync(task, toolbox), + v -> { + emitTimerMetric(RUN_TIME_METRIC, taskAction, actionRunTime.millisElapsed()); + return v; + }, + MoreExecutors.directExecutor() + ); + } else { + return Futures.immediateFuture(submit(taskAction)); + } + } + catch (Exception e) { + return Futures.immediateFailedFuture(e); } } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionClient.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionClient.java index a7a8a545e93c..457f14d013f7 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionClient.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionClient.java @@ -36,7 +36,7 @@ default ListenableFuture submitAsync(TaskAction task try { return Futures.immediateFuture(submit(taskAction)); } - catch (IOException e) { + catch (Exception e) { return Futures.immediateFailedFuture(e); } } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java index 73d5650345a5..82639b71df8c 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java @@ -72,6 +72,9 @@ public TaskActionToolbox( this.supervisorManager = supervisorManager; this.jsonMapper = jsonMapper; this.segmentAllocationQueue = segmentAllocationQueue; + + // This executor is currently used only for delayed segment publish actions. + // 4 threads are enough since each publish operation is expected to be fast. this.actionExec = scheduledExecutorFactory.create(4, "TaskActionToolbox-%s"); } @@ -186,8 +189,14 @@ public ListenableFuture publishSegmentsWhenReady( taskReadyToPublishFuture, readyToPublish -> { if (Boolean.TRUE.equals(readyToPublish)) { - // Task is already unblocked for publish, retrying will not fix offset mismatch - return doNotRetryOffsetMismatchFailure(publishAction.apply(task, this)); + final SegmentPublishResult result = publishAction.apply(task, this); + if (result.isOffsetMismatch() && !result.isSuccess() && result.isRetryable()) { + // Do not retry offset mismatch failures since task is already + // unblocked for publish and retrying will not fix the mismatch. + return SegmentPublishResult.fail(result.getErrorMsg()); + } else { + return result; + } } else { return SegmentPublishResult.retryableFailure("Task is not ready to publish yet"); } @@ -195,17 +204,4 @@ public ListenableFuture publishSegmentsWhenReady( actionExec ); } - - /** - * Converts the given publish result to a non-retryable failure if the publish - * had failed due to offset mismatch. - */ - private static SegmentPublishResult doNotRetryOffsetMismatchFailure(SegmentPublishResult result) - { - if (result.isSuccess() || !result.isOffsetMismatch() || !result.isRetryable()) { - return result; - } else { - return SegmentPublishResult.fail(result.getErrorMsg()); - } - } } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java index 9a4a09b1961a..169cdfc64b33 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java @@ -33,6 +33,7 @@ import com.google.common.util.concurrent.MoreExecutors; import com.google.inject.Inject; import com.sun.jersey.spi.container.ResourceFilters; +import io.netty.handler.codec.http.HttpResponseStatus; import org.apache.druid.audit.AuditEntry; import org.apache.druid.audit.AuditManager; import org.apache.druid.client.indexing.ClientTaskQuery; @@ -103,6 +104,7 @@ import javax.ws.rs.core.Response.Status; import java.io.InputStream; import java.util.Collections; +import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -517,7 +519,6 @@ public Response getWorkerConfigHistory( @ResourceFilters(StateResourceFilter.class) public void doAction( final TaskActionHolder holder, - @Nullable @PathParam("timeout") final Long timeoutMillis, @Context HttpServletRequest request ) { @@ -525,21 +526,20 @@ public void doAction( final Optional taskActionClient = taskMaster.getTaskActionClient(holder.getTask()); if (!taskActionClient.isPresent()) { // Encourage client to try again soon, when we'll likely have a redirect set up - completeRequest(asyncContext, Status.SERVICE_UNAVAILABLE.getStatusCode(), null); + completeAsyncRequest(asyncContext, Status.SERVICE_UNAVAILABLE.getStatusCode(), null); return; } final ListenableFuture future = taskActionClient.get().submitAsync(holder.getAction()); asyncContext.addListener( ServletResourceUtils.createAsyncTimeoutListener(event -> { - // HTTP 204 NO_CONTENT is sent to the client. future.cancel(true); - event.getAsyncContext().complete(); + completeAsyncRequest(event.getAsyncContext(), HttpResponseStatus.GATEWAY_TIMEOUT.code(), null); }) ); - // Use a default timeout of 5 minutes - asyncContext.setTimeout(Configs.valueOrDefault(timeoutMillis, 5 * 60_000)); + // Use a default timeout of 15 minutes + asyncContext.setTimeout(15 * 60_000); Futures.addCallback( future, @@ -548,20 +548,20 @@ public void doAction( @Override public void onSuccess(Object result) { - completeRequest( - asyncContext, - Status.OK.getStatusCode(), - result == null ? Map.of() : Map.of("result", result) - ); + // Use null-safe map since some task actions may return null result + final Map payload = new HashMap<>(); + payload.put("result", result); + + completeAsyncRequest(asyncContext, Status.OK.getStatusCode(), payload); } @Override public void onFailure(Throwable t) { if (t instanceof DruidException druidException) { - completeRequest(asyncContext, druidException.getStatusCode(), druidException.toErrorResponse()); + completeAsyncRequest(asyncContext, druidException.getStatusCode(), druidException.toErrorResponse()); } else { - completeRequest(asyncContext, Status.INTERNAL_SERVER_ERROR.getStatusCode(), null); + completeAsyncRequest(asyncContext, Status.INTERNAL_SERVER_ERROR.getStatusCode(), null); } } }, @@ -917,7 +917,7 @@ private List securedTaskStatusPlus( ); } - private void completeRequest(AsyncContext context, int statusCode, Object result) + private void completeAsyncRequest(AsyncContext context, int statusCode, Object result) { try { final HttpServletResponse response = (HttpServletResponse) context.getResponse(); From d0773ad6267212b66811f1daac09ca2ff003ac8c Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Fri, 9 Oct 2026 20:27:50 +0530 Subject: [PATCH 09/16] Add tests for OverlordResource.performTaskAction --- .../overlord/http/OverlordResource.java | 2 +- .../overlord/http/OverlordResourceTest.java | 153 ++++++++++++++++++ 2 files changed, 154 insertions(+), 1 deletion(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java index 169cdfc64b33..40fb288124d4 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java @@ -517,7 +517,7 @@ public Response getWorkerConfigHistory( @Path("/action") @Produces(MediaType.APPLICATION_JSON) @ResourceFilters(StateResourceFilter.class) - public void doAction( + public void performTaskAction( final TaskActionHolder holder, @Context HttpServletRequest request ) diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordResourceTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordResourceTest.java index 74ab86de2d83..346aa0c3da24 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordResourceTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordResourceTest.java @@ -25,6 +25,8 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.SettableFuture; import io.netty.handler.codec.http.HttpResponseStatus; import org.apache.druid.audit.AuditEntry; import org.apache.druid.audit.AuditManager; @@ -40,6 +42,9 @@ import org.apache.druid.indexing.common.TaskLock; import org.apache.druid.indexing.common.TaskLockType; import org.apache.druid.indexing.common.TimeChunkLock; +import org.apache.druid.indexing.common.actions.LockListAction; +import org.apache.druid.indexing.common.actions.TaskActionClient; +import org.apache.druid.indexing.common.actions.TaskActionHolder; import org.apache.druid.indexing.common.task.KillUnusedSegmentsTask; import org.apache.druid.indexing.common.task.NoopTask; import org.apache.druid.indexing.common.task.Task; @@ -55,6 +60,7 @@ import org.apache.druid.indexing.overlord.WorkerTaskRunnerQueryAdapter; import org.apache.druid.indexing.overlord.setup.WorkerBehaviorConfig; import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.Intervals; import org.apache.druid.java.util.common.RE; import org.apache.druid.java.util.common.UOE; @@ -64,6 +70,8 @@ import org.apache.druid.metadata.TaskLookup.CompleteTaskLookup; import org.apache.druid.metadata.TaskLookup.TaskLookupType; import org.apache.druid.segment.TestHelper; +import org.apache.druid.server.mocks.MockAsyncContext; +import org.apache.druid.server.mocks.MockHttpServletResponse; import org.apache.druid.server.security.Access; import org.apache.druid.server.security.Action; import org.apache.druid.server.security.AuthConfig; @@ -86,6 +94,8 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import javax.servlet.AsyncEvent; +import javax.servlet.AsyncListener; import javax.servlet.http.HttpServletRequest; import javax.ws.rs.WebApplicationException; import javax.ws.rs.core.Response; @@ -1477,6 +1487,149 @@ public void testGetTaskSegmentsReturns404() ); } + @Test + public void test_performTaskAction_returnsOk_whenFutureCompletes() + { + final NoopTask task = NoopTask.create(); + final TaskActionHolder holder = new TaskActionHolder(task, new LockListAction()); + final List result = List.of( + new TimeChunkLock(TaskLockType.EXCLUSIVE, "groupId", "ds", Intervals.ETERNITY, "v1", 0) + ); + + final TaskActionClient taskActionClient = EasyMock.createMock(TaskActionClient.class); + EasyMock.expect(taskMaster.getTaskActionClient(task)) + .andReturn(Optional.of(taskActionClient)) + .once(); + EasyMock.expect(taskActionClient.submitAsync(holder.getAction())) + .andReturn(Futures.immediateFuture(result)) + .once(); + + final MockAsyncContext asyncContext = new MockAsyncContext(); + final MockHttpServletResponse asyncResponse = new MockHttpServletResponse(); + asyncContext.response = asyncResponse; + EasyMock.expect(req.startAsync()).andReturn(asyncContext).once(); + + EasyMock.replay(taskActionClient); + replayAll(); + + overlordResource.performTaskAction(holder, req); + + Assertions.assertEquals(Status.OK.getStatusCode(), asyncResponse.getStatus()); + Assertions.assertTrue(asyncContext.isCompleted()); + + EasyMock.verify(taskActionClient); + } + + @Test + public void test_performTaskAction_returnsGatewayTimeout_whenRequestTimesOut() throws Exception + { + final NoopTask task = NoopTask.create(); + final TaskActionHolder holder = new TaskActionHolder(task, new LockListAction()); + + final TaskActionClient taskActionClient = EasyMock.createMock(TaskActionClient.class); + final SettableFuture future = SettableFuture.create(); + + EasyMock.expect(taskMaster.getTaskActionClient(task)) + .andReturn(Optional.of(taskActionClient)) + .once(); + EasyMock.expect(taskActionClient.submitAsync(holder.getAction())) + .andReturn(future) + .once(); + + final Capture listenerCapture = Capture.newInstance(); + final MockAsyncContext asyncContext = new MockAsyncContext() + { + @Override + public void addListener(AsyncListener listener) + { + listenerCapture.setValue(listener); + } + }; + final MockHttpServletResponse asyncResponse = new MockHttpServletResponse(); + asyncContext.response = asyncResponse; + EasyMock.expect(req.startAsync()).andReturn(asyncContext).once(); + + EasyMock.replay(taskActionClient); + replayAll(); + + // Perform the action and then invoke timeout on the async context + overlordResource.performTaskAction(holder, req); + + Assertions.assertTrue(listenerCapture.hasCaptured()); + final AsyncEvent timeoutEvent = EasyMock.createMock(AsyncEvent.class); + EasyMock.expect(timeoutEvent.getAsyncContext()).andReturn(asyncContext).anyTimes(); + EasyMock.replay(timeoutEvent); + + listenerCapture.getValue().onTimeout(timeoutEvent); + + Assertions.assertEquals(HttpResponseStatus.GATEWAY_TIMEOUT.code(), asyncResponse.getStatus()); + Assertions.assertTrue(asyncContext.isCompleted()); + Assertions.assertTrue(future.isCancelled()); + + EasyMock.verify(taskActionClient); + } + + @Test + public void test_performTaskAction_returnsErrorStatusCode_whenFutureFailsWithDruidException() + { + final NoopTask task = NoopTask.create(); + final TaskActionHolder holder = new TaskActionHolder(task, new LockListAction()); + final DruidException druidException = InvalidInput.exception("Task[%s] not found", task.getId()); + + final TaskActionClient taskActionClient = EasyMock.createMock(TaskActionClient.class); + EasyMock.expect(taskMaster.getTaskActionClient(task)) + .andReturn(Optional.of(taskActionClient)) + .once(); + EasyMock.expect(taskActionClient.submitAsync(holder.getAction())) + .andReturn(Futures.immediateFailedFuture(druidException)) + .once(); + + final MockAsyncContext asyncContext = new MockAsyncContext(); + final MockHttpServletResponse asyncResponse = new MockHttpServletResponse(); + asyncContext.response = asyncResponse; + EasyMock.expect(req.startAsync()).andReturn(asyncContext).once(); + + EasyMock.replay(taskActionClient); + replayAll(); + + overlordResource.performTaskAction(holder, req); + + Assertions.assertEquals(druidException.getStatusCode(), asyncResponse.getStatus()); + Assertions.assertTrue(asyncContext.isCompleted()); + + EasyMock.verify(taskActionClient); + } + + @Test + public void test_performTaskAction_returns500_whenFutureFailsWithGenericException() + { + final NoopTask task = NoopTask.create(); + final TaskActionHolder holder = new TaskActionHolder(task, new LockListAction()); + + final TaskActionClient taskActionClient = EasyMock.createMock(TaskActionClient.class); + EasyMock.expect(taskMaster.getTaskActionClient(task)) + .andReturn(Optional.of(taskActionClient)) + .once(); + EasyMock.expect(taskActionClient.submitAsync(holder.getAction())) + .andReturn(Futures.immediateFailedFuture(new ISE("Something went wrong"))) + .once(); + + final MockAsyncContext asyncContext = new MockAsyncContext(); + final MockHttpServletResponse asyncResponse = new MockHttpServletResponse(); + asyncContext.response = asyncResponse; + EasyMock.expect(req.startAsync()).andReturn(asyncContext).once(); + + EasyMock.replay(taskActionClient); + replayAll(); + + overlordResource.performTaskAction(holder, req); + + Assertions.assertEquals(Status.INTERNAL_SERVER_ERROR.getStatusCode(), asyncResponse.getStatus()); + Assertions.assertTrue(asyncContext.isCompleted()); + + EasyMock.verify(taskActionClient); + } + private void expectAuthorizationTokenCheck() { expectAuthorizationTokenCheck(Users.DRUID); From 173427e076462a94e5489f03b1a8b6c2858798ef Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Mon, 12 Oct 2026 17:49:59 +0530 Subject: [PATCH 10/16] Perform first attempt on executor --- .../indexing/kafka/KafkaIndexTaskTest.java | 2 +- .../common/actions/TaskActionToolbox.java | 28 ++++----- .../overlord/http/OverlordResource.java | 57 ++++++++++++------- .../SeekableStreamIndexTaskTestBase.java | 12 +++- 4 files changed, 59 insertions(+), 40 deletions(-) diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java index 5517a2f6773b..3908a0abb22e 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java @@ -2648,7 +2648,7 @@ public void testRunTransactionModeRollback() throws Exception @Test @Timeout(value = 60_000L, unit = TimeUnit.MILLISECONDS) - public void testRunUnTransactionMode() throws Exception + public void testRunInTransactionMode() throws Exception { Map configs = kafkaServer.consumerProperties(); configs.put("isolation.level", "read_uncommitted"); diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java index 82639b71df8c..4e1f0e807e5a 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java @@ -165,25 +165,17 @@ public ListenableFuture publishSegmentsWhenReady( BiFunction publishAction ) { - // First try publishing synchronously - try { - final SegmentPublishResult firstAttemptResult = publishAction.apply(task, this); - if (firstAttemptResult.isSuccess() - || !firstAttemptResult.isOffsetMismatch() - || !firstAttemptResult.isRetryable()) { - return Futures.immediateFuture(firstAttemptResult); - } - } - catch (Exception e) { - return Futures.immediateFailedFuture(e); - } - // Try publishing later if the failure was due to offset mismatch - final ListenableFuture taskReadyToPublishFuture = supervisorManager.isTaskReadyToPublishSegments( - supervisorId, - task.getId(), - startMetadata - ); + final ListenableFuture taskReadyToPublishFuture; + if (startMetadata == null || supervisorId == null) { + taskReadyToPublishFuture = Futures.immediateFuture(true); + } else { + taskReadyToPublishFuture = supervisorManager.isTaskReadyToPublishSegments( + supervisorId, + task.getId(), + startMetadata + ); + } return Futures.transform( taskReadyToPublishFuture, diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java index 40fb288124d4..812f996839d8 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java @@ -109,6 +109,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; /** * @@ -522,24 +523,24 @@ public void performTaskAction( @Context HttpServletRequest request ) { - final AsyncContext asyncContext = request.startAsync(); + final AsyncRequest asyncRequest = new AsyncRequest(request); final Optional taskActionClient = taskMaster.getTaskActionClient(holder.getTask()); if (!taskActionClient.isPresent()) { // Encourage client to try again soon, when we'll likely have a redirect set up - completeAsyncRequest(asyncContext, Status.SERVICE_UNAVAILABLE.getStatusCode(), null); + asyncRequest.complete(Status.SERVICE_UNAVAILABLE.getStatusCode(), null); return; } final ListenableFuture future = taskActionClient.get().submitAsync(holder.getAction()); - asyncContext.addListener( - ServletResourceUtils.createAsyncTimeoutListener(event -> { + asyncRequest.context.addListener( + ServletResourceUtils.createAsyncTimeoutListener(_ -> { + asyncRequest.complete(HttpResponseStatus.GATEWAY_TIMEOUT.code(), null); future.cancel(true); - completeAsyncRequest(event.getAsyncContext(), HttpResponseStatus.GATEWAY_TIMEOUT.code(), null); }) ); // Use a default timeout of 15 minutes - asyncContext.setTimeout(15 * 60_000); + asyncRequest.context.setTimeout(15 * 60_000); Futures.addCallback( future, @@ -552,16 +553,16 @@ public void onSuccess(Object result) final Map payload = new HashMap<>(); payload.put("result", result); - completeAsyncRequest(asyncContext, Status.OK.getStatusCode(), payload); + asyncRequest.complete(Status.OK.getStatusCode(), payload); } @Override public void onFailure(Throwable t) { if (t instanceof DruidException druidException) { - completeAsyncRequest(asyncContext, druidException.getStatusCode(), druidException.toErrorResponse()); + asyncRequest.complete(druidException.getStatusCode(), druidException.toErrorResponse()); } else { - completeAsyncRequest(asyncContext, Status.INTERNAL_SERVER_ERROR.getStatusCode(), null); + asyncRequest.complete(Status.INTERNAL_SERVER_ERROR.getStatusCode(), null); } } }, @@ -917,19 +918,35 @@ private List securedTaskStatusPlus( ); } - private void completeAsyncRequest(AsyncContext context, int statusCode, Object result) + private class AsyncRequest { - try { - final HttpServletResponse response = (HttpServletResponse) context.getResponse(); - response.setStatus(statusCode); - if (result != null) { - response.setContentType(MediaType.APPLICATION_JSON); - jsonMapper.writeValue(response.getOutputStream(), result); - } - context.complete(); + final AtomicBoolean completed = new AtomicBoolean(false); + final AsyncContext context; + + AsyncRequest(HttpServletRequest request) + { + this.context = request.startAsync(); } - catch (Exception e) { - log.noStackTrace().warn(e, "Request timed out or is closed already"); + + void complete(int statusCode, @Nullable Object result) + { + if (!completed.compareAndSet(false, true)) { + log.warn("Request has already been completed"); + return; + } + + try { + final HttpServletResponse response = (HttpServletResponse) context.getResponse(); + response.setStatus(statusCode); + if (result != null) { + response.setContentType(MediaType.APPLICATION_JSON); + jsonMapper.writeValue(response.getOutputStream(), result); + } + context.complete(); + } + catch (Exception e) { + log.noStackTrace().warn(e, "Request timed out or is closed already"); + } } } } diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskTestBase.java b/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskTestBase.java index f0d08b0d941a..55ee58a73817 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskTestBase.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskTestBase.java @@ -27,6 +27,7 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.io.Files; +import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.ListeningExecutorService; import org.apache.druid.client.cache.CacheConfig; @@ -619,7 +620,6 @@ public boolean checkPointDataSourceMetadata( @Nullable DataSourceMetadata previousDataSourceMetadata ) { - // log.info("Adding checkpoint hash to the set"); checkpointRequestsHash.add( Objects.hash( supervisorId, @@ -629,6 +629,16 @@ public boolean checkPointDataSourceMetadata( ); return true; } + + @Override + public ListenableFuture isTaskReadyToPublishSegments( + String supervisorId, + String taskId, + DataSourceMetadata startMetadata + ) + { + return Futures.immediateFuture(true); + } }, objectMapper ); From 5d01f3166e121f6c8564d4e21b63a1cd9035fdbf Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Tue, 13 Oct 2026 16:24:27 +0530 Subject: [PATCH 11/16] Fix flakiness in KafkaIndexTaskTest --- .../org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java index 3908a0abb22e..be1082629a0e 100644 --- a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java +++ b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java @@ -2689,6 +2689,11 @@ public void testRunInTransactionMode() throws Exception } Assertions.assertEquals(2, countEvents(task)); + + task.getRunner().stopGracefully(); + + // Wait for task to exit + Assertions.assertEquals(TaskState.SUCCESS, future.get().getStatusCode()); } @Test From 650e459b23dd0fc2853d186a394628ed0e8fcc82 Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Tue, 13 Oct 2026 17:14:48 +0530 Subject: [PATCH 12/16] Add tests for TaskActionToolbox --- .../common/actions/TaskActionToolbox.java | 2 +- .../common/actions/TaskActionToolboxTest.java | 171 +++++++++++++++++- 2 files changed, 171 insertions(+), 2 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java index 4e1f0e807e5a..8f19a3b8c77f 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java @@ -185,7 +185,7 @@ public ListenableFuture publishSegmentsWhenReady( if (result.isOffsetMismatch() && !result.isSuccess() && result.isRetryable()) { // Do not retry offset mismatch failures since task is already // unblocked for publish and retrying will not fix the mismatch. - return SegmentPublishResult.fail(result.getErrorMsg()); + return SegmentPublishResult.fail("Non-retryable metadata mismatch error: " + result.getErrorMsg()); } else { return result; } diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/common/actions/TaskActionToolboxTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/common/actions/TaskActionToolboxTest.java index 8df71215f368..e75ada0721fd 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/common/actions/TaskActionToolboxTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/common/actions/TaskActionToolboxTest.java @@ -19,20 +19,61 @@ package org.apache.druid.indexing.common.actions; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; +import org.apache.druid.common.guava.FutureUtils; +import org.apache.druid.indexing.common.task.NoopTask; +import org.apache.druid.indexing.common.task.Task; +import org.apache.druid.indexing.overlord.DataSourceMetadata; import org.apache.druid.indexing.overlord.ForkingTaskRunner; import org.apache.druid.indexing.overlord.ForkingTaskRunnerFactory; +import org.apache.druid.indexing.overlord.ObjectMetadata; +import org.apache.druid.indexing.overlord.SegmentPublishResult; +import org.apache.druid.indexing.overlord.supervisor.SupervisorManager; +import org.apache.druid.server.coordinator.simulate.BlockingExecutorService; +import org.apache.druid.server.coordinator.simulate.WrappingScheduledExecutorService; +import org.apache.druid.timeline.DataSegment; +import org.apache.druid.timeline.SegmentId; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import java.util.Set; + +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; public class TaskActionToolboxTest { + private TaskActionToolbox toolbox; + private SupervisorManager supervisorManager; + private BlockingExecutorService taskActionExec; + + @BeforeEach + public void setup() + { + taskActionExec = new BlockingExecutorService("test"); + supervisorManager = mock(SupervisorManager.class); + toolbox = new TaskActionToolbox( + null, + null, + null, + null, + null, + supervisorManager, + null, + (_, nameFormat) -> new WrappingScheduledExecutorService(nameFormat, taskActionExec, true) + ); + } @Test - public void testMakeCodeCoverageHappy() + public void test_setTaskRunnerFactory() { TaskActionToolbox toolbox = new TaskActionToolbox(null, null, null, null, null, null); assertFalse(toolbox.getTaskRunner().isPresent()); @@ -41,4 +82,132 @@ public void testMakeCodeCoverageHappy() toolbox.setTaskRunnerFactory(factory); assertTrue(toolbox.getTaskRunner().isPresent()); } + + @Test + public void test_publishSegmentsWhenReady_waits_untilTaskIsReady() + { + final Task task = NoopTask.create(); + final String supervisorId = "supervisor1"; + final DataSourceMetadata startMetadata = new ObjectMetadata("abc"); + + final SettableFuture taskReadyFuture = SettableFuture.create(); + + when(supervisorManager.isTaskReadyToPublishSegments(supervisorId, task.getId(), startMetadata)) + .thenReturn(taskReadyFuture); + + final Set segmentsToPublish = Set.of(DataSegment.builder(SegmentId.dummy("wiki")).build()); + final ListenableFuture publishFuture = toolbox.publishSegmentsWhenReady( + task, + supervisorId, + startMetadata, + (_, _) -> SegmentPublishResult.ok(segmentsToPublish) + ); + + assertFalse(taskReadyFuture.isDone()); + assertFalse(publishFuture.isDone()); + + // Mark the task as ready to publish and verify that publish proceeds + taskReadyFuture.set(true); + taskActionExec.finishAllPendingTasks(); + assertTrue(publishFuture.isDone()); + assertTrue(taskReadyFuture.isDone()); + assertEquals(segmentsToPublish, getResult(publishFuture).getSegments()); + } + + @Test + public void test_publishSegmentsWhenReady_proceedsImmediately_ifSupervisorIdIsNull() + { + final Task task = NoopTask.create(); + + final Set segmentsToPublish = Set.of(DataSegment.builder(SegmentId.dummy("wiki")).build()); + final ListenableFuture publishFuture = toolbox.publishSegmentsWhenReady( + task, + null, + null, + (_, _) -> SegmentPublishResult.ok(segmentsToPublish) + ); + + assertFalse(publishFuture.isDone()); + + // Verify that the publish proceeds immediately + taskActionExec.finishAllPendingTasks(); + assertTrue(publishFuture.isDone()); + assertEquals(segmentsToPublish, getResult(publishFuture).getSegments()); + + verify(supervisorManager, never()).isTaskReadyToPublishSegments(any(), eq(task.getId()), any()); + } + + @Test + public void test_publishSegmentsWhenReady_returnsNonRetryableFailure_ifOffsetMismatches() + { + final Task task = NoopTask.create(); + final String supervisorId = "supervisor1"; + final DataSourceMetadata startMetadata = new ObjectMetadata("abc"); + + final SettableFuture taskReadyFuture = SettableFuture.create(); + + when(supervisorManager.isTaskReadyToPublishSegments(supervisorId, task.getId(), startMetadata)) + .thenReturn(taskReadyFuture); + + final ListenableFuture publishFuture = toolbox.publishSegmentsWhenReady( + task, + supervisorId, + startMetadata, + (_, _) -> SegmentPublishResult.retryableOffsetMismatchFailure("offset mismatch") + ); + + assertFalse(taskReadyFuture.isDone()); + assertFalse(publishFuture.isDone()); + + // Mark the task as ready to publish and verify that publish proceeds + taskReadyFuture.set(true); + taskActionExec.finishAllPendingTasks(); + assertTrue(publishFuture.isDone()); + assertTrue(taskReadyFuture.isDone()); + + // Verify that the publish failure is non retryable + final SegmentPublishResult result = getResult(publishFuture); + assertFalse(result.isRetryable()); + assertEquals("Non-retryable metadata mismatch error: offset mismatch", result.getErrorMsg()); + } + + @Test + public void test_publishSegmentsWhenReady_returnsRetryableFailure_ifTaskCouldNotBecomeReady() + { + final Task task = NoopTask.create(); + final String supervisorId = "supervisor1"; + final DataSourceMetadata startMetadata = new ObjectMetadata("abc"); + + final SettableFuture taskReadyFuture = SettableFuture.create(); + + when(supervisorManager.isTaskReadyToPublishSegments(supervisorId, task.getId(), startMetadata)) + .thenReturn(taskReadyFuture); + + final Set segmentsToPublish = Set.of(DataSegment.builder(SegmentId.dummy("wiki")).build()); + final ListenableFuture publishFuture = toolbox.publishSegmentsWhenReady( + task, + supervisorId, + startMetadata, + (_, _) -> SegmentPublishResult.ok(segmentsToPublish) + ); + + assertFalse(taskReadyFuture.isDone()); + assertFalse(publishFuture.isDone()); + + // Mark the task as not ready to publish and verify that publish fails + taskReadyFuture.set(false); + taskActionExec.finishAllPendingTasks(); + assertTrue(publishFuture.isDone()); + assertTrue(taskReadyFuture.isDone()); + + assertEquals( + SegmentPublishResult.retryableFailure("Task is not ready to publish yet"), + getResult(publishFuture) + ); + } + + private SegmentPublishResult getResult(ListenableFuture publishFuture) + { + return FutureUtils.getUnchecked(publishFuture, true); + } } From 05dd75df86e8e8744dc3d2a3a0d1efb41d391307 Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Tue, 13 Oct 2026 17:55:00 +0530 Subject: [PATCH 13/16] Add tests for SeekableStreamSupervisor --- .../supervisor/SeekableStreamSupervisor.java | 2 +- .../SeekableStreamSupervisorStateTest.java | 54 +++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java index 88d5beb439fa..aaa770026d69 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java @@ -366,7 +366,7 @@ public String getBaseSequenceName() /** * Called when this TaskGroup is removed from {@link #pendingCompletionTaskGroups}. */ - private void onCompleted() + public void onCompleted() { if (completed.compareAndSet(false, true)) { runCompletionListeners(); diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorStateTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorStateTest.java index d7634a842db4..f394c40a4390 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorStateTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorStateTest.java @@ -29,6 +29,7 @@ import com.google.common.collect.ImmutableSet; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; +import org.apache.druid.common.guava.FutureUtils; import org.apache.druid.data.input.impl.ByteEntity; import org.apache.druid.data.input.impl.DimensionSchema; import org.apache.druid.data.input.impl.JsonInputFormat; @@ -510,6 +511,59 @@ public void testAddDiscoveredTaskToPendingCompletionMultipleTaskGroups() throws Assertions.assertEquals(3, taskGroups.get(1).tasks.size()); } + @Test + public void test_isTaskReadyToPublishSegments_returnsCompleted_ifNoConflictingTask() + { + EasyMock.expect(spec.isSuspended()).andReturn(false).anyTimes(); + EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasksByDatasource(DATASOURCE)).andReturn(List.of()).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); + + replayAll(); + + final TestSeekableStreamSupervisor supervisor = new TestSeekableStreamSupervisor(); + + final ListenableFuture future = supervisor.isTaskReadyToPublishSegments( + "task1", + Set.of("partition") + ); + + Assertions.assertTrue(future.isDone()); + Assertions.assertEquals(Boolean.TRUE, FutureUtils.getUnchecked(future, true)); + } + + @Test + public void test_isTaskReadyToPublishSegments_waitsForConflictingTask() + { + EasyMock.expect(spec.isSuspended()).andReturn(false).anyTimes(); + EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasksByDatasource(DATASOURCE)).andReturn(List.of()).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); + + replayAll(); + + final TestSeekableStreamSupervisor supervisor = new TestSeekableStreamSupervisor(); + + final Map startingPartitions = new HashMap<>(); + startingPartitions.put("partition", "offset"); + + supervisor.addDiscoveredTaskToPendingCompletionTaskGroups(0, "taskPendingCompletion", startingPartitions); + + final ListenableFuture future = supervisor.isTaskReadyToPublishSegments( + "waitingTask", + Set.of("partition") + ); + Assertions.assertFalse(future.isDone()); + + // Mark the earlier task group as completed + final SeekableStreamSupervisor.TaskGroup conflictingTaskGroup = + supervisor.getPendingCompletionTaskGroups(0).get(0); + conflictingTaskGroup.onCompleted(); + + Assertions.assertTrue(future.isDone()); + Assertions.assertEquals(Boolean.TRUE, FutureUtils.getUnchecked(future, true)); + } + @Test public void testConnectingToStreamFail() { From 7fe22363dacaf516a8df62c4c25dd447b9676f6d Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Tue, 13 Oct 2026 19:06:15 +0530 Subject: [PATCH 14/16] Fix response in case of generic error --- .../indexing/overlord/http/OverlordResource.java | 4 +++- .../indexing/overlord/http/OverlordResourceTest.java | 11 ++++++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java index 812f996839d8..d851cc874077 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java @@ -562,7 +562,9 @@ public void onFailure(Throwable t) if (t instanceof DruidException druidException) { asyncRequest.complete(druidException.getStatusCode(), druidException.toErrorResponse()); } else { - asyncRequest.complete(Status.INTERNAL_SERVER_ERROR.getStatusCode(), null); + final Map payload = new HashMap<>(); + payload.put("error", t.getMessage()); + asyncRequest.complete(Status.INTERNAL_SERVER_ERROR.getStatusCode(), payload); } } }, diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordResourceTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordResourceTest.java index 346aa0c3da24..ae1f50cc044d 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordResourceTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordResourceTest.java @@ -1601,7 +1601,7 @@ public void test_performTaskAction_returnsErrorStatusCode_whenFutureFailsWithDru } @Test - public void test_performTaskAction_returns500_whenFutureFailsWithGenericException() + public void test_performTaskAction_returns500_whenFutureFailsWithGenericException() throws Exception { final NoopTask task = NoopTask.create(); final TaskActionHolder holder = new TaskActionHolder(task, new LockListAction()); @@ -1627,6 +1627,15 @@ public void test_performTaskAction_returns500_whenFutureFailsWithGenericExceptio Assertions.assertEquals(Status.INTERNAL_SERVER_ERROR.getStatusCode(), asyncResponse.getStatus()); Assertions.assertTrue(asyncContext.isCompleted()); + final Map responsePayload = TestHelper.JSON_MAPPER.readValue( + asyncResponse.baos.toByteArray(), + new TypeReference<>() {} + ); + Assertions.assertEquals( + Map.of("error", "Something went wrong"), + responsePayload + ); + EasyMock.verify(taskActionClient); } From 4ef4e4442e0578f117201ee3100b10d50239c306 Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Tue, 13 Oct 2026 21:16:03 +0530 Subject: [PATCH 15/16] Wait on older task groups only --- .../overlord/http/OverlordResource.java | 1 + .../supervisor/SupervisorManager.java | 8 +-- .../supervisor/SeekableStreamSupervisor.java | 59 ++++++++++++------- .../supervisor/SupervisorManagerTest.java | 8 +-- .../SeekableStreamSupervisorStateTest.java | 47 +++++++++++---- 5 files changed, 81 insertions(+), 42 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java index d851cc874077..6c2efddd746e 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordResource.java @@ -559,6 +559,7 @@ public void onSuccess(Object result) @Override public void onFailure(Throwable t) { + log.warn(t, "Failed to perform task action"); if (t instanceof DruidException druidException) { asyncRequest.complete(druidException.getStatusCode(), druidException.toErrorResponse()); } else { diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java index fabae3a48247..ffec91c56d71 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManager.java @@ -626,14 +626,8 @@ public ListenableFuture isTaskReadyToPublishSegments( ); } - final Set partitionIds = Set.copyOf( - ((SeekableStreamDataSourceMetadata) startMetadata) - .getSeekableStreamSequenceNumbers() - .getPartitionSequenceNumberMap() - .keySet() - ); return ((SeekableStreamSupervisor) supervisor.lhs) - .isTaskReadyToPublishSegments(taskId, partitionIds); + .isTaskReadyToPublishSegments(taskId, startMetadata); } /** diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java index aaa770026d69..962679fecbde 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java @@ -2932,7 +2932,7 @@ private void verifyAndMergeCheckpoints( // OR if there is an older publishing task for another groupId, use the first sequenceId of this task // since the offsets persisted in metadata store might still be behind the checkpoints of this task - || (isAnotherTaskGroupPublishingToPartitions(taskId, Set.copyOf(taskGroup.startingSequences.keySet())) + || (isOlderTaskGroupPublishingToPartitions(taskId, Map.copyOf(taskGroup.startingSequences)) && earliestConsistentSequenceId.compareAndSet(-1, taskCheckpoints.firstKey())) ) { final SortedMap> latestCheckpoints = new TreeMap<>( @@ -2995,29 +2995,30 @@ private void verifyAndMergeCheckpoints( } /** - * Checks if there is a Task distinct from the given {@code taskId} or its replicas + * Checks if there is an older task group (which does not contain {@code taskId}) * publishing to any of the given partitions. If this method returns a non-empty map,it - * indicates that the current task would need to wait for the other tasks to + * indicates that the current task would need to wait for the older tasks to * finish publishing before it can publish its own offsets. * * @return Map containing the conflicting TaskGroups that are publishing to * overlapping partitions. */ - private Map> getOtherTaskGroupsPublishingToPartitions( + private Map> getOlderTaskGroupsPublishingToPartitions( String taskId, - Set partitions + Map taskStartOffsets ) { // Identify all the partitions that are being published by other taskGroups - final Map> partitionIdToPublishingGroups = new HashMap<>(); + final Map> partitionIdToPublishingGroups = new HashMap<>(); pendingCompletionTaskGroups.values().stream().flatMap(Collection::stream).forEach(group -> { // Do not consider this taskId or its replicas if (group.taskIds().contains(taskId)) { return; } - for (PartitionIdType partitionId : group.startingSequences.keySet()) { - partitionIdToPublishingGroups.computeIfAbsent(partitionId, p -> new HashSet<>()).add(group); - } + group.startingSequences.forEach( + (partitionId, offset) -> partitionIdToPublishingGroups + .computeIfAbsent(partitionId, p -> new HashMap<>()).put(group, offset) + ); }); if (partitionIdToPublishingGroups.isEmpty()) { @@ -3026,26 +3027,29 @@ private Map> getOtherTaskGroupsPublishingToParti // Check if any of the partitions of this taskId are being published by other taskGroups final Map> blockingTaskGroups = new HashMap<>(); - for (Object partition : partitions) { - @SuppressWarnings("unchecked") - final PartitionIdType partitionId = (PartitionIdType) partition; + for (PartitionIdType partitionId : taskStartOffsets.keySet()) { if (partitionIdToPublishingGroups.containsKey(partitionId)) { - for (TaskGroup taskGroup : partitionIdToPublishingGroups.get(partitionId)) { - blockingTaskGroups.computeIfAbsent(taskGroup, _ -> new HashSet<>()) - .add(partitionId); - } + final SequenceOffsetType taskOffset = taskStartOffsets.get(partitionId); + partitionIdToPublishingGroups.get(partitionId).forEach((taskGroup, pendingGroupOffset) -> { + // Consider only task groups which are reading from offsets strictly + // smaller than the candidate taskId + if (!isOffsetAtOrBeyond(pendingGroupOffset, taskOffset)) { + blockingTaskGroups.computeIfAbsent(taskGroup, _ -> new HashSet<>()) + .add(partitionId); + } + }); } } return blockingTaskGroups; } - private boolean isAnotherTaskGroupPublishingToPartitions( + private boolean isOlderTaskGroupPublishingToPartitions( String taskId, - Set partitions + Map taskStartOffsets ) { - return !getOtherTaskGroupsPublishingToPartitions(taskId, partitions).isEmpty(); + return !getOlderTaskGroupsPublishingToPartitions(taskId, taskStartOffsets).isEmpty(); } /** @@ -3056,10 +3060,23 @@ private boolean isAnotherTaskGroupPublishingToPartitions( * its segments since all previous tasks that were reading from an overlapping * set of partitions have already published their offsets. */ - public ListenableFuture isTaskReadyToPublishSegments(String taskId, Set partitions) + public ListenableFuture isTaskReadyToPublishSegments(String taskId, DataSourceMetadata startMetadata) { + if (!(startMetadata instanceof SeekableStreamDataSourceMetadata streamStartMetadata)) { + throw InvalidInput.exception( + "Start metadata[%s] of type[%s] is not valid streaming metadata", + startMetadata, startMetadata == null ? null : startMetadata.getClass() + ); + } + + @SuppressWarnings("unchecked") + final Map taskStartOffsets = + (Map) streamStartMetadata + .getSeekableStreamSequenceNumbers() + .getPartitionSequenceNumberMap(); + final Map> blockingTaskGroups = - getOtherTaskGroupsPublishingToPartitions(taskId, partitions); + getOlderTaskGroupsPublishingToPartitions(taskId, taskStartOffsets); // No tasks writing to overlapping partitions, task may proceed with publish if (blockingTaskGroups.isEmpty()) { diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java index a9688d79dfc2..bec5e1850748 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/supervisor/SupervisorManagerTest.java @@ -1148,7 +1148,7 @@ public void testRegisterUpgradedPendingSegmentOnSupervisor() } @Test - public void test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifSupervisorIdIsNull() + public void test_isTaskReadyToPublishSegments_throwsException_ifSupervisorIdIsNull() { final DruidException exception = Assertions.assertThrows( DruidException.class, @@ -1161,7 +1161,7 @@ public void test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifSupe } @Test - public void test_isTaskReadyToPublish_Segments_throwsException_ifSupervisorNotFound() + public void test_isTaskReadyToPublishSegments_throwsException_ifSupervisorNotFound() { final DruidException exception = Assertions.assertThrows( DruidException.class, @@ -1194,7 +1194,7 @@ public void test_isTaskReadyToPublishSegments_returnsCompleteFuture_forNonStream } @Test - public void test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifMetadataIsNull() + public void test_isTaskReadyToPublishSegments_throwsException_ifMetadataIsNull() { final String supervisorId = "supervisor1"; final SeekableStreamSupervisor seekableStreamSupervisor = @@ -1224,7 +1224,7 @@ public void test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifMeta } @Test - public void test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifMetadataIsInvalid() + public void test_isTaskReadyToPublishSegments_throwsException_ifMetadataIsInvalid() { final String supervisorId = "supervisor1"; final SeekableStreamSupervisor seekableStreamSupervisor = diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorStateTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorStateTest.java index f394c40a4390..912455f8ba2f 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorStateTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisorStateTest.java @@ -44,6 +44,7 @@ import org.apache.druid.indexing.common.task.Task; import org.apache.druid.indexing.overlord.DataSourceMetadata; import org.apache.druid.indexing.overlord.IndexerMetadataStorageCoordinator; +import org.apache.druid.indexing.overlord.ObjectMetadata; import org.apache.druid.indexing.overlord.TaskMaster; import org.apache.druid.indexing.overlord.TaskQueue; import org.apache.druid.indexing.overlord.TaskRunner; @@ -511,6 +512,26 @@ public void testAddDiscoveredTaskToPendingCompletionMultipleTaskGroups() throws Assertions.assertEquals(3, taskGroups.get(1).tasks.size()); } + @Test + public void test_isTaskReadyToPublishSegments_throwsException_ifMetadataIsInvalid() + { + EasyMock.expect(spec.isSuspended()).andReturn(false).anyTimes(); + EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).anyTimes(); + EasyMock.expect(taskStorage.getActiveTasksByDatasource(DATASOURCE)).andReturn(List.of()).anyTimes(); + EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes(); + + replayAll(); + + final TestSeekableStreamSupervisor supervisor = new TestSeekableStreamSupervisor(); + Assertions.assertThrows( + DruidException.class, + () -> supervisor.isTaskReadyToPublishSegments( + "task1", + new ObjectMetadata("abc") + ) + ); + } + @Test public void test_isTaskReadyToPublishSegments_returnsCompleted_ifNoConflictingTask() { @@ -525,7 +546,9 @@ public void test_isTaskReadyToPublishSegments_returnsCompleted_ifNoConflictingTa final ListenableFuture future = supervisor.isTaskReadyToPublishSegments( "task1", - Set.of("partition") + new TestSeekableStreamDataSourceMetadata( + new SeekableStreamStartSequenceNumbers<>(STREAM, Map.of("p1", "100"), Set.of()) + ) ); Assertions.assertTrue(future.isDone()); @@ -533,7 +556,7 @@ public void test_isTaskReadyToPublishSegments_returnsCompleted_ifNoConflictingTa } @Test - public void test_isTaskReadyToPublishSegments_waitsForConflictingTask() + public void test_isTaskReadyToPublishSegments_waitsForEarlierTask() { EasyMock.expect(spec.isSuspended()).andReturn(false).anyTimes(); EasyMock.expect(recordSupplier.getPartitionIds(STREAM)).andReturn(ImmutableSet.of(SHARD_ID)).anyTimes(); @@ -543,25 +566,29 @@ public void test_isTaskReadyToPublishSegments_waitsForConflictingTask() replayAll(); final TestSeekableStreamSupervisor supervisor = new TestSeekableStreamSupervisor(); - - final Map startingPartitions = new HashMap<>(); - startingPartitions.put("partition", "offset"); - - supervisor.addDiscoveredTaskToPendingCompletionTaskGroups(0, "taskPendingCompletion", startingPartitions); + supervisor.addDiscoveredTaskToPendingCompletionTaskGroups(0, "earlierGroupPendingCompletion", Map.of("p1", "0")); + supervisor.addDiscoveredTaskToPendingCompletionTaskGroups(1, "laterGroupPendingCompletion", Map.of("p1", "10000")); final ListenableFuture future = supervisor.isTaskReadyToPublishSegments( "waitingTask", - Set.of("partition") + new TestSeekableStreamDataSourceMetadata( + new SeekableStreamStartSequenceNumbers<>(STREAM, Map.of("p1", "100"), Set.of()) + ) ); Assertions.assertFalse(future.isDone()); // Mark the earlier task group as completed - final SeekableStreamSupervisor.TaskGroup conflictingTaskGroup = + final SeekableStreamSupervisor.TaskGroup earlierTaskGroup = supervisor.getPendingCompletionTaskGroups(0).get(0); - conflictingTaskGroup.onCompleted(); + earlierTaskGroup.onCompleted(); + // Verify that the task is now ready to publish even though the later task group is not complete yet Assertions.assertTrue(future.isDone()); Assertions.assertEquals(Boolean.TRUE, FutureUtils.getUnchecked(future, true)); + + final SeekableStreamSupervisor.TaskGroup laterTaskGroup = + supervisor.getPendingCompletionTaskGroups(1).get(0); + Assertions.assertFalse(laterTaskGroup.completed.get()); } @Test From 93a75b5f95a7560c8c4280511b7ffcdfae3587d4 Mon Sep 17 00:00:00 2001 From: Kashif Faraz Date: Wed, 14 Oct 2026 09:56:22 +0530 Subject: [PATCH 16/16] Use entrySet instead of keySet --- .../supervisor/SeekableStreamSupervisor.java | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java index 962679fecbde..472fa01bfe20 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java @@ -3027,19 +3027,18 @@ private Map> getOlderTaskGroupsPublishingToParti // Check if any of the partitions of this taskId are being published by other taskGroups final Map> blockingTaskGroups = new HashMap<>(); - for (PartitionIdType partitionId : taskStartOffsets.keySet()) { + taskStartOffsets.forEach((partitionId, taskOffset) -> { if (partitionIdToPublishingGroups.containsKey(partitionId)) { - final SequenceOffsetType taskOffset = taskStartOffsets.get(partitionId); partitionIdToPublishingGroups.get(partitionId).forEach((taskGroup, pendingGroupOffset) -> { - // Consider only task groups which are reading from offsets strictly - // smaller than the candidate taskId + // Consider only task groups which are reading from offsets + // strictly smaller than the candidate taskId if (!isOffsetAtOrBeyond(pendingGroupOffset, taskOffset)) { blockingTaskGroups.computeIfAbsent(taskGroup, _ -> new HashSet<>()) .add(partitionId); } }); } - } + }); return blockingTaskGroups; }