Skip to content

Fix task executor worker identity reconciliation races #874

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
Andyz26 wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Fix task executor worker identity reconciliation races #874

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ef38242
Fix task executor worker identity races
Andyz26 Sep 1, 2026
eadc133
Handle task cancellation cleanup failures
Andyz26 Sep 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions ...ore/src/main/java/io/mantisrx/server/master/resourcecluster/TaskExecutorRegistration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,9 @@
@ToString
@EqualsAndHashCode
public class TaskExecutorRegistration {
public static final String ACCEPTED_TASK_RESERVATION_ATTRIBUTE =
"mantis_task_executor_accepted_task_reservation";

@NonNull
TaskExecutorID taskExecutorID;

Expand Down Expand Up @@ -149,6 +152,13 @@ public Optional<String> getAttributeByKey(String attributeKey) {
return Optional.empty();
}

@JsonIgnore
public boolean reservesAcceptedTask() {
return getAttributeByKey(ACCEPTED_TASK_RESERVATION_ATTRIBUTE)
.map(Boolean::parseBoolean)
.orElse(false);
}

@JsonIgnore
public Map<String, String> getSchedulingAttributes() {
return taskExecutorAttributes.entrySet().stream()
Expand Down
6 changes: 5 additions & 1 deletion ...antis-control-plane-core/src/main/java/io/mantisrx/server/worker/TaskExecutorGateway.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,9 +75,13 @@ public TaskAlreadyRunningException(WorkerId workerId) {
}

public TaskAlreadyRunningException(WorkerId workerId, Throwable cause) {
super(cause);
super(String.format("Task executor is already running %s", workerId), cause);
this.currentlyRunningWorkerTask = workerId;
}

public WorkerId getCurrentlyRunningWorkerTask() {
return currentlyRunningWorkerTask;
}
}

class TaskNotFoundException extends Exception {
Expand Down
85 changes: 76 additions & 9 deletions ...plane-server/src/main/java/io/mantisrx/master/resourcecluster/AssignmentHandlerActor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,15 @@
import io.mantisrx.server.master.ExecuteStageRequestFactory;
import io.mantisrx.server.master.scheduler.JobMessageRouter;
import io.mantisrx.server.worker.TaskExecutorGateway;
import io.mantisrx.server.worker.TaskExecutorGateway.TaskAlreadyRunningException;
import io.mantisrx.shaded.com.fasterxml.jackson.annotation.JsonCreator;
import io.mantisrx.shaded.com.fasterxml.jackson.annotation.JsonProperty;
import org.apache.flink.util.ExceptionUtils;
import java.time.Duration;
import java.time.Instant;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicBoolean;
import javax.annotation.Nullable;
import lombok.Value;
import lombok.extern.slf4j.Slf4j;
Expand Down Expand Up @@ -96,6 +98,8 @@ private void onTaskExecutorAssignmentRequest(TaskExecutorAssignmentRequest reque
request, request.getAttempt(), maxAssignmentRetries);
try {
TaskExecutorRegistration registration = request.getRegistration();
// Gateway completion and timeout race for the right to start submission.
AtomicBoolean submitGateClaimed = new AtomicBoolean();
// Use the gateway future from the request
CompletableFuture<TaskExecutorGateway> gatewayFut = request.getGatewayFuture();

Expand All @@ -104,6 +108,11 @@ private void onTaskExecutorAssignmentRequest(TaskExecutorAssignmentRequest reque
.<Object>thenComposeAsync(gateway -> {
log.debug("Successfully obtained gateway for task executor {}",
registration.getTaskExecutorID());
if (!submitGateClaimed.compareAndSet(false, true)) {
return CompletableFuture.failedFuture(
new java.util.concurrent.TimeoutException(
"Assignment attempt expired before submit started"));
}
return gateway
.submitTask(
executeStageRequestFactory.of(
Expand All @@ -119,38 +128,64 @@ private void onTaskExecutorAssignmentRequest(TaskExecutorAssignmentRequest reque
log.error("[Submit Task] failed for {}: {}",
registration.getTaskExecutorID(), throwable.getMessage());
return new TaskExecutorAssignmentFailedEvent(
request, ExceptionUtils.stripCompletionException(throwable));
request,
unwrapAssignmentFailure(throwable),
AssignmentFailureType.MayHaveRun);
});
})
.exceptionally(throwable -> {
log.warn("Failed to obtain gateway for task executor {}",
registration.getTaskExecutorID(), throwable);
return new TaskExecutorAssignmentFailedEvent(
request,
ExceptionUtils.stripCompletionException(throwable));
unwrapAssignmentFailure(throwable),
submitGateClaimed.get()
? AssignmentFailureType.MayHaveRun
: AssignmentFailureType.NotSent);
})
.toCompletableFuture()
.orTimeout(
assignmentTimeout.toMillis(),
java.util.concurrent.TimeUnit.MILLISECONDS)
.exceptionally(throwable -> {
if (throwable instanceof java.util.concurrent.TimeoutException) {
Throwable cause = unwrapAssignmentFailure(throwable);
if (cause instanceof java.util.concurrent.TimeoutException) {
boolean preventedSubmit = submitGateClaimed.compareAndSet(false, true);
log.warn("Assignment timeout for task executor {} after {}ms",
registration.getTaskExecutorID(), assignmentTimeout.toMillis());
return new TaskExecutorAssignmentFailedEvent(
request,
throwable);
cause,
preventedSubmit
? AssignmentFailureType.NotSent
: AssignmentFailureType.MayHaveRun);
}
return new TaskExecutorAssignmentFailedEvent(
request,
ExceptionUtils.stripCompletionException(throwable));
cause,
submitGateClaimed.get()
? AssignmentFailureType.MayHaveRun
: AssignmentFailureType.NotSent);
});

akka.pattern.Patterns.pipe(ackFuture, getContext().getDispatcher()).to(self());
} catch (Exception e) {
log.error("Exception during task executor assignment for {}",
request.getRegistration().getTaskExecutorID(), e);
self().tell(new TaskExecutorAssignmentFailedEvent(request, e), self());
self().tell(new TaskExecutorAssignmentFailedEvent(
request, unwrapAssignmentFailure(e), AssignmentFailureType.NotSent), self());
}
}

