This is an automated email from the ASF dual-hosted git repository.
kfaraz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new 7c89096bd1f feat: Block segment publish action for a streaming task
until prior task groups are complete (#20505)
7c89096bd1f is described below
commit 7c89096bd1faefb14a4efef1cbd684719bf2aa83
Author: Kashif Faraz <[email protected]>
AuthorDate: Wed Oct 14 13:57:41 2026 +0530
feat: Block segment publish action for a streaming task until prior task
groups are complete (#20505)
The changes in #19091 have shown that tasks often need to wait for their
priors
to finish publishing as we often see logs like:
```
Task[xyz] needs to wait before publishing as other taskGroups[abc, def]
are currently publishing to partition[123].
```
Before this patch, if publish from streaming task A is slow, a later
generation task B keeps
trying to publish repeatedly until eventually failing once the number of
retries is exhausted.
It is difficult to fix the number of retries since publish times may vary
across supervisors
and over time. It is also pointless to retry if a prior task is yet to
finish publishing.
This patch addresses these problems by blocking the publish of later task B
until A is complete.
Changes
---------
- Add `TaskActionClient.submitAsync` which returns a `ListenableFuture`
- Update API `/druid/indexer/v1/action` in `OverlordResource` to be async
by employing
the new `submitAsync()` method
- Update task actions `segmentTransactionalAppend` and
`segmentTransactionalReplace`
to implement `performAsync()`
- Update `LocalTaskActionClient.submitAsync()` to use `performAsync()` when
supported
- Attempt segment publish only when the task is ready to publish
- Complete async API request only when the publish attempt is done
- For non-streaming tasks, attempt to publish immediately
- Add method `isTaskReadyToPublish()` in `SeekableStreamSupervisor`.
- Return a `ListenableFuture` that completes when all prior task groups
have completed
- Update `SegmentAllocateAction` and `SegmentAllocationQueue` to use
`ListenableFuture`
---
.../druid/indexing/kafka/KafkaIndexTaskTest.java | 7 +-
.../common/actions/LocalTaskActionClient.java | 31 +++-
.../common/actions/SegmentAllocateAction.java | 4 +-
.../common/actions/SegmentAllocationQueue.java | 22 +--
.../actions/SegmentTransactionalAppendAction.java | 18 ++-
.../actions/SegmentTransactionalInsertAction.java | 18 ++-
.../druid/indexing/common/actions/TaskAction.java | 8 +-
.../indexing/common/actions/TaskActionClient.java | 16 ++
.../indexing/common/actions/TaskActionToolbox.java | 63 ++++++--
.../indexing/overlord/http/OverlordResource.java | 110 ++++++++++---
.../overlord/supervisor/SupervisorManager.java | 34 ++--
.../supervisor/SeekableStreamSupervisor.java | 176 ++++++++++++++++++---
.../SegmentTransactionalInsertActionTest.java | 3 +-
.../indexing/common/actions/TaskActionTestKit.java | 3 +-
.../common/actions/TaskActionToolboxTest.java | 171 +++++++++++++++++++-
.../overlord/http/OverlordResourceTest.java | 167 ++++++++++++++++++-
.../druid/indexing/overlord/http/OverlordTest.java | 4 +-
.../overlord/supervisor/SupervisorManagerTest.java | 50 +++---
.../SeekableStreamIndexTaskTestBase.java | 12 +-
.../SeekableStreamSupervisorStateTest.java | 81 ++++++++++
.../indexing/overlord/SegmentPublishResult.java | 27 +++-
.../IndexerSQLMetadataStorageCoordinator.java | 2 +-
22 files changed, 880 insertions(+), 147 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 5517a2f6773..be1082629a0 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 class KafkaIndexTaskTest extends
SeekableStreamIndexTaskTestBase
@Test
@Timeout(value = 60_000L, unit = TimeUnit.MILLISECONDS)
- public void testRunUnTransactionMode() throws Exception
+ public void testRunInTransactionMode() throws Exception
{
Map<String, Object> configs = kafkaServer.consumerProperties();
configs.put("isolation.level", "read_uncommitted");
@@ -2689,6 +2689,11 @@ public class KafkaIndexTaskTest extends
SeekableStreamIndexTaskTestBase
}
Assertions.assertEquals(2, countEvents(task));
+
+ task.getRunner().stopGracefully();
+
+ // Wait for task to exit
+ Assertions.assertEquals(TaskState.SUCCESS, future.get().getStatusCode());
}
@Test
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 1d0059335ed..49c29fdde12 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,33 @@ public class LocalTaskActionClient implements
TaskActionClient
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 <RetType> ListenableFuture<RetType> submitAsync(TaskAction<RetType>
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);
+ }
+ }
+
private <R> R performAction(TaskAction<R> 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 3f240e0aa96..859d60bdd16 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.JsonCreator;
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;
@@ -50,7 +51,6 @@ import javax.annotation.Nullable;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
-import java.util.concurrent.Future;
import java.util.concurrent.ThreadLocalRandom;
/**
@@ -194,7 +194,7 @@ public class SegmentAllocateAction implements
TaskAction<SegmentIdWithShardSpec>
}
@Override
- public Future<SegmentIdWithShardSpec> performAsync(Task task,
TaskActionToolbox toolbox)
+ public ListenableFuture<SegmentIdWithShardSpec> 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 7b504a336b4..7606bde2ff9 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.Objects;
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 class SegmentAllocationQueue
* Queues a SegmentAllocateRequest. The returned future may complete
successfully
* with a non-null value or with a non-null value.
*/
- public Future<SegmentIdWithShardSpec> add(SegmentAllocateRequest request)
+ public ListenableFuture<SegmentIdWithShardSpec> add(SegmentAllocateRequest
request)
{
if (!isLeader.get()) {
throw new ISE("Cannot allocate segment if not leader.");
@@ -220,7 +220,7 @@ public class SegmentAllocationQueue
}
final AllocateRequestKey requestKey = new AllocateRequestKey(request);
- final AtomicReference<Future<SegmentIdWithShardSpec>> futureReference =
new AtomicReference<>();
+ final AtomicReference<ListenableFuture<SegmentIdWithShardSpec>>
futureReference = new AtomicReference<>();
// Possible race condition:
// t1 -> new batch is added to queue or batch already exists in queue
@@ -644,7 +644,7 @@ public class SegmentAllocationQueue
* Map from allocate requests (represents a single SegmentAllocateAction)
* to the future of allocated segment id.
*/
- private final Map<SegmentAllocateRequest,
CompletableFuture<SegmentIdWithShardSpec>>
+ private final Map<SegmentAllocateRequest,
SettableFuture<SegmentIdWithShardSpec>>
requestToFuture = new HashMap<>();
AllocateRequestBatch(AllocateRequestKey key)
@@ -692,10 +692,10 @@ public class SegmentAllocationQueue
return size() >= MAX_BATCH_SIZE;
}
- Future<SegmentIdWithShardSpec> add(SegmentAllocateRequest request)
+ ListenableFuture<SegmentIdWithShardSpec> 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 @@ public class SegmentAllocationQueue
{
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 @@ public class SegmentAllocationQueue
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 @@ public class SegmentAllocationQueue
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 @@ public class SegmentAllocationQueue
+ " 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 15001fcd4a2..77a62cbd3e9 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,7 @@ package org.apache.druid.indexing.common.actions;
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.ListenableFuture;
import org.apache.druid.error.DruidException;
import org.apache.druid.error.InvalidInput;
import org.apache.druid.indexing.common.TaskLock;
@@ -219,12 +220,19 @@ public class SegmentTransactionalAppendAction implements
TaskAction<SegmentPubli
}
IndexTaskUtils.emitSegmentPublishMetrics(retVal, task, toolbox);
+ return retVal;
+ }
- if (toolbox.shouldFailSegmentPublishImmediately(retVal, task,
supervisorId, startMetadata)) {
- return SegmentPublishResult.fail(retVal.getErrorMsg());
- } else {
- return retVal;
- }
+ @Override
+ public boolean canPerformAsync(Task task, TaskActionToolbox toolbox)
+ {
+ return supervisorId != null && startMetadata != null;
+ }
+
+ @Override
+ public ListenableFuture<SegmentPublishResult> performAsync(Task task,
TaskActionToolbox toolbox)
+ {
+ 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 bf121812573..2035ff5a796 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,7 @@ import com.fasterxml.jackson.annotation.JsonProperty;
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.ListenableFuture;
import org.apache.druid.common.config.Configs;
import org.apache.druid.indexing.common.LockGranularity;
import org.apache.druid.indexing.common.TaskLock;
@@ -259,12 +260,19 @@ public class SegmentTransactionalInsertAction implements
TaskAction<SegmentPubli
}
IndexTaskUtils.emitSegmentPublishMetrics(retVal, task, toolbox);
+ return retVal;
+ }
- if (toolbox.shouldFailSegmentPublishImmediately(retVal, task,
supervisorId, startMetadata)) {
- return SegmentPublishResult.fail(retVal.getErrorMsg());
- } else {
- return retVal;
- }
+ @Override
+ public boolean canPerformAsync(Task task, TaskActionToolbox toolbox)
+ {
+ return supervisorId != null && startMetadata != null;
+ }
+
+ @Override
+ public ListenableFuture<SegmentPublishResult> performAsync(Task task,
TaskActionToolbox toolbox)
+ {
+ 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/TaskAction.java
b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskAction.java
index 1deab619f75..1f9b62c991d 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 @@ package org.apache.druid.indexing.common.actions;
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 @@ public interface TaskAction<RetType>
return false;
}
- default Future<RetType> performAsync(Task task, TaskActionToolbox toolbox)
+ default ListenableFuture<RetType> 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 fb0e8c51400..457f14d013f 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,25 @@
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> RetType submit(TaskAction<RetType> taskAction) throws IOException;
+
+ default <RetType> ListenableFuture<RetType> submitAsync(TaskAction<RetType>
taskAction)
+ {
+ try {
+ 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/TaskActionToolbox.java
b/indexing-service/src/main/java/org/apache/druid/indexing/common/actions/TaskActionToolbox.java
index 6294710f167..8f19a3b8c77 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,6 +21,8 @@ package org.apache.druid.indexing.common.actions;
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;
@@ -32,8 +34,13 @@ 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;
@@ -43,6 +50,7 @@ public class TaskActionToolbox
private final ServiceEmitter emitter;
private final SupervisorManager supervisorManager;
private final ObjectMapper jsonMapper;
+ private final ScheduledExecutorService actionExec;
private Optional<TaskRunnerFactory> factory = Optional.absent();
@Inject
@@ -53,7 +61,8 @@ public class TaskActionToolbox
SegmentAllocationQueue segmentAllocationQueue,
ServiceEmitter emitter,
SupervisorManager supervisorManager,
- @Json ObjectMapper jsonMapper
+ @Json ObjectMapper jsonMapper,
+ ScheduledExecutorFactory scheduledExecutorFactory
)
{
this.taskLockbox = taskLockbox;
@@ -63,6 +72,10 @@ public class 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");
}
public TaskActionToolbox(
@@ -81,7 +94,8 @@ public class TaskActionToolbox
null,
emitter,
supervisorManager,
- jsonMapper
+ jsonMapper,
+ ScheduledExecutors::fixed
);
}
@@ -140,25 +154,46 @@ public class TaskActionToolbox
}
/**
- * 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.
+ * 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 boolean shouldFailSegmentPublishImmediately(
- SegmentPublishResult result,
+ public ListenableFuture<SegmentPublishResult> publishSegmentsWhenReady(
Task task,
String supervisorId,
- DataSourceMetadata startMetadata
+ DataSourceMetadata startMetadata,
+ BiFunction<Task, TaskActionToolbox, SegmentPublishResult> publishAction
)
{
- if (result.isSuccess() || !result.isRetryable() || startMetadata == null) {
- return false;
+ // Try publishing later if the failure was due to offset mismatch
+ final ListenableFuture<Boolean> taskReadyToPublishFuture;
+ if (startMetadata == null || supervisorId == null) {
+ taskReadyToPublishFuture = Futures.immediateFuture(true);
+ } else {
+ taskReadyToPublishFuture =
supervisorManager.isTaskReadyToPublishSegments(
+ supervisorId,
+ task.getId(),
+ startMetadata
+ );
}
- return !getSupervisorManager().isAnotherTaskGroupPublishingToPartitions(
- supervisorId,
- task.getId(),
- startMetadata
+ return Futures.transform(
+ taskReadyToPublishFuture,
+ readyToPublish -> {
+ if (Boolean.TRUE.equals(readyToPublish)) {
+ 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("Non-retryable metadata
mismatch error: " + result.getErrorMsg());
+ } else {
+ return result;
+ }
+ } else {
+ return SegmentPublishResult.retryableFailure("Task is not ready to
publish yet");
+ }
+ },
+ actionExec
);
}
}
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 a12a9ea22fe..6c2efddd746 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,8 +27,13 @@ import com.google.common.collect.ImmutableList;
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 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;
@@ -79,7 +85,9 @@ import org.apache.druid.utils.CollectionUtils;
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;
@@ -101,6 +109,7 @@ import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
/**
*
@@ -121,6 +130,7 @@ public class OverlordResource
private final WorkerTaskRunnerQueryAdapter workerTaskRunnerQueryAdapter;
private final AuthConfig authConfig;
+ private final ObjectMapper jsonMapper;
private static final List<String> API_TASK_STATES =
ImmutableList.of("pending", "waiting", "running", "complete");
private static final Set<String> AUDITED_TASK_TYPES
@@ -137,7 +147,8 @@ public class OverlordResource
AuditManager auditManager,
AuthorizerMapper authorizerMapper,
WorkerTaskRunnerQueryAdapter workerTaskRunnerQueryAdapter,
- AuthConfig authConfig
+ AuthConfig authConfig,
+ ObjectMapper jsonMapper
)
{
this.overlord = overlord;
@@ -150,6 +161,7 @@ public class OverlordResource
this.authorizerMapper = authorizerMapper;
this.workerTaskRunnerQueryAdapter = workerTaskRunnerQueryAdapter;
this.authConfig = authConfig;
+ this.jsonMapper = jsonMapper;
}
/**
@@ -506,34 +518,58 @@ public class OverlordResource
@Path("/action")
@Produces(MediaType.APPLICATION_JSON)
@ResourceFilters(StateResourceFilter.class)
- public Response doAction(final TaskActionHolder holder)
+ public void performTaskAction(
+ final TaskActionHolder holder,
+ @Context HttpServletRequest request
+ )
{
- return asLeaderWith(
- taskMaster.getTaskActionClient(holder.getTask()),
- new Function<>()
+ final AsyncRequest asyncRequest = new AsyncRequest(request);
+ final Optional<TaskActionClient> taskActionClient =
taskMaster.getTaskActionClient(holder.getTask());
+ if (!taskActionClient.isPresent()) {
+ // Encourage client to try again soon, when we'll likely have a redirect
set up
+ asyncRequest.complete(Status.SERVICE_UNAVAILABLE.getStatusCode(), null);
+ return;
+ }
+
+ final ListenableFuture<?> future =
taskActionClient.get().submitAsync(holder.getAction());
+ asyncRequest.context.addListener(
+ ServletResourceUtils.createAsyncTimeoutListener(_ -> {
+ asyncRequest.complete(HttpResponseStatus.GATEWAY_TIMEOUT.code(),
null);
+ future.cancel(true);
+ })
+ );
+
+ // Use a default timeout of 15 minutes
+ asyncRequest.context.setTimeout(15 * 60_000);
+
+ Futures.addCallback(
+ future,
+ new FutureCallback<Object>()
{
@Override
- public Response apply(TaskActionClient taskActionClient)
+ public void onSuccess(Object result)
{
- final Map<String, Object> retMap;
+ // Use null-safe map since some task actions may return null result
+ final Map<String, Object> payload = new HashMap<>();
+ payload.put("result", result);
- // 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.
+ asyncRequest.complete(Status.OK.getStatusCode(), payload);
+ }
- 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)
+ {
+ log.warn(t, "Failed to perform task action");
+ if (t instanceof DruidException druidException) {
+ asyncRequest.complete(druidException.getStatusCode(),
druidException.toErrorResponse());
+ } else {
+ final Map<String, String> payload = new HashMap<>();
+ payload.put("error", t.getMessage());
+
asyncRequest.complete(Status.INTERNAL_SERVER_ERROR.getStatusCode(), payload);
}
-
- return Response.ok().entity(retMap).build();
}
- }
+ },
+ MoreExecutors.directExecutor()
);
}
@@ -884,4 +920,36 @@ public class OverlordResource
)
);
}
+
+ private class AsyncRequest
+ {
+ final AtomicBoolean completed = new AtomicBoolean(false);
+ final AsyncContext context;
+
+ AsyncRequest(HttpServletRequest request)
+ {
+ this.context = request.startAsync();
+ }
+
+ 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/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 3510ae7cfb6..ffec91c56d7 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.fasterxml.jackson.databind.ObjectMapper;
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 class SupervisorManager implements
SupervisorStatsProvider
}
/**
- * 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<Boolean> isTaskReadyToPublishSegments(
String supervisorId,
String taskId,
DataSourceMetadata startMetadata
@@ -616,7 +616,7 @@ public class SupervisorManager implements
SupervisorStatsProvider
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,8 @@ public class SupervisorManager implements
SupervisorStatsProvider
);
}
- try {
- final Set<Object> 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;
- }
+ return ((SeekableStreamSupervisor<?, ?, ?>) supervisor.lhs)
+ .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 727a9e3a1bd..35eed987a51 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.ListenableFuture;
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.TreeSet;
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 abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
boolean handoffEarly = false; // set by
SupervisorManager.stopTaskGroupEarly
+ final AtomicBoolean completed = new AtomicBoolean(false);
+ final List<Runnable> completionListeners = new ArrayList<>();
+
TaskGroup(
int groupId,
ImmutableMap<PartitionIdType, SequenceOffsetType> startingSequences,
@@ -358,6 +363,33 @@ public abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
return baseSequenceName;
}
+ /**
+ * Called when this TaskGroup is removed from {@link
#pendingCompletionTaskGroups}.
+ */
+ public 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 abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
}
}
+ /**
+ * 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<Boolean> readyToPublish = SettableFuture.create();
+ final ConcurrentHashMap<TaskGroup, Set<PartitionIdType>>
blockingTaskGroups;
+
+ TaskWaitingToPublish(String taskId, Map<TaskGroup, Set<PartitionIdType>>
blockingTaskGroups)
+ {
+ this.latch = new CountDownLatch(blockingTaskGroups.size());
+ this.blockingTaskGroups = new ConcurrentHashMap<>(blockingTaskGroups);
+
+ tasksWaitingToPublish.put(taskId, this);
+
+ for (TaskGroup group : blockingTaskGroups.keySet()) {
+ group.addCompletionListener(() -> {
+ latch.countDown();
+ this.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 abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
// Map<{group id}, List<{pending completion task groups}>>
private final ConcurrentHashMap<Integer, CopyOnWriteArrayList<TaskGroup>>
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<String, TaskWaitingToPublish>
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
@@ -2871,7 +2941,7 @@ public abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
// 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<Integer, Map<PartitionIdType, SequenceOffsetType>>
latestCheckpoints = new TreeMap<>(
@@ -2934,47 +3004,104 @@ public abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
}
/**
- * 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
- * indicates that the current task would need to wait for the other tasks to
+ * 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 older 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<Object> partitions)
+ private Map<TaskGroup, Set<PartitionIdType>>
getOlderTaskGroupsPublishingToPartitions(
+ String taskId,
+ Map<PartitionIdType, SequenceOffsetType> taskStartOffsets
+ )
{
// Identify all the partitions that are being published by other taskGroups
- final Map<PartitionIdType, Set<TaskGroup>> partitionIdToPublishingGroups =
new HashMap<>();
+ final Map<PartitionIdType, Map<TaskGroup, SequenceOffsetType>>
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()) {
- return false;
+ return Map.of();
}
// Check if any of the partitions of this taskId are being published by
other taskGroups
- for (Object partition : partitions) {
- @SuppressWarnings("unchecked")
- final PartitionIdType partitionId = (PartitionIdType) partition;
+ final Map<TaskGroup, Set<PartitionIdType>> blockingTaskGroups = new
HashMap<>();
+ taskStartOffsets.forEach((partitionId, taskOffset) -> {
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;
+ 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 false;
+ return blockingTaskGroups;
}
+ private boolean isOlderTaskGroupPublishingToPartitions(
+ String taskId,
+ Map<PartitionIdType, SequenceOffsetType> taskStartOffsets
+ )
+ {
+ return !getOlderTaskGroupsPublishingToPartitions(taskId,
taskStartOffsets).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<Boolean> 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<PartitionIdType, SequenceOffsetType> taskStartOffsets =
+ (Map<PartitionIdType, SequenceOffsetType>) streamStartMetadata
+ .getSeekableStreamSequenceNumbers()
+ .getPartitionSequenceNumberMap();
+
+ final Map<TaskGroup, Set<PartitionIdType>> blockingTaskGroups =
+ getOlderTaskGroupsPublishingToPartitions(taskId, taskStartOffsets);
+
+ // 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(
@@ -4211,7 +4338,7 @@ public abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
// 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
@@ -4220,7 +4347,7 @@ public abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
// 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);
@@ -4228,6 +4355,7 @@ public abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
}
taskGroupList.removeAll(toRemove);
+ toRemove.forEach(TaskGroup::onCompleted);
}
// Ignore return value; just await.
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 4e9ae27f887..eb09e017acb 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 class SegmentTransactionalInsertActionTest
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."
),
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 6b9c88727ae..73900ea2c2e 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 class TaskActionTestKit implements
BeforeEachCallback, AfterEachCallback
segmentAllocationQueue,
emitter,
supervisorManager,
- objectMapper
+ objectMapper,
+ ScheduledExecutors::fixed
);
taskActionDelegate = new HashMap<>();
testDerbyConnector.createDataSourceTable();
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 8df71215f36..e75ada0721f 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 class TaskActionToolboxTest
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<Boolean> taskReadyFuture = SettableFuture.create();
+
+ when(supervisorManager.isTaskReadyToPublishSegments(supervisorId,
task.getId(), startMetadata))
+ .thenReturn(taskReadyFuture);
+
+ final Set<DataSegment> segmentsToPublish =
Set.of(DataSegment.builder(SegmentId.dummy("wiki")).build());
+ final ListenableFuture<SegmentPublishResult> 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<DataSegment> segmentsToPublish =
Set.of(DataSegment.builder(SegmentId.dummy("wiki")).build());
+ final ListenableFuture<SegmentPublishResult> 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<Boolean> taskReadyFuture = SettableFuture.create();
+
+ when(supervisorManager.isTaskReadyToPublishSegments(supervisorId,
task.getId(), startMetadata))
+ .thenReturn(taskReadyFuture);
+
+ final ListenableFuture<SegmentPublishResult> 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<Boolean> taskReadyFuture = SettableFuture.create();
+
+ when(supervisorManager.isTaskReadyToPublishSegments(supervisorId,
task.getId(), startMetadata))
+ .thenReturn(taskReadyFuture);
+
+ final Set<DataSegment> segmentsToPublish =
Set.of(DataSegment.builder(SegmentId.dummy("wiki")).build());
+ final ListenableFuture<SegmentPublishResult> 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<SegmentPublishResult> publishFuture)
+ {
+ return FutureUtils.getUnchecked(publishFuture, true);
+ }
}
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 b65465a6e9f..ae1f50cc044 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.base.Optional;
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.indexer.TaskStatusPlus;
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.TaskStorage;
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.ActiveTaskLookup;
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.Assertions;
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;
@@ -184,7 +194,8 @@ public class OverlordResourceTest
auditManager,
authMapper,
workerTaskRunnerQueryAdapter,
- authConfig
+ authConfig,
+ TestHelper.JSON_MAPPER
);
}
@@ -1463,7 +1474,7 @@ public class OverlordResourceTest
{
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(
@@ -1476,6 +1487,158 @@ public class OverlordResourceTest
);
}
+ @Test
+ public void test_performTaskAction_returnsOk_whenFutureCompletes()
+ {
+ final NoopTask task = NoopTask.create();
+ final TaskActionHolder holder = new TaskActionHolder(task, new
LockListAction());
+ final List<TaskLock> 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<Object> future = SettableFuture.create();
+
+ EasyMock.expect(taskMaster.getTaskActionClient(task))
+ .andReturn(Optional.of(taskActionClient))
+ .once();
+ EasyMock.expect(taskActionClient.submitAsync(holder.getAction()))
+ .andReturn(future)
+ .once();
+
+ final Capture<AsyncListener> 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() throws
Exception
+ {
+ 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());
+
+ final Map<String, String> responsePayload =
TestHelper.JSON_MAPPER.readValue(
+ asyncResponse.baos.toByteArray(),
+ new TypeReference<>() {}
+ );
+ Assertions.assertEquals(
+ Map.of("error", "Something went wrong"),
+ responsePayload
+ );
+
+ EasyMock.verify(taskActionClient);
+ }
+
private void expectAuthorizationTokenCheck()
{
expectAuthorizationTokenCheck(Users.DRUID);
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 1bca81b8759..822362cf59e 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.common.concurrent.Execs;
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 class OverlordTest
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 a93686972a1..bec5e185074 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.base.Optional;
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;
@@ -1146,11 +1148,11 @@ public class SupervisorManagerTest extends
EasyMockSupport
}
@Test
- public void
test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifSupervisorIdIsNull()
+ public void
test_isTaskReadyToPublishSegments_throwsException_ifSupervisorIdIsNull()
{
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 class SupervisorManagerTest extends
EasyMockSupport
}
@Test
- public void
test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifSupervisorNotFound()
+ public void
test_isTaskReadyToPublishSegments_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 class SupervisorManagerTest extends EasyMockSupport
}
@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,13 +1188,13 @@ public class SupervisorManagerTest extends
EasyMockSupport
manager.start();
- Assertions.assertFalse(
- manager.isAnotherTaskGroupPublishingToPartitions(supervisorId,
"task1", null)
+ Assertions.assertTrue(
+ manager.isTaskReadyToPublishSegments(supervisorId, "task1", null).get()
);
}
@Test
- public void
test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifMetadataIsNull()
+ public void
test_isTaskReadyToPublishSegments_throwsException_ifMetadataIsNull()
{
final String supervisorId = "supervisor1";
final SeekableStreamSupervisor<Integer, String, ByteEntity>
seekableStreamSupervisor =
@@ -1209,7 +1212,7 @@ public class SupervisorManagerTest extends EasyMockSupport
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());
@@ -1221,7 +1224,7 @@ public class SupervisorManagerTest extends EasyMockSupport
}
@Test
- public void
test_isAnotherTaskGroupPublishingToPartitions_throwsException_ifMetadataIsInvalid()
+ public void
test_isTaskReadyToPublishSegments_throwsException_ifMetadataIsInvalid()
{
final String supervisorId = "supervisor1";
final SeekableStreamSupervisor<Integer, String, ByteEntity>
seekableStreamSupervisor =
@@ -1239,7 +1242,7 @@ public class SupervisorManagerTest extends EasyMockSupport
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 class SupervisorManagerTest extends EasyMockSupport
}
@Test
- public void test_isAnotherTaskGroupPublishingToPartitions()
+ public void test_isTaskReadyToPublishSegments() throws Exception
{
final String supervisorId = "supervisor1";
final SeekableStreamSupervisor<Integer, String, ByteEntity>
seekableStreamSupervisor =
@@ -1268,20 +1271,20 @@ public class SupervisorManagerTest extends
EasyMockSupport
// 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 class SupervisorManagerTest extends
EasyMockSupport
"topic",
Map.of("0", "10")
);
- Assertions.assertTrue(
- manager.isAnotherTaskGroupPublishingToPartitions(supervisorId,
conflictingTaskId, startMetadata)
- );
- Assertions.assertFalse(
- manager.isAnotherTaskGroupPublishingToPartitions(supervisorId,
readyTaskId, startMetadata)
- );
+
+ final ListenableFuture<Boolean> conflictingTaskFuture =
+ manager.isTaskReadyToPublishSegments(supervisorId, conflictingTaskId,
startMetadata);
+ Assertions.assertFalse(conflictingTaskFuture.isDone());
+
+ final ListenableFuture<Boolean> readyTaskFuture =
+ manager.isTaskReadyToPublishSegments(supervisorId, readyTaskId,
startMetadata);
+ Assertions.assertTrue(readyTaskFuture.isDone());
+ Assertions.assertTrue(readyTaskFuture.get());
}
@Test
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 f0d08b0d941..55ee58a7381 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.ImmutableList;
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 abstract class SeekableStreamIndexTaskTestBase
extends EasyMockSupport
@Nullable DataSourceMetadata previousDataSourceMetadata
)
{
- // log.info("Adding checkpoint hash to the set");
checkpointRequestsHash.add(
Objects.hash(
supervisorId,
@@ -629,6 +629,16 @@ public abstract class SeekableStreamIndexTaskTestBase
extends EasyMockSupport
);
return true;
}
+
+ @Override
+ public ListenableFuture<Boolean> isTaskReadyToPublishSegments(
+ String supervisorId,
+ String taskId,
+ DataSourceMetadata startMetadata
+ )
+ {
+ return Futures.immediateFuture(true);
+ }
},
objectMapper
);
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 d7634a842db..912455f8ba2 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.ImmutableMap;
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;
@@ -43,6 +44,7 @@ import org.apache.druid.indexing.common.TestUtils;
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;
@@ -510,6 +512,85 @@ public class SeekableStreamSupervisorStateTest extends
EasyMockSupport
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()
+ {
+ 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<Boolean> future =
supervisor.isTaskReadyToPublishSegments(
+ "task1",
+ new TestSeekableStreamDataSourceMetadata(
+ new SeekableStreamStartSequenceNumbers<>(STREAM, Map.of("p1",
"100"), Set.of())
+ )
+ );
+
+ Assertions.assertTrue(future.isDone());
+ Assertions.assertEquals(Boolean.TRUE, FutureUtils.getUnchecked(future,
true));
+ }
+
+ @Test
+ public void test_isTaskReadyToPublishSegments_waitsForEarlierTask()
+ {
+ 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();
+ supervisor.addDiscoveredTaskToPendingCompletionTaskGroups(0,
"earlierGroupPendingCompletion", Map.of("p1", "0"));
+ supervisor.addDiscoveredTaskToPendingCompletionTaskGroups(1,
"laterGroupPendingCompletion", Map.of("p1", "10000"));
+
+ final ListenableFuture<Boolean> future =
supervisor.isTaskReadyToPublishSegments(
+ "waitingTask",
+ 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 earlierTaskGroup =
+ supervisor.getPendingCompletionTaskGroups(0).get(0);
+ 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
public void testConnectingToStreamFail()
{
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 04b745c812e..c34553a617d 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<DataSegment> segments;
private final boolean success;
private final boolean retryable;
+ private final boolean offsetMismatch;
private final String errorMsg;
private final List<PendingSegmentRecord> upgradedPendingSegments;
@@ -51,17 +52,22 @@ public class SegmentPublishResult
public static SegmentPublishResult ok(Set<DataSegment> segments,
List<PendingSegmentRecord> 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);
}
@JsonCreator
@@ -72,7 +78,7 @@ public class SegmentPublishResult
@JsonProperty("errorMsg") @Nullable String errorMsg
)
{
- this(segments, success, retryable, errorMsg, null);
+ this(segments, success, retryable, errorMsg, null, false);
}
private SegmentPublishResult(
@@ -80,7 +86,8 @@ public class SegmentPublishResult
boolean success,
boolean retryable,
@Nullable String errorMsg,
- List<PendingSegmentRecord> upgradedPendingSegments
+ List<PendingSegmentRecord> upgradedPendingSegments,
+ boolean offsetMismatch
)
{
this.segments = Preconditions.checkNotNull(segments, "segments");
@@ -88,6 +95,7 @@ public class 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 +137,15 @@ public class SegmentPublishResult
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<PendingSegmentRecord> 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 f469b0e7ebb..ee8933cf32c 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 @@ public class IndexerSQLMetadataStorageCoordinator
implements IndexerMetadataStor
// 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
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