private static Throwable unwrapAssignmentFailure(Throwable throwable) {
Throwable current = throwable;
while (true) {
Throwable unwrapped = ExceptionUtils.stripCompletionException(
ExceptionUtils.stripExecutionException(current));
if (unwrapped == current) {
return current;
}
current = unwrapped;
}
}

Expand All @@ -170,16 +205,24 @@ private void onAssignmentFailed(TaskExecutorAssignmentFailedEvent event) {
maxAssignmentRetries,
event.getThrowable().getMessage());

if (request.getAttempt() >= maxAssignmentRetries) {
Throwable failure = event.getThrowable();
AssignmentFailureType failureType = failure instanceof TaskAlreadyRunningException
? AssignmentFailureType.Conflict
: event.getFailureType();

if (failureType != AssignmentFailureType.NotSent
|| request.getAttempt() >= maxAssignmentRetries) {
log.error("Assignment failed for {} after {} attempts, giving up",
registration.getTaskExecutorID(), maxAssignmentRetries);

// Send assignmentFailure event to parent after max retries
getContext().parent().tell(new TaskExecutorAssignmentFailAndTerminate(
registration.getTaskExecutorID(),
request.getAllocationRequest(),
event.getThrowable(),
request.getAttempt()
failure,
request.getAttempt(),
request.getAssignmentEpoch(),
failureType
), self());
} else {
log.info("Retrying assignment for {} in {} (attempt {}/{})",
Expand Down Expand Up @@ -256,6 +299,7 @@ public static class TaskExecutorAssignmentRequest {
TaskExecutorRegistration registration;
CompletableFuture<TaskExecutorGateway> gatewayFuture;
int attempt;
long assignmentEpoch;

/*
Deprecated field.
Expand All @@ -272,6 +316,7 @@ public TaskExecutorAssignmentRequest(
@JsonProperty("registration") TaskExecutorRegistration registration,
@JsonProperty("gatewayFuture") CompletableFuture<TaskExecutorGateway> gatewayFuture,
@JsonProperty("attempt") int attempt,
@JsonProperty("assignmentEpoch") long assignmentEpoch,
@JsonProperty("previousFailure") @Nullable Throwable previousFailure,
@JsonProperty("requestTime") Instant requestTime
) {
Expand All @@ -280,6 +325,7 @@ public TaskExecutorAssignmentRequest(
this.registration = registration;
this.gatewayFuture = gatewayFuture;
this.attempt = attempt;
this.assignmentEpoch = assignmentEpoch;
this.previousFailure = previousFailure;
this.requestTime = requestTime;
}
Expand All @@ -289,13 +335,24 @@ public static TaskExecutorAssignmentRequest of(
TaskExecutorID taskExecutorID,
TaskExecutorRegistration registration,
CompletableFuture<TaskExecutorGateway> gatewayFuture
) {
return of(allocationRequest, taskExecutorID, registration, gatewayFuture, 0L);
}

public static TaskExecutorAssignmentRequest of(
TaskExecutorAllocationRequest allocationRequest,
TaskExecutorID taskExecutorID,
TaskExecutorRegistration registration,
CompletableFuture<TaskExecutorGateway> gatewayFuture,
long assignmentEpoch
) {
return new TaskExecutorAssignmentRequest(
allocationRequest,
taskExecutorID,
registration,
gatewayFuture,
1,
assignmentEpoch,
null,
Instant.now()
);
Expand All @@ -316,6 +373,7 @@ public TaskExecutorAssignmentRequest onRetry(CompletableFuture<TaskExecutorGatew
registration,
freshGatewayFuture, // Use fresh future instead of reusing the old one
attempt + 1,
assignmentEpoch,
null,
requestTime
);
Expand Down Expand Up @@ -343,6 +401,13 @@ private static class TaskExecutorAssignmentSucceededEvent {
private static class TaskExecutorAssignmentFailedEvent {
TaskExecutorAssignmentRequest request;
Throwable throwable;
AssignmentFailureType failureType;
}

enum AssignmentFailureType {
NotSent,
MayHaveRun,
Conflict,
}

@Value
Expand All @@ -351,6 +416,8 @@ public static class TaskExecutorAssignmentFailAndTerminate {
TaskExecutorAllocationRequest allocationRequest;
Throwable throwable;
int attemptCount;
long assignmentEpoch;
AssignmentFailureType failureType;
}

@Value
Expand Down
Loading
Loading