diff --git a/mantis-control-plane/mantis-control-plane-core/src/main/java/io/mantisrx/server/master/resourcecluster/TaskExecutorRegistration.java b/mantis-control-plane/mantis-control-plane-core/src/main/java/io/mantisrx/server/master/resourcecluster/TaskExecutorRegistration.java index a1457db24..4e11c0a80 100644 --- a/mantis-control-plane/mantis-control-plane-core/src/main/java/io/mantisrx/server/master/resourcecluster/TaskExecutorRegistration.java +++ b/mantis-control-plane/mantis-control-plane-core/src/main/java/io/mantisrx/server/master/resourcecluster/TaskExecutorRegistration.java @@ -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; @@ -149,6 +152,13 @@ public Optional getAttributeByKey(String attributeKey) { return Optional.empty(); } + @JsonIgnore + public boolean reservesAcceptedTask() { + return getAttributeByKey(ACCEPTED_TASK_RESERVATION_ATTRIBUTE) + .map(Boolean::parseBoolean) + .orElse(false); + } + @JsonIgnore public Map getSchedulingAttributes() { return taskExecutorAttributes.entrySet().stream() diff --git a/mantis-control-plane/mantis-control-plane-core/src/main/java/io/mantisrx/server/worker/TaskExecutorGateway.java b/mantis-control-plane/mantis-control-plane-core/src/main/java/io/mantisrx/server/worker/TaskExecutorGateway.java index fdf9f3611..32f4c8052 100644 --- a/mantis-control-plane/mantis-control-plane-core/src/main/java/io/mantisrx/server/worker/TaskExecutorGateway.java +++ b/mantis-control-plane/mantis-control-plane-core/src/main/java/io/mantisrx/server/worker/TaskExecutorGateway.java @@ -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 { diff --git a/mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/resourcecluster/AssignmentHandlerActor.java b/mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/resourcecluster/AssignmentHandlerActor.java index 855271cbf..80d823d7f 100644 --- a/mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/resourcecluster/AssignmentHandlerActor.java +++ b/mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/resourcecluster/AssignmentHandlerActor.java @@ -11,6 +11,7 @@ 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; @@ -18,6 +19,7 @@ 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; @@ -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 gatewayFut = request.getGatewayFuture(); @@ -104,6 +108,11 @@ private void onTaskExecutorAssignmentRequest(TaskExecutorAssignmentRequest reque .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( @@ -119,7 +128,9 @@ 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 -> { @@ -127,30 +138,54 @@ private void onTaskExecutorAssignmentRequest(TaskExecutorAssignmentRequest reque 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; } } @@ -170,7 +205,13 @@ 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); @@ -178,8 +219,10 @@ private void onAssignmentFailed(TaskExecutorAssignmentFailedEvent event) { 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 {}/{})", @@ -256,6 +299,7 @@ public static class TaskExecutorAssignmentRequest { TaskExecutorRegistration registration; CompletableFuture gatewayFuture; int attempt; + long assignmentEpoch; /* Deprecated field. @@ -272,6 +316,7 @@ public TaskExecutorAssignmentRequest( @JsonProperty("registration") TaskExecutorRegistration registration, @JsonProperty("gatewayFuture") CompletableFuture gatewayFuture, @JsonProperty("attempt") int attempt, + @JsonProperty("assignmentEpoch") long assignmentEpoch, @JsonProperty("previousFailure") @Nullable Throwable previousFailure, @JsonProperty("requestTime") Instant requestTime ) { @@ -280,6 +325,7 @@ public TaskExecutorAssignmentRequest( this.registration = registration; this.gatewayFuture = gatewayFuture; this.attempt = attempt; + this.assignmentEpoch = assignmentEpoch; this.previousFailure = previousFailure; this.requestTime = requestTime; } @@ -289,6 +335,16 @@ public static TaskExecutorAssignmentRequest of( TaskExecutorID taskExecutorID, TaskExecutorRegistration registration, CompletableFuture gatewayFuture + ) { + return of(allocationRequest, taskExecutorID, registration, gatewayFuture, 0L); + } + + public static TaskExecutorAssignmentRequest of( + TaskExecutorAllocationRequest allocationRequest, + TaskExecutorID taskExecutorID, + TaskExecutorRegistration registration, + CompletableFuture gatewayFuture, + long assignmentEpoch ) { return new TaskExecutorAssignmentRequest( allocationRequest, @@ -296,6 +352,7 @@ public static TaskExecutorAssignmentRequest of( registration, gatewayFuture, 1, + assignmentEpoch, null, Instant.now() ); @@ -316,6 +373,7 @@ public TaskExecutorAssignmentRequest onRetry(CompletableFuture gatewayFuture = taskExecutorState.getGatewayAsync(); + CompletableFuture gatewayFuture; + try { + gatewayFuture = taskExecutorState.getGatewayAsync(); + } catch (IllegalStateException e) { + self().tell( + new AssignmentHandlerActor.TaskExecutorAssignmentFailAndTerminate( + taskExecutorID, + allocationRequest, + e, + 1, + assignmentEpoch, + AssignmentHandlerActor.AssignmentFailureType.NotSent), + self()); + return; + } // Delegate actual assignment logic to AssignmentHandlerActor AssignmentHandlerActor.TaskExecutorAssignmentRequest assignmentRequest = @@ -632,7 +643,8 @@ private void assignTaskExecutor(TaskExecutorAllocationRequest allocationRequest, allocationRequest, taskExecutorID, registration, - gatewayFuture + gatewayFuture, + assignmentEpoch ); assignmentHandlerActor.tell(assignmentRequest, self()); @@ -650,26 +662,36 @@ private void onTaskExecutorAssignmentFailure(AssignmentHandlerActor.TaskExecutor if (state == null) { log.error("[TaskExecutorAssignmentFailure] TaskExecutor lost during task assignment: {}", request); + return; + } + WorkerId expectedWorker = request.getAllocationRequest().getWorkerId(); + if (!state.isCurrentAssignment(expectedWorker, request.getAssignmentEpoch())) { + log.info("Ignoring stale assignment failure for {} epoch {} on executor {}", + expectedWorker, request.getAssignmentEpoch(), request.getTaskExecutorID()); + return; + } + + jobMessageRouter.routeWorkerEvent(new WorkerLaunchFailed( + expectedWorker, + request.getAllocationRequest().getStageNum(), + "Failed to assign worker to task executor " + request.getTaskExecutorID())); + + if (request.getFailureType() == AssignmentHandlerActor.AssignmentFailureType.NotSent) { + state.onUnassignment(); + disconnectTaskExecutor(request.getTaskExecutorID()); + return; } - else if (state.isRunningTask()) { - log.warn("[onTaskExecutorAssignmentFailure] TaskExecutor {} entered running state already; no need to act", - request.getTaskExecutorID()); + + WorkerId cancellationTarget = expectedWorker; + if (request.getThrowable() instanceof TaskAlreadyRunningException) { + cancellationTarget = ((TaskAlreadyRunningException) request.getThrowable()) + .getCurrentlyRunningWorkerTask(); + state.quarantineWorkerOnTask(cancellationTarget); } else { - log.error("[onTaskExecutorAssignmentFailure] TaskExecutor {} failed to accept assignment: {}", - request.getTaskExecutorID(), request.getAllocationRequest()); - try - { - jobMessageRouter.routeWorkerEvent(new WorkerLaunchFailed( - request.getAllocationRequest().getWorkerId(), - request.getAllocationRequest().getStageNum(), - "Failed to assign worker to task executor " + request.getTaskExecutorID())); - state.onUnassignment(); - // disconnect the TE since it cannot be assigned. - disconnectTaskExecutor(request.getTaskExecutorID()); - } catch (IllegalStateException e) { - log.error("Failed to un-assign taskExecutor {}", request.getTaskExecutorID(), e); - } + state.setCancelledWorkerOnTask(cancellationTarget); } + this.delegate.tryMarkUnavailable(request.getTaskExecutorID()); + cancelTaskOnExecutor(state, request.getTaskExecutorID(), cancellationTarget); } private void onTaskExecutorDisconnection(TaskExecutorDisconnection disconnection) { @@ -684,6 +706,9 @@ private void onTaskExecutorDisconnection(TaskExecutorDisconnection disconnection private void disconnectTaskExecutor(TaskExecutorID taskExecutorID) { final TaskExecutorState state = this.delegate.get(taskExecutorID); + if (state.isRegistered()) { + this.delegate.tryMarkUnavailable(taskExecutorID); + } boolean stateChange = state.onDisconnection(); if (stateChange) { this.delegate.archive(taskExecutorID); @@ -716,6 +741,14 @@ private void setupTaskExecutorStateIfNecessary(TaskExecutorID taskExecutorID) { TaskExecutorState.of(clock, rpcService, jobMessageRouter)); } + private void syncAvailabilityIndex(TaskExecutorID taskExecutorID, TaskExecutorState state) { + if (state.isAvailable()) { + this.delegate.tryMarkAvailable(taskExecutorID); + } else { + this.delegate.tryMarkUnavailable(taskExecutorID); + } + } + private void updateHeartbeatTimeout(TaskExecutorID taskExecutorID) { final TaskExecutorState state = this.delegate.get(taskExecutorID); getTimers().startSingleTimer( @@ -949,16 +982,7 @@ private void onTerminateWorkerRequest(TerminateWorkerRequest request) { state.setCancelledWorkerOnTask(request.getWorkerId()); // Proactively call cancelTask on gateway - state.getGatewayAsync().thenAccept(gateway -> { - gateway.cancelTask(request.getWorkerId()) - .whenComplete((ack, throwable) -> { - if (throwable != null) { - log.warn("Failed to proactively cancel task {} on executor {}", request.getWorkerId(), taskExecutorID, throwable); - } else { - log.info("Successfully proactively cancelled task {} on executor {}", request.getWorkerId(), taskExecutorID); - } - }); - }); + cancelTaskOnExecutor(state, taskExecutorID, request.getWorkerId()); sender().tell(Ack.getInstance(), self()); } else { log.info("Cannot find executor to proactively terminate worker {}", request.getWorkerId()); @@ -966,6 +990,25 @@ private void onTerminateWorkerRequest(TerminateWorkerRequest request) { } } + private void cancelTaskOnExecutor( + TaskExecutorState state, + TaskExecutorID taskExecutorID, + WorkerId workerId) { + try { + state.getGatewayAsync() + .thenCompose(gateway -> gateway.cancelTask(workerId)) + .whenComplete((ack, throwable) -> { + if (throwable != null) { + log.warn("Failed to cancel task {} on executor {}", workerId, taskExecutorID, throwable); + } else { + log.info("Cancellation accepted for task {} on executor {}", workerId, taskExecutorID); + } + }); + } catch (IllegalStateException e) { + log.warn("Cannot cancel task {} on unregistered executor {}", workerId, taskExecutorID, e); + } + } + private void onResourceOverviewRequest(ResourceOverviewRequest request) { sender().tell(getResourceOverview(), self()); } diff --git a/mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/resourcecluster/ExecutorStateManagerImpl.java b/mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/resourcecluster/ExecutorStateManagerImpl.java index 2de99bfa5..f834482ec 100644 --- a/mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/resourcecluster/ExecutorStateManagerImpl.java +++ b/mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/resourcecluster/ExecutorStateManagerImpl.java @@ -45,6 +45,7 @@ import io.mantisrx.server.master.resourcecluster.TaskExecutorRegistration.TaskExecutorGroupKey; import io.mantisrx.shaded.com.google.common.cache.Cache; import io.mantisrx.shaded.com.google.common.cache.CacheBuilder; +import io.mantisrx.shaded.com.google.common.cache.RemovalCause; import io.mantisrx.shaded.com.google.common.cache.RemovalListener; import java.time.Duration; @@ -134,7 +135,9 @@ private TaskExecutorGroupKey findBestFitGroupOrDefault(SchedulingConstraints con .expireAfterWrite(24, TimeUnit.HOURS) .removalListener(notification -> { TaskExecutorState state = (TaskExecutorState) notification.getValue(); - boolean teIsDisabled = state != null && state.onNodeDisabled(); + boolean teIsDisabled = notification.getCause() != RemovalCause.EXPLICIT + && state != null + && state.onNodeDisabled(); log.info("Archived TaskExecutor: {} with disabled state: {} removed due to: {}", notification.getKey(), teIsDisabled, notification.getCause()); }) .build(); @@ -180,13 +183,15 @@ private TaskExecutorGroupKey findBestFitGroupOrDefault(SchedulingConstraints con @Override public void trackIfAbsent(TaskExecutorID taskExecutorID, TaskExecutorState state) { - this.taskExecutorStateMap.putIfAbsent(taskExecutorID, state); - if (this.archivedState.getIfPresent(taskExecutorID) != null) { + TaskExecutorState archived = this.archivedState.getIfPresent(taskExecutorID); + TaskExecutorState stateToTrack = archived == null ? state : archived; + this.taskExecutorStateMap.putIfAbsent(taskExecutorID, stateToTrack); + if (archived != null) { log.info("Reviving archived executor: {}", taskExecutorID); this.archivedState.invalidate(taskExecutorID); } - tryMarkAvailable(taskExecutorID, state); + tryMarkAvailable(taskExecutorID, this.taskExecutorStateMap.get(taskExecutorID)); } /** diff --git a/mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/resourcecluster/TaskExecutorState.java b/mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/resourcecluster/TaskExecutorState.java index 5d59809a0..e685733bc 100644 --- a/mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/resourcecluster/TaskExecutorState.java +++ b/mantis-control-plane/mantis-control-plane-server/src/main/java/io/mantisrx/master/resourcecluster/TaskExecutorState.java @@ -38,6 +38,7 @@ import java.time.Duration; import java.time.Instant; import java.util.Map; +import java.util.Objects; import java.util.concurrent.CompletableFuture; import javax.annotation.Nullable; import lombok.AllArgsConstructor; @@ -54,6 +55,13 @@ enum RegistrationState { Unregistered, } + private enum ReconciliationState { + None, + Cancelling, + Quarantined, + Verifying, + } + private RegistrationState state; @Nullable private TaskExecutorRegistration registration; @@ -73,10 +81,10 @@ enum RegistrationState { private final RpcService rpcService; private final JobMessageRouter jobMessageRouter; - // isTaskCancelled: this state is to mark the current assigned worker has been cancelled and this executor need to - // stop the task and re-register. + private ReconciliationState reconciliationState; @Nullable private WorkerId cancelledWorkerOnTask; + private long assignmentEpoch; // previousWorkerId: tracks the last WorkerId this executor was running before disconnection // This enables targeted notifications when the executor reconnects @@ -94,7 +102,9 @@ static TaskExecutorState of(Clock clock, RpcService rpcService, JobMessageRouter clock, rpcService, jobMessageRouter, + ReconciliationState.None, null, + 0L, null); } @@ -116,12 +126,30 @@ WorkerId getCancelledWorkerId() { } void setCancelledWorkerOnTask(WorkerId cancelledWorkerOnTask) { + if (!isRegistered()) { + throwNotRegistered(String.format("cancellation for %s", cancelledWorkerOnTask)); + } + this.reconciliationState = ReconciliationState.Cancelling; this.cancelledWorkerOnTask = cancelledWorkerOnTask; } + void quarantineWorkerOnTask(WorkerId workerId) { + if (!isRegistered()) { + throwNotRegistered(String.format("quarantine for %s", workerId)); + } + this.reconciliationState = ReconciliationState.Quarantined; + this.cancelledWorkerOnTask = workerId; + setAvailabilityState(AvailabilityState.running(workerId)); + } + boolean onRegistration(TaskExecutorRegistration registration) { if (state == RegistrationState.Registered) { - return false; + if (Objects.equals(this.registration, registration)) { + return false; + } + this.registration = registration; + updateTicker(); + return true; } else { this.state = RegistrationState.Registered; this.registration = registration; @@ -136,9 +164,10 @@ boolean onDisconnection() { } else { state = RegistrationState.Unregistered; registration = null; - // Store the current WorkerId as previousWorkerId for potential reconnection notification previousWorkerId = getWorkerId(); - setAvailabilityState(null); + if (previousWorkerId == null) { + setAvailabilityState(null); + } updateTicker(); return true; } @@ -163,7 +192,11 @@ boolean onAssignment(WorkerId workerId) throws IllegalStateException { throw new IllegalStateException("availability state was null when unassignmentas was issued"); } - return setAvailabilityState(this.availabilityState.onAssignment(workerId)); + boolean changed = setAvailabilityState(this.availabilityState.onAssignment(workerId)); + if (changed) { + assignmentEpoch++; + } + return changed; } boolean onUnassignment() throws IllegalStateException { @@ -174,6 +207,16 @@ boolean onUnassignment() throws IllegalStateException { return setAvailabilityState(this.availabilityState.onUnassignment()); } + long getAssignmentEpoch() { + return assignmentEpoch; + } + + boolean isCurrentAssignment(WorkerId workerId, long expectedEpoch) { + return isAssigned() + && assignmentEpoch == expectedEpoch + && workerId.equals(getWorkerId()); + } + boolean onNodeDisabled() { if (!this.disabled) { this.disabled = true; @@ -201,31 +244,8 @@ boolean onHeartbeat(TaskExecutorHeartbeat heartbeat) throwNotRegistered(String.format("heartbeat %s", heartbeat)); } - TaskExecutorReport report = heartbeat.getTaskExecutorReport(); - if (this.availabilityState instanceof Running && report instanceof Available) { - WorkerId runningWorkerId = this.availabilityState.getWorkerId(); - if (runningWorkerId != null) { - log.warn("Heartbeat indicates available while running {}. Marking worker as lost.", runningWorkerId); - jobMessageRouter.routeWorkerEvent( - new WorkerTerminate(runningWorkerId, WorkerState.Failed, JobCompletedReason.Lost)); - } - } - if (this.cancelledWorkerOnTask != null) { - if (report instanceof Occupied && ((Occupied) report).getWorkerId().equals(this.cancelledWorkerOnTask)) { - log.warn("{} cancelled, request cancel on heartbeat.", this.cancelledWorkerOnTask); - throw new TaskExecutorTaskCancelledException( - String.format( - "heartbeat from %s has cancelled task %s", - heartbeat.getTaskExecutorID(), - this.cancelledWorkerOnTask), - this.cancelledWorkerOnTask); - } else { - log.info("{} cancelled but executor is no longer occupied by it.", this.cancelledWorkerOnTask); - this.cancelledWorkerOnTask = null; - } - } - - boolean result = handleStatusChange(report); + boolean result = handleReport( + heartbeat.getTaskExecutorReport(), heartbeat.getTaskExecutorID()); updateTicker(); return result; } @@ -235,11 +255,113 @@ boolean onTaskExecutorStatusChange(TaskExecutorStatusChange statusChange) { throwNotRegistered(String.format("status change %s", statusChange)); } - boolean result = handleStatusChange(statusChange.getTaskExecutorReport()); + boolean result; + try { + result = handleReport(statusChange.getTaskExecutorReport(), null); + } catch (TaskExecutorTaskCancelledException e) { + throw new IllegalStateException("Status changes cannot request cancellation", e); + } updateTicker(); return result; } + private boolean handleReport( + TaskExecutorReport report, + @Nullable io.mantisrx.server.master.resourcecluster.TaskExecutorID heartbeatTaskExecutorID) + throws TaskExecutorTaskCancelledException { + boolean heartbeat = heartbeatTaskExecutorID != null; + + if (reconciliationState != ReconciliationState.None) { + if (report instanceof Available) { + if (!heartbeat) { + clearReconciliation(); + return setAvailabilityState(AvailabilityState.pending()); + } + if (registration.reservesAcceptedTask()) { + if (reconciliationState == ReconciliationState.Verifying) { + clearReconciliation(); + return setAvailabilityState(AvailabilityState.pending()); + } + reconciliationState = ReconciliationState.Verifying; + } + return false; + } + + WorkerId reportedWorker = ((Occupied) report).getWorkerId(); + if (reconciliationState == ReconciliationState.Cancelling + && !reportedWorker.equals(cancelledWorkerOnTask)) { + log.info( + "Executor expected cancelled worker {} but reports {}; clearing stale cancellation.", + cancelledWorkerOnTask, + reportedWorker); + clearReconciliation(); + return reconcileOccupiedWorker(reportedWorker); + } + + reconciliationState = ReconciliationState.Quarantined; + cancelledWorkerOnTask = reportedWorker; + boolean changed = reconcileOccupiedWorker(reportedWorker); + if (heartbeat) { + throw cancellationException(heartbeatTaskExecutorID, reportedWorker); + } + return changed; + } + + if (report instanceof Occupied && availabilityState != null) { + WorkerId reportedWorker = ((Occupied) report).getWorkerId(); + WorkerId expectedWorker = availabilityState.getWorkerId(); + if (expectedWorker != null && !expectedWorker.equals(reportedWorker)) { + log.warn( + "Executor expected worker {} but reports {}; quarantining reported worker.", + expectedWorker, + reportedWorker); + jobMessageRouter.routeWorkerEvent( + new WorkerTerminate(expectedWorker, WorkerState.Failed, JobCompletedReason.Lost)); + quarantineWorkerOnTask(reportedWorker); + if (heartbeat) { + throw cancellationException(heartbeatTaskExecutorID, reportedWorker); + } + return true; + } + } + + if (heartbeat && availabilityState instanceof Running && report instanceof Available) { + WorkerId runningWorkerId = availabilityState.getWorkerId(); + log.warn( + "Heartbeat indicates available while running {}. Preserving ownership for reconciliation.", + runningWorkerId); + reconciliationState = registration.reservesAcceptedTask() + ? ReconciliationState.Verifying + : ReconciliationState.Quarantined; + cancelledWorkerOnTask = runningWorkerId; + jobMessageRouter.routeWorkerEvent( + new WorkerTerminate(runningWorkerId, WorkerState.Failed, JobCompletedReason.Lost)); + return false; + } + + return handleStatusChange(report); + } + + private TaskExecutorTaskCancelledException cancellationException( + io.mantisrx.server.master.resourcecluster.TaskExecutorID taskExecutorID, + WorkerId workerId) { + return new TaskExecutorTaskCancelledException( + String.format("heartbeat from %s reports cancelled task %s", taskExecutorID, workerId), + workerId); + } + + private boolean reconcileOccupiedWorker(WorkerId workerId) { + if (availabilityState instanceof Running && workerId.equals(availabilityState.getWorkerId())) { + return false; + } + return setAvailabilityState(AvailabilityState.running(workerId)); + } + + private void clearReconciliation() { + reconciliationState = ReconciliationState.None; + cancelledWorkerOnTask = null; + } + private boolean handleStatusChange(TaskExecutorReport report) throws IllegalStateException { if (availabilityState == null) { return setAvailabilityState(from(report)); @@ -281,7 +403,7 @@ private void updateTicker() { } boolean isAvailable() { - return this.availabilityState instanceof Pending && !isDisabled(); + return isRegistered() && this.availabilityState instanceof Pending && !isDisabled(); } boolean isRunningTask() { diff --git a/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/AssignmentHandlerActorTest.java b/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/AssignmentHandlerActorTest.java index a3e6f4d98..1897daf66 100644 --- a/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/AssignmentHandlerActorTest.java +++ b/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/AssignmentHandlerActorTest.java @@ -46,9 +46,11 @@ import io.mantisrx.server.master.resourcecluster.TaskExecutorRegistration; 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.google.common.collect.ImmutableMap; import java.time.Duration; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; import org.junit.AfterClass; import org.junit.Before; import org.junit.BeforeClass; @@ -187,19 +189,16 @@ public void testAssignmentRetryAndSuccess() { ActorRef parent = actorSystem.actorOf(Props.create(ForwarderParent.class, props, probe.getRef(), taskExecutorGateway)); ActorRef actor = probe.expectMsgClass(ActorRef.class); - // First call fails, second succeeds - CompletableFuture failedFuture = new CompletableFuture<>(); - failedFuture.completeExceptionally(new RuntimeException("fail")); - + CompletableFuture failedGatewayFuture = new CompletableFuture<>(); + failedGatewayFuture.completeExceptionally(new RuntimeException("fail")); when(taskExecutorGateway.submitTask(any())) - .thenReturn(failedFuture) .thenReturn(CompletableFuture.completedFuture(Ack.getInstance())); TaskExecutorAssignmentRequest request = TaskExecutorAssignmentRequest.of( createAllocationRequest(), taskExecutorID, createRegistration(), - CompletableFuture.completedFuture(taskExecutorGateway) + failedGatewayFuture ); actor.tell(request, probe.getRef()); @@ -207,8 +206,8 @@ public void testAssignmentRetryAndSuccess() { // Should eventually succeed and send no failure message probe.expectNoMessage(Duration.ofMillis(500)); - // Verify retry happened - note: retry gets a fresh gateway, so we verify at least 2 calls - verify(taskExecutorGateway, times(2)).submitTask(any()); + // The first, pre-submit gateway failure retries with a fresh gateway. + verify(taskExecutorGateway, times(1)).submitTask(any()); } @Test @@ -222,27 +221,63 @@ public void testAssignmentMaxRetriesExceeded() { 2, // 2 retries (total 2 attempts) Duration.ofMillis(50) ); - ActorRef parent = actorSystem.actorOf(Props.create(ForwarderParent.class, props, probe.getRef(), taskExecutorGateway)); + ActorRef parent = actorSystem.actorOf(Props.create(ForwarderParent.class, props, probe.getRef())); ActorRef actor = probe.expectMsgClass(ActorRef.class); - CompletableFuture failedFuture = new CompletableFuture<>(); - failedFuture.completeExceptionally(new RuntimeException("fail")); - - when(taskExecutorGateway.submitTask(any())).thenReturn(failedFuture); + CompletableFuture failedGatewayFuture = new CompletableFuture<>(); + failedGatewayFuture.completeExceptionally(new RuntimeException("fail")); TaskExecutorAssignmentRequest request = TaskExecutorAssignmentRequest.of( createAllocationRequest(), taskExecutorID, createRegistration(), - CompletableFuture.completedFuture(taskExecutorGateway) + failedGatewayFuture ); actor.tell(request, probe.getRef()); TaskExecutorAssignmentFailAndTerminate msg = probe.expectMsgClass(TaskExecutorAssignmentFailAndTerminate.class); assertEquals(taskExecutorID, msg.getTaskExecutorID()); + assertEquals(2, msg.getAttemptCount()); + + verify(taskExecutorGateway, times(0)).submitTask(any()); + } + + @Test + public void testTaskAlreadyRunningFailureIsUnwrappedAndNotRetried() { + TestKit probe = new TestKit(actorSystem); + Props props = AssignmentHandlerActor.props( + clusterID, + jobMessageRouter, + Duration.ofSeconds(10), + executeStageRequestFactory, + 3, + Duration.ofMillis(50)); + ActorRef parent = actorSystem.actorOf( + Props.create(ForwarderParent.class, props, probe.getRef(), taskExecutorGateway)); + ActorRef actor = probe.expectMsgClass(ActorRef.class); - verify(taskExecutorGateway, times(2)).submitTask(any()); + WorkerId runningWorkerId = WorkerId.fromIdUnsafe("job-1-worker-1-2"); + CompletableFuture failedFuture = new CompletableFuture<>(); + failedFuture.completeExceptionally( + new ExecutionException(new TaskAlreadyRunningException(runningWorkerId))); + when(taskExecutorGateway.submitTask(any())).thenReturn(failedFuture); + + actor.tell( + TaskExecutorAssignmentRequest.of( + createAllocationRequest(), + taskExecutorID, + createRegistration(), + CompletableFuture.completedFuture(taskExecutorGateway), + 1L), + probe.getRef()); + + TaskExecutorAssignmentFailAndTerminate message = + probe.expectMsgClass(TaskExecutorAssignmentFailAndTerminate.class); + assertEquals(1, message.getAttemptCount()); + assertTrue(message.getThrowable() instanceof TaskAlreadyRunningException); + probe.expectNoMessage(Duration.ofMillis(200)); + verify(taskExecutorGateway, times(1)).submitTask(any()); } @Test @@ -277,6 +312,39 @@ public void testAssignmentTimeout() { assertTrue(msg.getThrowable() instanceof java.util.concurrent.TimeoutException); } + @Test + public void lateGatewayCannotSubmitAfterAssignmentTimeout() { + TestKit probe = new TestKit(actorSystem); + Duration timeout = Duration.ofMillis(100); + Props props = AssignmentHandlerActor.props( + clusterID, + jobMessageRouter, + timeout, + executeStageRequestFactory, + 1, + Duration.ofMillis(50)); + ActorRef parent = actorSystem.actorOf( + Props.create(ForwarderParent.class, props, probe.getRef(), taskExecutorGateway)); + ActorRef actor = probe.expectMsgClass(ActorRef.class); + CompletableFuture delayedGateway = new CompletableFuture<>(); + + actor.tell( + TaskExecutorAssignmentRequest.of( + createAllocationRequest(), + taskExecutorID, + createRegistration(), + delayedGateway), + probe.getRef()); + + TaskExecutorAssignmentFailAndTerminate message = + probe.expectMsgClass(TaskExecutorAssignmentFailAndTerminate.class); + assertEquals(AssignmentHandlerActor.AssignmentFailureType.NotSent, message.getFailureType()); + + delayedGateway.complete(taskExecutorGateway); + probe.expectNoMessage(Duration.ofMillis(200)); + verify(taskExecutorGateway, times(0)).submitTask(any()); + } + public static class ForwarderParent extends AbstractActor { private final ActorRef probe; private final Props childProps; @@ -308,9 +376,9 @@ public Receive createReceive() { if (taskExecutorGateway != null) { sender().tell(CompletableFuture.completedFuture(taskExecutorGateway), self()); } else { - // Fallback: create a completed future with a mock gateway - CompletableFuture gatewayFuture = CompletableFuture.completedFuture( - mock(TaskExecutorGateway.class)); + CompletableFuture gatewayFuture = new CompletableFuture<>(); + gatewayFuture.completeExceptionally( + new RuntimeException("gateway unavailable")); sender().tell(gatewayFuture, self()); } }) @@ -319,4 +387,3 @@ public Receive createReceive() { } } } - diff --git a/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/DisableReconnectExpireStuckReproTest.java b/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/DisableReconnectExpireStuckReproTest.java index 0561cb76b..54707243d 100644 --- a/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/DisableReconnectExpireStuckReproTest.java +++ b/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/DisableReconnectExpireStuckReproTest.java @@ -226,6 +226,9 @@ private void runReproSequence( assertEquals(Ack.getInstance(), resourceCluster.heartBeatFromTaskExecutor( new TaskExecutorHeartbeat(TE_ID, CLUSTER_ID, TaskExecutorReport.available())).get()); + assertEquals(Ack.getInstance(), + resourceCluster.heartBeatFromTaskExecutor( + new TaskExecutorHeartbeat(TE_ID, CLUSTER_ID, TaskExecutorReport.available())).get()); // 5) Wait past expiry so ExpireDisableTaskExecutorsRequest fires. long elapsedMs = Instant.now().toEpochMilli() - disableAt.toEpochMilli(); @@ -287,7 +290,8 @@ private TaskExecutorRegistration buildRegistration(TaskExecutorID id) { .taskExecutorAttributes( ImmutableMap.of( WorkerConstants.WORKER_CONTAINER_DEFINITION_ID, SKU_ID.getResourceID(), - "repro", "attr")) + "repro", "attr", + TaskExecutorRegistration.ACCEPTED_TASK_RESERVATION_ATTRIBUTE, "true")) .build(); } } diff --git a/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/ReservationRegistryActorIntegrationTest.java b/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/ReservationRegistryActorIntegrationTest.java index b8e414510..069883b4b 100644 --- a/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/ReservationRegistryActorIntegrationTest.java +++ b/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/ReservationRegistryActorIntegrationTest.java @@ -52,6 +52,7 @@ import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import io.mantisrx.server.core.ExecuteStageRequest; @@ -345,12 +346,11 @@ public void testReservationCancellation() throws Exception { public void testAssignmentRetry() throws Exception { TestKit probe = new TestKit(system); - // Reset gateway mock to have different behavior for this test - // 1. Fail first - // 2. Succeed second - when(gateway.submitTask(any())) - .thenReturn(CompletableFuture.failedFuture(new RuntimeException("Connection failed"))) - .thenReturn(CompletableFuture.completedFuture(Ack.getInstance())); + AtomicInteger connectionAttempts = new AtomicInteger(); + rpcService.setRpcGatewayFutureFunction(ignored -> + connectionAttempts.getAndIncrement() == 0 + ? CompletableFuture.failedFuture(new RuntimeException("Connection failed")) + : CompletableFuture.completedFuture(gateway)); TaskExecutorRegistration registration = TaskExecutorRegistration.builder() .taskExecutorID(TASK_EXECUTOR_ID_1) @@ -390,8 +390,8 @@ public void testAssignmentRetry() throws Exception { assertTrue(view.getGroups().isEmpty() || view.getGroups().values().stream() .allMatch(g -> g.getReservationCount() == 0)); - // Verify gateway was called twice - verify(gateway, timeout(5000).times(2)).submitTask(any()); + assertTrue(connectionAttempts.get() >= 2); + verify(gateway, timeout(5000).times(1)).submitTask(any()); // Verify WorkerLaunched event was routed (eventually success) verify(jobMessageRouter, timeout(5000).times(1)).routeWorkerEvent(any(WorkerLaunched.class)); @@ -401,9 +401,11 @@ public void testAssignmentRetry() throws Exception { public void testAssignmentFailure() throws Exception { TestKit probe = new TestKit(system); - // Fail always - when(gateway.submitTask(any())) - .thenReturn(CompletableFuture.failedFuture(new RuntimeException("Permanent failure"))); + AtomicInteger connectionAttempts = new AtomicInteger(); + rpcService.setRpcGatewayFutureFunction(ignored -> { + connectionAttempts.incrementAndGet(); + return CompletableFuture.failedFuture(new RuntimeException("Permanent failure")); + }); TaskExecutorRegistration registration = TaskExecutorRegistration.builder() .taskExecutorID(TASK_EXECUTOR_ID_2) @@ -437,8 +439,8 @@ public void testAssignmentFailure() throws Exception { // Wait for retries to exhaust Thread.sleep(1000); - // Verify gateway called multiple times (3 retries) - verify(gateway, timeout(5000).atLeast(3)).submitTask(any()); + assertTrue(connectionAttempts.get() >= 3); + verify(gateway, times(0)).submitTask(any()); // Verify WorkerLaunchFailed event was routed verify(jobMessageRouter, timeout(5000).times(1)).routeWorkerEvent(any(WorkerLaunchFailed.class)); diff --git a/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/ResourceClusterActorTest.java b/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/ResourceClusterActorTest.java index f8cea27ff..3869c4089 100644 --- a/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/ResourceClusterActorTest.java +++ b/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/ResourceClusterActorTest.java @@ -67,6 +67,7 @@ import io.mantisrx.server.master.resourcecluster.TaskExecutorRegistration; import io.mantisrx.server.master.resourcecluster.TaskExecutorReport; import io.mantisrx.server.master.resourcecluster.TaskExecutorStatusChange; +import io.mantisrx.server.master.resourcecluster.TaskExecutorTaskCancelledException; import io.mantisrx.server.master.scheduler.JobMessageRouter; import io.mantisrx.server.master.scheduler.WorkerEvent; import io.mantisrx.server.master.scheduler.WorkerOnDisabledVM; @@ -170,6 +171,8 @@ public class ResourceClusterActorTest { private static final WorkerId WORKER_ID = WorkerId.fromIdUnsafe("late-sine-function-tutorial-1-worker-0-1"); + private static final WorkerId WORKER_ID_2 = + WorkerId.fromIdUnsafe("late-sine-function-tutorial-1-worker-0-2"); private static final JobMetadata JOB_METADATA = new JobMetadata(WORKER_ID.getJobId(), null, null, 1, "testuser", null, ImmutableList.of(), -1, -1, -1); private static ActorSystem actorSystem; @@ -309,11 +312,11 @@ public void testMarkTaskCancelled() throws Exception { tEStatus = resourceCluster.getTaskExecutorState(TASK_EXECUTOR_ID).get(); assertEquals(WORKER_ID, tEStatus.getCancelledWorkerId()); - // new heartbeat with available state reset the cancelled worker state. + // The agent's status change is emitted after cancellation cleanup completes. assertEquals(Ack.getInstance(), resourceCluster - .heartBeatFromTaskExecutor( - new TaskExecutorHeartbeat( + .notifyTaskExecutorStatusChange( + new TaskExecutorStatusChange( TASK_EXECUTOR_ID, CLUSTER_ID, TaskExecutorReport.available())).get()); @@ -322,6 +325,159 @@ public void testMarkTaskCancelled() throws Exception { assertEquals(null, tEStatus.getCancelledWorkerId()); } + @Test + public void cancellationOwnershipSurvivesDisconnectAndReconnect() throws Exception { + assertEquals(Ack.getInstance(), resourceCluster.registerTaskExecutor(TASK_EXECUTOR_REGISTRATION).get()); + assertEquals( + Ack.getInstance(), + resourceCluster.initializeTaskExecutor(TASK_EXECUTOR_ID, WORKER_ID).get()); + assertEquals(Ack.getInstance(), resourceCluster.markTaskExecutorWorkerCancelled(WORKER_ID).get()); + assertEquals( + WORKER_ID, + resourceCluster.getTaskExecutorState(TASK_EXECUTOR_ID).get().getCancelledWorkerId()); + + assertEquals( + Ack.getInstance(), + resourceCluster.disconnectTaskExecutor( + new TaskExecutorDisconnection(TASK_EXECUTOR_ID, CLUSTER_ID)).get()); + assertEquals(Ack.getInstance(), resourceCluster.registerTaskExecutor(TASK_EXECUTOR_REGISTRATION).get()); + + assertEquals( + "disconnect/reconnect discarded cancellation ownership", + WORKER_ID, + resourceCluster.getTaskExecutorState(TASK_EXECUTOR_ID).get().getCancelledWorkerId()); + } + + @Test + public void delayedAssignmentFailureCannotDisconnectNewerAssignment() throws Exception { + assertEquals(Ack.getInstance(), resourceCluster.registerTaskExecutor(TASK_EXECUTOR_REGISTRATION).get()); + assertEquals( + Ack.getInstance(), + resourceCluster.heartBeatFromTaskExecutor( + new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, + CLUSTER_ID, + TaskExecutorReport.available())).get()); + + TaskExecutorAllocationRequest firstAssignment = TaskExecutorAllocationRequest.of( + WORKER_ID, + SchedulingConstraints.of(MACHINE_DEFINITION), + JOB_METADATA, + 0, + MantisJobDurationType.Perpetual); + assertEquals( + TASK_EXECUTOR_ID, + resourceCluster.getTaskExecutorsFor(Collections.singleton(firstAssignment)) + .get().values().iterator().next()); + + ActorRef executorStateManagerActor = actorSystem.actorSelection( + resourceClusterActor.path().child("executorStateManager-clusterId")) + .resolveOne(Duration.ofSeconds(2)).toCompletableFuture().get(); + TestKit probe = new TestKit(actorSystem); + executorStateManagerActor.tell( + new AssignmentHandlerActor.TaskExecutorAssignmentFailAndTerminate( + TASK_EXECUTOR_ID, + firstAssignment, + new RuntimeException("first assignment was not sent"), + 1, + 1L, + AssignmentHandlerActor.AssignmentFailureType.NotSent), + probe.getRef()); + + long disconnectedDeadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2); + while (!resourceCluster.getRegisteredTaskExecutors().get().isEmpty() + && System.nanoTime() < disconnectedDeadline) { + Thread.sleep(20); + } + assertEquals(Ack.getInstance(), resourceCluster.registerTaskExecutor(TASK_EXECUTOR_REGISTRATION).get()); + assertEquals( + Ack.getInstance(), + resourceCluster.heartBeatFromTaskExecutor( + new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, + CLUSTER_ID, + TaskExecutorReport.available())).get()); + Thread.sleep(150); + + TaskExecutorAllocationRequest newerAssignment = TaskExecutorAllocationRequest.of( + WORKER_ID_2, + SchedulingConstraints.of(MACHINE_DEFINITION), + new JobMetadata(WORKER_ID_2.getJobId(), null, null, 1, "testuser", null, + ImmutableList.of(), -1, -1, -1), + 0, + MantisJobDurationType.Perpetual); + assertEquals( + TASK_EXECUTOR_ID, + resourceCluster.getTaskExecutorsFor(Collections.singleton(newerAssignment)) + .get().values().iterator().next()); + + executorStateManagerActor.tell( + new AssignmentHandlerActor.TaskExecutorAssignmentFailAndTerminate( + TASK_EXECUTOR_ID, + firstAssignment, + new RuntimeException("delayed failure for first assignment"), + 1, + 1L, + AssignmentHandlerActor.AssignmentFailureType.NotSent), + probe.getRef()); + executorStateManagerActor.tell( + new GetTaskExecutorStatusRequest(TASK_EXECUTOR_ID, CLUSTER_ID), + probe.getRef()); + TaskExecutorStatus statusAfterDelayedFailure = + probe.expectMsgClass(TaskExecutorStatus.class); + + assertEquals( + "a delayed terminal result cleared the newer assignment", + WORKER_ID_2, + statusAfterDelayedFailure.getWorkerId()); + assertTrue(statusAfterDelayedFailure.isRegistered()); + assertEquals(TASK_EXECUTOR_ID, resourceCluster.getTaskExecutorAssignedFor(WORKER_ID_2).get()); + } + + @Test + public void assignmentTimeoutKeepsAmbiguousWorkerFenced() throws Exception { + CompletableFuture hangingFuture = new CompletableFuture<>(); + when(gateway.submitTask(ArgumentMatchers.any())).thenReturn(hangingFuture); + + assertEquals(Ack.getInstance(), resourceCluster.registerTaskExecutor(TASK_EXECUTOR_REGISTRATION).get()); + assertEquals(Ack.getInstance(), resourceCluster.heartBeatFromTaskExecutor( + new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, CLUSTER_ID, TaskExecutorReport.available())).get()); + + TaskExecutorAllocationRequest assignment = TaskExecutorAllocationRequest.of( + WORKER_ID, + SchedulingConstraints.of(MACHINE_DEFINITION), + JOB_METADATA, + 0, + MantisJobDurationType.Perpetual); + assertEquals( + TASK_EXECUTOR_ID, + resourceCluster.getTaskExecutorsFor(Collections.singleton(assignment)) + .get().values().iterator().next()); + + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3); + TaskExecutorStatus status = resourceCluster.getTaskExecutorState(TASK_EXECUTOR_ID).get(); + while (!WORKER_ID.equals(status.getCancelledWorkerId()) && System.nanoTime() < deadline) { + Thread.sleep(20); + status = resourceCluster.getTaskExecutorState(TASK_EXECUTOR_ID).get(); + } + + assertEquals(WORKER_ID, status.getWorkerId()); + assertEquals(WORKER_ID, status.getCancelledWorkerId()); + assertEquals(ImmutableList.of(), resourceCluster.getAvailableTaskExecutors().get()); + + try { + resourceCluster.heartBeatFromTaskExecutor(new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, + CLUSTER_ID, + TaskExecutorReport.occupied(WORKER_ID))).get(); + Assert.fail("Expected the ambiguously accepted worker to be cancelled"); + } catch (ExecutionException e) { + assertTrue(ExceptionUtils.stripExecutionException(e) + instanceof TaskExecutorTaskCancelledException); + } + } + @Test public void testInitializationAfterRestart() throws Exception { doReturn(ImmutableList.of()) @@ -566,19 +722,13 @@ public void testAssignmentTimeout() throws Exception { assertEquals(ImmutableList.of(), resourceCluster.getAvailableTaskExecutors().get()); Thread.sleep(2000); - assertEquals(ImmutableList.of(), resourceCluster.getRegisteredTaskExecutors().get()); + assertEquals(ImmutableList.of(TASK_EXECUTOR_ID), resourceCluster.getRegisteredTaskExecutors().get()); assertEquals(ImmutableList.of(), resourceCluster.getAvailableTaskExecutors().get()); - // Restore the submitTask mock to return a completed future for re-registration - when(gateway.submitTask(ArgumentMatchers.any())).thenReturn(CompletableFuture.completedFuture(Ack.getInstance())); - - // re-register TE - when(mantisJobStore.getTaskExecutor(TASK_EXECUTOR_ID)) - .thenReturn(new TaskExecutorRegistration(TASK_EXECUTOR_ID, CLUSTER_ID, "", "", null, MACHINE_DEFINITION, ATTRIBUTES)); assertEquals(Ack.getInstance(), resourceCluster - .heartBeatFromTaskExecutor( - new TaskExecutorHeartbeat( + .notifyTaskExecutorStatusChange( + new TaskExecutorStatusChange( TASK_EXECUTOR_ID, CLUSTER_ID, TaskExecutorReport.available())).get()); diff --git a/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/TaskExecutorStateTest.java b/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/TaskExecutorStateTest.java index 213e052ee..b944cc236 100644 --- a/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/TaskExecutorStateTest.java +++ b/mantis-control-plane/mantis-control-plane-server/src/test/java/io/mantisrx/master/resourcecluster/TaskExecutorStateTest.java @@ -19,11 +19,15 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import io.mantisrx.common.WorkerPorts; import io.mantisrx.common.util.DelegateClock; +import io.mantisrx.master.jobcluster.job.worker.WorkerTerminate; import io.mantisrx.runtime.MachineDefinition; +import io.mantisrx.server.core.JobCompletedReason; import io.mantisrx.server.core.TestingRpcService; import io.mantisrx.server.core.domain.WorkerId; import io.mantisrx.server.master.resourcecluster.ClusterID; @@ -34,6 +38,7 @@ import io.mantisrx.server.master.resourcecluster.TaskExecutorStatusChange; import io.mantisrx.server.master.resourcecluster.TaskExecutorTaskCancelledException; import io.mantisrx.server.master.scheduler.JobMessageRouter; +import io.mantisrx.server.master.scheduler.WorkerEvent; import io.mantisrx.server.worker.TaskExecutorGateway; import io.mantisrx.shaded.com.google.common.collect.ImmutableList; import io.mantisrx.shaded.com.google.common.collect.ImmutableMap; @@ -45,6 +50,7 @@ import java.util.concurrent.atomic.AtomicReference; import org.junit.Before; import org.junit.Test; +import org.mockito.ArgumentCaptor; public class TaskExecutorStateTest { private final AtomicReference actual = @@ -68,6 +74,7 @@ public class TaskExecutorStateTest { private static final Map ATTRIBUTES = ImmutableMap.of("attr1", "attr2"); private static final WorkerId WORKER_ID = WorkerId.fromIdUnsafe("late-sine-function-tutorial-1-worker-0-1"); + private static final WorkerId WORKER_ID_2 = WorkerId.fromIdUnsafe("late-sine-function-tutorial-1-worker-0-2"); @Before public void setup() { @@ -157,6 +164,233 @@ public void testInitializationLifecycle() throws TaskExecutorTaskCancelledExcept assertEquals(currentTime, state.getLastActivity()); } + @Test + public void availableHeartbeatCannotReleaseStillOwnedExecutor() + throws TaskExecutorTaskCancelledException { + registerAndStartWorker(WORKER_ID); + + assertFalse(state.onHeartbeat( + new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, + CLUSTER_ID, + TaskExecutorReport.available()))); + + assertEquals("stale Available heartbeat released the owned executor", WORKER_ID, state.getWorkerId()); + assertEquals(WORKER_ID, state.getCancelledWorkerId()); + assertFalse(state.isAvailable()); + + ArgumentCaptor eventCaptor = ArgumentCaptor.forClass(WorkerEvent.class); + verify(router).routeWorkerEvent(eventCaptor.capture()); + WorkerTerminate termination = (WorkerTerminate) eventCaptor.getValue(); + assertEquals(WORKER_ID, termination.getWorkerId()); + assertEquals(JobCompletedReason.Lost, termination.getReason()); + } + + @Test + public void authoritativeAvailableStatusReleasesQuarantinedExecutor() + throws TaskExecutorTaskCancelledException { + registerAndStartWorker(WORKER_ID); + state.onHeartbeat(new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, CLUSTER_ID, TaskExecutorReport.available())); + + assertTrue(state.onTaskExecutorStatusChange(new TaskExecutorStatusChange( + TASK_EXECUTOR_ID, CLUSTER_ID, TaskExecutorReport.available()))); + + assertTrue(state.isAvailable()); + assertEquals(null, state.getWorkerId()); + assertEquals(null, state.getCancelledWorkerId()); + } + + @Test + public void reservationAwareSecondAvailableHeartbeatReleasesQuarantine() + throws TaskExecutorTaskCancelledException { + registerAndStartWorker(WORKER_ID, true); + + assertFalse(state.onHeartbeat(new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, CLUSTER_ID, TaskExecutorReport.available()))); + assertTrue(state.onHeartbeat(new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, CLUSTER_ID, TaskExecutorReport.available()))); + + assertTrue(state.isAvailable()); + assertEquals(null, state.getCancelledWorkerId()); + } + + @Test + public void legacyRepeatedAvailableHeartbeatsRemainQuarantined() + throws TaskExecutorTaskCancelledException { + registerAndStartWorker(WORKER_ID); + + assertFalse(state.onHeartbeat(new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, CLUSTER_ID, TaskExecutorReport.available()))); + assertFalse(state.onHeartbeat(new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, CLUSTER_ID, TaskExecutorReport.available()))); + + assertEquals(WORKER_ID, state.getWorkerId()); + assertEquals(WORKER_ID, state.getCancelledWorkerId()); + assertFalse(state.isAvailable()); + } + + @Test + public void availableHeartbeatCannotClearLegacyPreparationCancellation() + throws TaskExecutorTaskCancelledException { + registerAndAssignWorker(WORKER_ID); + state.setCancelledWorkerOnTask(WORKER_ID); + + state.onHeartbeat( + new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, + CLUSTER_ID, + TaskExecutorReport.available())); + + assertEquals( + "Available alone lost cancellation ownership while the accepted task can still start", + WORKER_ID, + state.getCancelledWorkerId()); + assertEquals(WORKER_ID, state.getWorkerId()); + assertFalse(state.isAvailable()); + } + + @Test + public void occupiedWorkerMismatchWhileAssignedUsesReportedIdentity() + throws TaskExecutorTaskCancelledException { + registerAndAssignWorker(WORKER_ID); + + state.onTaskExecutorStatusChange( + new TaskExecutorStatusChange( + TASK_EXECUTOR_ID, + CLUSTER_ID, + TaskExecutorReport.occupied(WORKER_ID_2))); + + assertEquals( + "Occupied(B) was silently attributed to assigned worker A", + WORKER_ID_2, + state.getWorkerId()); + } + + @Test + public void occupiedWorkerMismatchWhileRunningUsesReportedIdentity() + throws TaskExecutorTaskCancelledException { + registerAndStartWorker(WORKER_ID); + + try { + state.onHeartbeat( + new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, + CLUSTER_ID, + TaskExecutorReport.occupied(WORKER_ID_2))); + fail("mismatched worker should be cancelled"); + } catch (TaskExecutorTaskCancelledException e) { + assertTrue(e.getMessage().contains(WORKER_ID_2.toString())); + } + + assertEquals( + "Occupied(B) was silently attributed to running worker A", + WORKER_ID_2, + state.getWorkerId()); + } + + @Test + public void staleCancellationDoesNotCancelDifferentReportedWorker() + throws TaskExecutorTaskCancelledException { + registerAndStartWorker(WORKER_ID); + state.setCancelledWorkerOnTask(WORKER_ID); + + assertTrue(state.onHeartbeat( + new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, + CLUSTER_ID, + TaskExecutorReport.occupied(WORKER_ID_2)))); + + assertEquals(WORKER_ID_2, state.getWorkerId()); + assertEquals(null, state.getCancelledWorkerId()); + } + + @Test + public void availableHeartbeatRestoresDisconnectedIdleExecutor() + throws TaskExecutorTaskCancelledException { + assertTrue(state.onRegistration(registration())); + assertTrue(state.onHeartbeat(new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, CLUSTER_ID, TaskExecutorReport.available()))); + + assertTrue(state.onDisconnection()); + assertTrue(state.onRegistration(registration())); + assertTrue(state.onHeartbeat(new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, CLUSTER_ID, TaskExecutorReport.available()))); + + assertTrue(state.isAvailable()); + } + + @Test(expected = TaskExecutorTaskCancelledException.class) + public void cancellationOwnershipSurvivesDisconnectAndReconnect() + throws TaskExecutorTaskCancelledException { + registerAndStartWorker(WORKER_ID); + state.setCancelledWorkerOnTask(WORKER_ID); + + assertTrue(state.onDisconnection()); + assertEquals(WORKER_ID, state.getCancelledWorkerId()); + assertTrue(state.onRegistration(registration())); + assertEquals(WORKER_ID, state.getCancelledWorkerId()); + + state.onHeartbeat( + new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, + CLUSTER_ID, + TaskExecutorReport.occupied(WORKER_ID))); + } + + private void registerAndStartWorker(WorkerId workerId) + throws TaskExecutorTaskCancelledException { + registerAndStartWorker(workerId, false); + } + + private void registerAndStartWorker(WorkerId workerId, boolean reservesAcceptedTask) + throws TaskExecutorTaskCancelledException { + registerAndAssignWorker(workerId, reservesAcceptedTask); + assertTrue(state.onTaskExecutorStatusChange( + new TaskExecutorStatusChange( + TASK_EXECUTOR_ID, + CLUSTER_ID, + TaskExecutorReport.occupied(workerId)))); + } + + private void registerAndAssignWorker(WorkerId workerId) + throws TaskExecutorTaskCancelledException { + registerAndAssignWorker(workerId, false); + } + + private void registerAndAssignWorker(WorkerId workerId, boolean reservesAcceptedTask) + throws TaskExecutorTaskCancelledException { + assertTrue(state.onRegistration(registration(reservesAcceptedTask))); + assertTrue(state.onHeartbeat( + new TaskExecutorHeartbeat( + TASK_EXECUTOR_ID, + CLUSTER_ID, + TaskExecutorReport.available()))); + assertTrue(state.onAssignment(workerId)); + } + + private TaskExecutorRegistration registration() { + return registration(false); + } + + private TaskExecutorRegistration registration(boolean reservesAcceptedTask) { + Map attributes = ImmutableMap.builder() + .putAll(ATTRIBUTES) + .put( + TaskExecutorRegistration.ACCEPTED_TASK_RESERVATION_ATTRIBUTE, + Boolean.toString(reservesAcceptedTask)) + .build(); + return TaskExecutorRegistration.builder() + .taskExecutorID(TASK_EXECUTOR_ID) + .clusterID(CLUSTER_ID) + .taskExecutorAddress(TASK_EXECUTOR_ADDRESS) + .hostname(HOST_NAME) + .workerPorts(WORKER_PORTS) + .machineDefinition(MACHINE_DEFINITION) + .taskExecutorAttributes(attributes) + .build(); + } + private Instant tick() { return actual.updateAndGet(currentTime -> Clock.offset(currentTime, Duration.ofSeconds(1))).instant(); } diff --git a/mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/TaskExecutor.java b/mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/TaskExecutor.java index 8f2cc779a..737b547a4 100644 --- a/mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/TaskExecutor.java +++ b/mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/TaskExecutor.java @@ -55,7 +55,9 @@ import io.mantisrx.shaded.com.google.common.util.concurrent.Service; import io.mantisrx.shaded.com.google.common.util.concurrent.Service.State; import io.mantisrx.shaded.org.apache.curator.shaded.com.google.common.annotations.VisibleForTesting; +import java.io.Closeable; import java.io.File; +import java.io.IOException; import java.util.Map; import java.util.Optional; import java.util.concurrent.Callable; @@ -64,6 +66,8 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; import java.util.stream.Collectors; import javax.annotation.Nullable; @@ -87,6 +91,10 @@ */ @Slf4j public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway { + private static final String LEGACY_REGISTRATION_STATE_FILE = "rmCxnState.txt"; + private static final String REGISTRATION_STATE_FILE = + "rmCxnState-accepted-task-reservation-v1.txt"; + @Getter private final TaskExecutorID taskExecutorID; @Getter @@ -116,10 +124,28 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway { private final TaskFactory taskFactory; - private RuntimeTask currentTask; - private ExecuteStageRequest currentRequest; + private final AtomicReference taskSlot = new AtomicReference<>(); + private volatile boolean stopping; private final DurableBooleanState registeredState; + private static final class TaskSlot { + private final ExecuteStageRequest request; + private final WorkerId workerId; + private final AtomicBoolean cancelled = new AtomicBoolean(); + private final AtomicBoolean cleanupStarted = new AtomicBoolean(); + private final CompletableFuture terminated = new CompletableFuture<>(); + @Nullable + private volatile RuntimeTask task; + @Nullable + private volatile UserCodeClassLoader userCodeClassLoader; + private volatile boolean prepared; + + private TaskSlot(ExecuteStageRequest request) { + this.request = request; + this.workerId = request.getWorkerId(); + } + } + @VisibleForTesting public TaskExecutor( RpcService rpcService, @@ -161,6 +187,18 @@ public TaskExecutor( MachineDefinition machineDefinition = MachineDefinitionUtils.from(workerConfiguration, workerPorts); String hostName = workerConfiguration.getExternalAddress(); + Map taskExecutorAttributes = + workerConfiguration + .getTaskExecutorAttributes() + .entrySet() + .stream() + .collect(Collectors., String, String>toMap( + kv -> kv.getKey().toLowerCase(), + kv -> kv.getValue().toLowerCase())); + taskExecutorAttributes.put( + TaskExecutorRegistration.ACCEPTED_TASK_RESERVATION_ATTRIBUTE, + Boolean.TRUE.toString()); + this.taskExecutorRegistration = TaskExecutorRegistration.builder() .machineDefinition(machineDefinition) @@ -169,14 +207,7 @@ public TaskExecutor( .hostname(hostName) .taskExecutorAddress(getAddress()) .workerPorts(workerPorts) - .taskExecutorAttributes(ImmutableMap.copyOf( - workerConfiguration - .getTaskExecutorAttributes() - .entrySet() - .stream() - .collect(Collectors., String, String>toMap( - kv -> kv.getKey().toLowerCase(), - kv -> kv.getValue().toLowerCase())))) + .taskExecutorAttributes(ImmutableMap.copyOf(taskExecutorAttributes)) .build(); log.info("Starting executor registration: {}", this.taskExecutorRegistration); @@ -190,9 +221,16 @@ public TaskExecutor( this.resourceManagerCxnIdx = 0; this.taskFactory = taskFactory == null ? new SingleTaskOnlyFactory() : taskFactory; + File registrationStore = workerConfiguration.getRegistrationStoreDir(); + DurableBooleanState legacyRegisteredState = new DurableBooleanState( + new File(registrationStore, LEGACY_REGISTRATION_STATE_FILE).getAbsolutePath()); + try { + legacyRegisteredState.setState(false); + } catch (IOException e) { + throw new IllegalStateException("Failed to reset legacy registration marker", e); + } this.registeredState = new DurableBooleanState( - new File(workerConfiguration.getRegistrationStoreDir(), - "rmCxnState.txt").getAbsolutePath()); + new File(registrationStore, REGISTRATION_STATE_FILE).getAbsolutePath()); this.rpcCallTimeoutMsDp = ConfigUtils.getDynamicPropertyLong("heartbeatTimeoutMs", WorkerConfiguration.class, workerConfiguration.heartbeatTimeoutMs(), this.dynamicPropertiesLoader); @@ -349,10 +387,11 @@ private ExecutorService getRuntimeExecutor() { CompletableFuture getCurrentReport() { return callAsync(() -> { - if (this.currentTask == null) { + TaskSlot slot = this.taskSlot.get(); + if (slot == null) { return TaskExecutorReport.available(); } else { - return TaskExecutorReport.occupied(WorkerId.fromIdUnsafe(currentTask.getWorkerId())); + return TaskExecutorReport.occupied(slot.workerId); } }, Time.milliseconds(this.rpcCallTimeoutMsDp.getValue())); } @@ -367,17 +406,31 @@ CompletableFuture callInMainThread(Callable> tSuppli public CompletableFuture submitTask(ExecuteStageRequest request) { log.info("Received request {} for execution", request); - if (currentTask != null) { - if (currentTask.getWorkerId().equals(request.getWorkerId().getId())) { - return CompletableFuture.completedFuture(Ack.getInstance()); - } else { + TaskSlot newSlot; + while (true) { + TaskSlot existingSlot = taskSlot.get(); + if (existingSlot != null) { + if (existingSlot.workerId.equals(request.getWorkerId())) { + return CompletableFuture.completedFuture(Ack.getInstance()); + } return CompletableFutures.exceptionallyCompletedFuture( - new TaskAlreadyRunningException(WorkerId.fromIdUnsafe(currentTask.getWorkerId()))); + new TaskAlreadyRunningException(existingSlot.workerId)); + } + + newSlot = new TaskSlot(request); + if (taskSlot.compareAndSet(null, newSlot)) { + break; } } - // do not wait for task processing (e.g. artifact download). - getIOExecutor().execute(() -> this.prepareTask(request)); + setStatus(TaskExecutorReport.occupied(newSlot.workerId)); + final TaskSlot slotToPrepare = newSlot; + try { + getIOExecutor().execute(() -> prepareTask(slotToPrepare)); + } catch (RuntimeException e) { + clearTaskSlot(newSlot); + return CompletableFutures.exceptionallyCompletedFuture(e); + } return CompletableFuture.completedFuture(Ack.getInstance()); } @@ -389,77 +442,156 @@ public CompletableFuture cacheJobArtifacts(CacheJobArtifactsRequest request return CompletableFuture.completedFuture(Ack.getInstance()); } - private void prepareTask(ExecuteStageRequest request) { + private void prepareTask(TaskSlot slot) { try { - this.currentRequest = request; - - UserCodeClassLoader userCodeClassLoader = this.taskFactory.getUserCodeClassLoader( - request, classLoaderHandle); - ClassLoader cl = userCodeClassLoader.asClassLoader(); + slot.userCodeClassLoader = this.taskFactory.getUserCodeClassLoader( + slot.request, classLoaderHandle); + ClassLoader cl = slot.userCodeClassLoader.asClassLoader(); // There should only be 1 task implementation provided by mantis-server-worker. JsonSerializer ser = new JsonSerializer(); - String executeRequest = ser.toJson(request); + String executeRequest = ser.toJson(slot.request); String configString = ser.toJson(WorkerConfigurationUtils.toWritable(workerConfiguration)); - RuntimeTask task = this.taskFactory.getRuntimeTaskInstance(request, cl); + slot.task = this.taskFactory.getRuntimeTaskInstance(slot.request, cl); - task.initialize( + slot.task.initialize( executeRequest, configString, - userCodeClassLoader); - - scheduleRunAsync(() -> { - setCurrentTask(task); - startCurrentTask(); - }, 0, TimeUnit.MILLISECONDS); - } catch (Exception ex) { - log.error("Failed to submit task, request: {}", request, ex); - final Status failedStatus = new Status(currentRequest.getJobId(), currentRequest.getStage(), currentRequest.getWorkerIndex(), currentRequest.getWorkerNumber(), - Status.TYPE.INFO, "stage " + currentRequest.getStage() + " worker index=" + currentRequest.getWorkerIndex() + " number=" + currentRequest.getWorkerNumber() + " failed during initialization", + slot.userCodeClassLoader); + + slot.prepared = true; + scheduleRunAsync(() -> onTaskPrepared(slot), 0, TimeUnit.MILLISECONDS); + } catch (Throwable ex) { + slot.prepared = true; + log.error("Failed to submit task, request: {}", slot.request, ex); + final Status failedStatus = new Status(slot.request.getJobId(), slot.request.getStage(), slot.request.getWorkerIndex(), slot.request.getWorkerNumber(), + Status.TYPE.INFO, "stage " + slot.request.getStage() + " worker index=" + slot.request.getWorkerIndex() + " number=" + slot.request.getWorkerNumber() + " failed during initialization", MantisJobState.Failed); updateExecutionStatus(failedStatus); - listeners.enqueue(getTaskFailedEvent(null, ex)); + scheduleRunAsync(() -> onTaskPreparationFailed(slot, ex), 0, TimeUnit.MILLISECONDS); } finally { getIOExecutor().execute(listeners::dispatch); } + } + private void onTaskPrepared(TaskSlot slot) { + validateRunsInMainThread(); + if (taskSlot.get() != slot || slot.cancelled.get() || stopping) { + disposeTaskSlot(slot, true); + return; + } + startTask(slot); + } + private void onTaskPreparationFailed(TaskSlot slot, Throwable throwable) { + validateRunsInMainThread(); + if (!slot.cancelled.get()) { + listeners.enqueue(getTaskFailedEvent(slot.task, throwable)); + getIOExecutor().execute(listeners::dispatch); + } + disposeTaskSlot(slot, slot.cancelled.get()); } - private void startCurrentTask() { + private void startTask(TaskSlot slot) { validateRunsInMainThread(); + RuntimeTask task = slot.task; + if (task == null) { + onTaskPreparationFailed(slot, new IllegalStateException("Prepared task was null")); + return; + } - if (currentTask.state().equals(State.NEW)) { - listeners.enqueue(getTaskStartingEvent(currentTask)); + if (task.state().equals(State.NEW)) { + listeners.enqueue(getTaskStartingEvent(task)); getIOExecutor().execute(listeners::dispatch); - CompletableFuture currentTaskSuccessfullyStartFuture = - Services.startAsync(currentTask, getRuntimeExecutor()); + CompletableFuture startFuture = Services.startAsync(task, getRuntimeExecutor()); - currentTaskSuccessfullyStartFuture + startFuture .whenCompleteAsync((dontCare, throwable) -> { if (throwable != null) { - // okay failed to start task successfully - // lets stop it log.error("TaskExecutor failed to start", throwable); - RuntimeTask task = currentTask; - setCurrentTask(null); setPreviousFailure(throwable); listeners.enqueue(getTaskFailedEvent(task, throwable)); getIOExecutor().execute(listeners::dispatch); + disposeTaskSlot(slot, false); } }, getMainThreadExecutor()); } } - private void setCurrentTask(@Nullable RuntimeTask task) { + private CompletableFuture disposeTaskSlot(TaskSlot slot, boolean cancellation) { validateRunsInMainThread(); + if (!slot.cleanupStarted.compareAndSet(false, true)) { + return slot.terminated; + } - this.currentTask = task; - if (task == null) { - setStatus(TaskExecutorReport.available()); + RuntimeTask task = slot.task; + if (cancellation && task != null) { + listeners.enqueue(getTaskCancellingEvent(task)); + getIOExecutor().execute(listeners::dispatch); + } + + CompletableFuture stopFuture; + try { + stopFuture = task != null && task.state().ordinal() <= Service.State.RUNNING.ordinal() + ? Services.stopAsync(task, getRuntimeExecutor()) + : CompletableFuture.completedFuture(null); + } catch (Exception e) { + stopFuture = CompletableFutures.exceptionallyCompletedFuture(e); + } + + stopFuture.whenCompleteAsync((ignored, throwable) -> { + if (throwable != null) { + setPreviousFailure(throwable); + } + + getIOExecutor().execute(() -> { + closeTaskClassLoader(slot); + scheduleRunAsync( + () -> finishTaskCleanup(slot, task, cancellation, throwable), + 0, + TimeUnit.MILLISECONDS); + }); + }, getMainThreadExecutor()); + return slot.terminated; + } + + private void finishTaskCleanup( + TaskSlot slot, + @Nullable RuntimeTask task, + boolean cancellation, + @Nullable Throwable cleanupFailure) { + validateRunsInMainThread(); + if (cancellation && task != null) { + listeners.enqueue(getTaskCancelledEvent(task, cleanupFailure)); + getIOExecutor().execute(listeners::dispatch); + } + clearTaskSlot(slot); + if (cleanupFailure == null) { + slot.terminated.complete(null); } else { - setStatus(TaskExecutorReport.occupied(WorkerId.fromIdUnsafe(task.getWorkerId()))); + slot.terminated.completeExceptionally(cleanupFailure); + } + } + + private void clearTaskSlot(TaskSlot slot) { + if (taskSlot.compareAndSet(slot, null) && !stopping) { + setStatus(TaskExecutorReport.available()); + } + } + + private void closeTaskClassLoader(TaskSlot slot) { + UserCodeClassLoader userCodeClassLoader = slot.userCodeClassLoader; + if (userCodeClassLoader == null) { + return; + } + ClassLoader classLoader = userCodeClassLoader.asClassLoader(); + if (classLoader instanceof Closeable) { + try { + ((Closeable) classLoader).close(); + } catch (IOException e) { + log.warn("Failed to close classloader for {}", slot.workerId, e); + } } } @@ -496,49 +628,32 @@ private void setStatus(TaskExecutorReport newReport) { @Override public CompletableFuture cancelTask(WorkerId workerId) { log.info("TaskExecutor cancelTask requested for {}", workerId); - if (this.currentTask == null) { + TaskSlot slot = taskSlot.get(); + if (slot == null) { return CompletableFutures.exceptionallyCompletedFuture(new TaskNotFoundException(workerId)); - } else if (!this.currentTask.getWorkerId().equals(workerId.getId())) { - log.error("my current worker id is {} while expected worker id is {}", currentTask.getWorkerId(), workerId); + } else if (!slot.workerId.equals(workerId)) { + log.error("my current worker id is {} while expected worker id is {}", slot.workerId, workerId); return CompletableFutures.exceptionallyCompletedFuture(new TaskNotFoundException(workerId)); - } else { - scheduleRunAsync(this::stopCurrentTask, 0, TimeUnit.MILLISECONDS); - return CompletableFuture.completedFuture(Ack.getInstance()); } + + slot.cancelled.set(true); + scheduleRunAsync(() -> { + if (slot.prepared) { + disposeTaskSlot(slot, true); + } + }, 0, TimeUnit.MILLISECONDS); + return CompletableFuture.completedFuture(Ack.getInstance()); } private CompletableFuture stopCurrentTask() { log.info("TaskExecutor stopCurrentTask."); validateRunsInMainThread(); - if (this.currentTask != null) { - try { - if (this.currentTask.state().ordinal() <= Service.State.RUNNING.ordinal()) { - listeners.enqueue(getTaskCancellingEvent(currentTask)); - CompletableFuture stopTaskFuture = - Services.stopAsync(this.currentTask, getRuntimeExecutor()); - - return stopTaskFuture - .whenCompleteAsync((dontCare, throwable) -> { - RuntimeTask t = this.currentTask; - setCurrentTask(null); - if (throwable != null) { - setPreviousFailure(throwable); - } - listeners.enqueue(getTaskCancelledEvent(t, throwable)); - getIOExecutor().execute(listeners::dispatch); - }, getMainThreadExecutor()); - } else { - return CompletableFuture.completedFuture(null); - } - } catch (Exception e) { - log.error("stopping current task failed", e); - return CompletableFutures.exceptionallyCompletedFuture(e); - } finally { - getIOExecutor().execute(listeners::dispatch); - } - } else { + TaskSlot slot = taskSlot.get(); + if (slot == null) { return CompletableFuture.completedFuture(null); } + slot.cancelled.set(true); + return slot.prepared ? disposeTaskSlot(slot, true) : slot.terminated; } private CompletableFuture stopResourceManager() { @@ -578,7 +693,11 @@ protected CompletableFuture onStop() { validateRunsInMainThread(); log.info("TaskExecutor onStop."); - final CompletableFuture runningTaskCompletionFuture = stopCurrentTask(); + stopping = true; + final CompletableFuture runningTaskCompletionFuture = + stopCurrentTask() + .thenApply(ignored -> (Void) null) + .orTimeout(this.rpcCallTimeoutMsDp.getValue(), TimeUnit.MILLISECONDS); return runningTaskCompletionFuture .handleAsync((dontCare, throwable) -> { diff --git a/mantis-server/mantis-server-agent/src/test/java/io/mantisrx/server/agent/RuntimeTaskImplExecutorTest.java b/mantis-server/mantis-server-agent/src/test/java/io/mantisrx/server/agent/RuntimeTaskImplExecutorTest.java index d2281ff3d..772103ccb 100644 --- a/mantis-server/mantis-server-agent/src/test/java/io/mantisrx/server/agent/RuntimeTaskImplExecutorTest.java +++ b/mantis-server/mantis-server-agent/src/test/java/io/mantisrx/server/agent/RuntimeTaskImplExecutorTest.java @@ -18,9 +18,11 @@ import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; +import static org.mockito.Mockito.timeout; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -33,18 +35,22 @@ import io.mantisrx.common.Ack; import io.mantisrx.common.JsonSerializer; import io.mantisrx.common.WorkerPorts; +import io.mantisrx.common.properties.DefaultMantisPropertiesLoader; +import io.mantisrx.common.properties.MantisPropertiesLoader; import io.mantisrx.runtime.MachineDefinition; import io.mantisrx.runtime.MantisJobDurationType; import io.mantisrx.runtime.MantisJobState; import io.mantisrx.runtime.descriptor.SchedulingInfo; import io.mantisrx.runtime.loader.ClassLoaderHandle; import io.mantisrx.runtime.loader.RuntimeTask; +import io.mantisrx.runtime.loader.TaskFactory; import io.mantisrx.runtime.loader.config.WorkerConfiguration; import io.mantisrx.runtime.source.http.HttpServerProvider; import io.mantisrx.runtime.source.http.HttpSources; import io.mantisrx.runtime.source.http.impl.HttpClientFactories; import io.mantisrx.runtime.source.http.impl.HttpRequestFactories; import io.mantisrx.server.agent.TaskExecutor.Listener; +import io.mantisrx.server.agent.utils.DurableBooleanState; import io.mantisrx.server.core.ExecuteStageRequest; import io.mantisrx.server.core.JobSchedulingInfo; import io.mantisrx.server.core.PostJobStatusRequest; @@ -58,15 +64,21 @@ import io.mantisrx.server.master.client.ResourceLeaderConnection; import io.mantisrx.server.master.resourcecluster.RequestThrottledException; import io.mantisrx.server.master.resourcecluster.ResourceClusterGateway; +import io.mantisrx.server.master.resourcecluster.TaskExecutorRegistration; import io.mantisrx.server.master.resourcecluster.TaskExecutorReport; import io.mantisrx.server.master.resourcecluster.TaskExecutorStatusChange; +import io.mantisrx.server.worker.TaskExecutorGateway.TaskAlreadyRunningException; import io.mantisrx.server.worker.config.StaticPropertiesConfigurationFactory; import io.mantisrx.shaded.com.fasterxml.jackson.databind.ObjectMapper; import io.mantisrx.shaded.com.google.common.base.Preconditions; import io.mantisrx.shaded.com.google.common.collect.ImmutableMap; import io.mantisrx.shaded.com.google.common.collect.Lists; +import io.mantisrx.shaded.com.google.common.util.concurrent.AbstractIdleService; import io.mantisrx.shaded.com.google.common.util.concurrent.MoreExecutors; +import io.mantisrx.shaded.com.google.common.util.concurrent.Service.State; import java.io.BufferedReader; +import java.io.Closeable; +import java.io.File; import java.io.IOException; import java.io.InputStreamReader; import java.io.OutputStream; @@ -77,10 +89,12 @@ import java.util.Optional; import java.util.Properties; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import javax.annotation.Nullable; import lombok.Getter; @@ -88,6 +102,7 @@ import mantis.io.reactivex.netty.client.RxClient.ServerInfo; import org.apache.flink.api.common.time.Time; import org.apache.flink.runtime.rpc.RpcService; +import org.apache.flink.util.UserCodeClassLoader; import org.junit.After; import org.junit.Assert; import org.junit.Before; @@ -166,15 +181,34 @@ public void setUp() throws IOException { } private void start() throws Exception { - - listener = new CollectingTaskLifecycleListener(); - taskExecutor = + startTaskExecutor( new TestingTaskExecutor( rpcService, workerConfiguration, highAvailabilityServices, - classLoaderHandle - ); + classLoaderHandle)); + } + + private void start(TaskFactory taskFactory) throws Exception { + start(taskFactory, new DefaultMantisPropertiesLoader(System.getProperties())); + } + + private void start( + TaskFactory taskFactory, + MantisPropertiesLoader propertiesLoader) throws Exception { + startTaskExecutor( + new TaskExecutor( + rpcService, + workerConfiguration, + propertiesLoader, + highAvailabilityServices, + classLoaderHandle, + taskFactory)); + } + + private void startTaskExecutor(TaskExecutor executor) throws Exception { + listener = new CollectingTaskLifecycleListener(); + taskExecutor = executor; taskExecutor.addListener(listener, MoreExecutors.directExecutor()); taskExecutor.start(); taskExecutor.awaitRunning().get(2, TimeUnit.SECONDS); @@ -182,7 +216,9 @@ private void start() throws Exception { @After public void tearDown() throws Exception { - taskExecutor.close(); + if (taskExecutor != null) { + taskExecutor.close(); + } this.localApiServer.stop(0); } @@ -269,6 +305,367 @@ public Observable getServersToRemove() { assertFalse(listener.isFailedCalled()); } + @Test + public void acceptedTaskReportsOccupiedWhilePreparing() throws Exception { + WorkerId workerId = new WorkerId("jobId-0", 0, 1); + CountDownLatch preparationStarted = new CountDownLatch(1); + CountDownLatch releasePreparation = new CountDownLatch(1); + BlockingRuntimeTask task = + new BlockingRuntimeTask(workerId, preparationStarted, releasePreparation); + UserCodeClassLoader userCodeClassLoader = mock(UserCodeClassLoader.class); + when(userCodeClassLoader.asClassLoader()).thenReturn(getClass().getClassLoader()); + + start(singleTaskFactory(task, userCodeClassLoader)); + try { + taskExecutor.callInMainThread( + () -> taskExecutor.submitTask(createExecuteStageRequest(1)), + Time.seconds(1)).get(); + assertTrue(preparationStarted.await(2, TimeUnit.SECONDS)); + + Assert.assertEquals( + TaskExecutorReport.occupied(workerId), + taskExecutor.getCurrentReport().get(1, TimeUnit.SECONDS)); + } finally { + releasePreparation.countDown(); + } + } + + @Test + public void secondWorkerIsRejectedWhileAcceptedWorkerCanStillStart() throws Exception { + WorkerId firstWorker = new WorkerId("jobId-0", 0, 1); + CountDownLatch preparationStarted = new CountDownLatch(1); + CountDownLatch releasePreparation = new CountDownLatch(1); + BlockingRuntimeTask firstTask = + new BlockingRuntimeTask(firstWorker, preparationStarted, releasePreparation); + UserCodeClassLoader userCodeClassLoader = mock(UserCodeClassLoader.class); + when(userCodeClassLoader.asClassLoader()).thenReturn(getClass().getClassLoader()); + + start(singleTaskFactory(firstTask, userCodeClassLoader)); + try { + taskExecutor.callInMainThread( + () -> taskExecutor.submitTask(createExecuteStageRequest(1)), + Time.seconds(1)).get(); + assertTrue(preparationStarted.await(2, TimeUnit.SECONDS)); + + CompletableFuture secondAssignment = taskExecutor.callInMainThread( + () -> taskExecutor.submitTask(createExecuteStageRequest(2)), + Time.seconds(1)); + Throwable failure = secondAssignment.handle((ack, throwable) -> throwable).get(); + + assertTrue( + "second worker was accepted while the first accepted worker could still start", + failure != null); + assertTrue(rootCause(failure) instanceof TaskAlreadyRunningException); + } finally { + releasePreparation.countDown(); + } + } + + @Test + public void cancelDuringPreparationDisposesTaskAndClassLoader() throws Exception { + WorkerId workerId = new WorkerId("jobId-0", 0, 1); + CountDownLatch preparationStarted = new CountDownLatch(1); + CountDownLatch releasePreparation = new CountDownLatch(1); + BlockingRuntimeTask runtimeTask = + new BlockingRuntimeTask(workerId, preparationStarted, releasePreparation); + CloseTrackingClassLoader taskClassLoader = + new CloseTrackingClassLoader(getClass().getClassLoader()); + UserCodeClassLoader userCodeClassLoader = mock(UserCodeClassLoader.class); + when(userCodeClassLoader.asClassLoader()).thenReturn(taskClassLoader); + + start(singleTaskFactory(runtimeTask, userCodeClassLoader)); + taskExecutor.callInMainThread( + () -> taskExecutor.submitTask(createExecuteStageRequest(1)), Time.seconds(1)).get(); + assertTrue(preparationStarted.await(2, TimeUnit.SECONDS)); + + CompletableFuture cancelFuture = taskExecutor.callInMainThread( + () -> taskExecutor.cancelTask(workerId), Time.seconds(1)); + cancelFuture.get(1, TimeUnit.SECONDS); + Assert.assertEquals( + TaskExecutorReport.occupied(workerId), + taskExecutor.getCurrentReport().get(1, TimeUnit.SECONDS)); + + releasePreparation.countDown(); + + assertTrue(runtimeTask.preparationFinished.await(2, TimeUnit.SECONDS)); + verify(resourceManagerGateway, timeout(2000).times(1)).notifyTaskExecutorStatusChange( + new TaskExecutorStatusChange( + taskExecutor.getTaskExecutorID(), + taskExecutor.getClusterID(), + TaskExecutorReport.available())); + Assert.assertEquals(State.TERMINATED, runtimeTask.state()); + assertFalse(runtimeTask.started.get()); + assertTrue(taskClassLoader.closed.get()); + } + + @Test + public void stopFailureStillReleasesTaskSlot() throws Exception { + WorkerId workerId = new WorkerId("jobId-0", 0, 1); + CountDownLatch taskStarted = new CountDownLatch(1); + FailingStopRuntimeTask runtimeTask = + new FailingStopRuntimeTask(workerId, taskStarted); + CloseTrackingClassLoader taskClassLoader = + new CloseTrackingClassLoader(getClass().getClassLoader()); + UserCodeClassLoader userCodeClassLoader = mock(UserCodeClassLoader.class); + when(userCodeClassLoader.asClassLoader()).thenReturn(taskClassLoader); + + start(singleTaskFactory(runtimeTask, userCodeClassLoader)); + taskExecutor.callInMainThread( + () -> taskExecutor.submitTask(createExecuteStageRequest(1)), Time.seconds(1)).get(); + assertTrue(taskStarted.await(2, TimeUnit.SECONDS)); + + taskExecutor.callInMainThread( + () -> taskExecutor.cancelTask(workerId), Time.seconds(1)).get(1, TimeUnit.SECONDS); + + verify(resourceManagerGateway, timeout(2000).times(1)).notifyTaskExecutorStatusChange( + new TaskExecutorStatusChange( + taskExecutor.getTaskExecutorID(), + taskExecutor.getClusterID(), + TaskExecutorReport.available())); + Assert.assertEquals( + TaskExecutorReport.available(), + taskExecutor.getCurrentReport().get(1, TimeUnit.SECONDS)); + Assert.assertEquals(State.FAILED, runtimeTask.state()); + assertTrue(taskClassLoader.closed.get()); + } + + @Test + public void shutdownDoesNotWaitIndefinitelyForPreparation() throws Exception { + WorkerId workerId = new WorkerId("jobId-0", 0, 1); + CountDownLatch preparationStarted = new CountDownLatch(1); + CountDownLatch releasePreparation = new CountDownLatch(1); + BlockingRuntimeTask runtimeTask = + new BlockingRuntimeTask(workerId, preparationStarted, releasePreparation); + UserCodeClassLoader userCodeClassLoader = mock(UserCodeClassLoader.class); + when(userCodeClassLoader.asClassLoader()).thenReturn(getClass().getClassLoader()); + MantisPropertiesLoader shortTimeout = (name, defaultValue) -> + "mantis.taskexecutor.heartbeats.timeout.ms".equals(name) ? "100" : defaultValue; + + start(singleTaskFactory(runtimeTask, userCodeClassLoader), shortTimeout); + taskExecutor.callInMainThread( + () -> taskExecutor.submitTask(createExecuteStageRequest(1)), Time.seconds(1)).get(); + assertTrue(preparationStarted.await(2, TimeUnit.SECONDS)); + + TaskExecutor executor = taskExecutor; + CompletableFuture closeFuture = CompletableFuture.runAsync(() -> { + try { + executor.close(); + } catch (Exception e) { + throw new CompletionException(e); + } + }); + try { + closeFuture.get(2, TimeUnit.SECONDS); + taskExecutor = null; + assertFalse(runtimeTask.started.get()); + } finally { + releasePreparation.countDown(); + } + assertTrue(runtimeTask.preparationFinished.await(2, TimeUnit.SECONDS)); + } + + @Test + public void shutdownDuringPreparationDisposesTaskAndClassLoader() throws Exception { + WorkerId workerId = new WorkerId("jobId-0", 0, 1); + CountDownLatch preparationStarted = new CountDownLatch(1); + CountDownLatch releasePreparation = new CountDownLatch(1); + BlockingRuntimeTask runtimeTask = + new BlockingRuntimeTask(workerId, preparationStarted, releasePreparation); + CloseTrackingClassLoader taskClassLoader = + new CloseTrackingClassLoader(getClass().getClassLoader()); + UserCodeClassLoader userCodeClassLoader = mock(UserCodeClassLoader.class); + when(userCodeClassLoader.asClassLoader()).thenReturn(taskClassLoader); + + start(singleTaskFactory(runtimeTask, userCodeClassLoader)); + taskExecutor.callInMainThread( + () -> taskExecutor.submitTask(createExecuteStageRequest(1)), Time.seconds(1)).get(); + assertTrue(preparationStarted.await(2, TimeUnit.SECONDS)); + + TaskExecutor executor = taskExecutor; + CompletableFuture closeFuture = CompletableFuture.runAsync(() -> { + try { + executor.close(); + } catch (Exception e) { + throw new CompletionException(e); + } + }); + releasePreparation.countDown(); + closeFuture.get(2, TimeUnit.SECONDS); + taskExecutor = null; + + Assert.assertEquals(State.TERMINATED, runtimeTask.state()); + assertTrue(taskClassLoader.closed.get()); + } + + @Test + public void legacyRegistrationMarkerIsResetForCapabilityRefresh() throws Exception { + File legacyStateFile = new File( + workerConfiguration.getRegistrationStoreDir(), "rmCxnState.txt"); + DurableBooleanState legacyState = new DurableBooleanState(legacyStateFile.getAbsolutePath()); + legacyState.setState(true); + + start(); + + verify(resourceManagerGateway, timeout(2000).times(1)).registerTaskExecutor( + argThat(TaskExecutorRegistration::reservesAcceptedTask)); + assertFalse(new DurableBooleanState(legacyStateFile.getAbsolutePath()).getState()); + } + + private TaskFactory singleTaskFactory(RuntimeTask runtimeTask, UserCodeClassLoader userCodeClassLoader) { + return new TaskFactory() { + @Override + public RuntimeTask getRuntimeTaskInstance(ExecuteStageRequest request, ClassLoader ignored) { + return runtimeTask; + } + + @Override + public UserCodeClassLoader getUserCodeClassLoader( + ExecuteStageRequest request, + ClassLoaderHandle ignored) { + return userCodeClassLoader; + } + }; + } + + private ExecuteStageRequest createExecuteStageRequest(int workerNumber) throws Exception { + return new ExecuteStageRequest( + "jobName", + "jobId-0", + 0, + workerNumber, + new URL("https://example.invalid/job.zip"), + 1, + 1, + ImmutableList.of(100), + 100L, + 1, + ImmutableList.of(), + new SchedulingInfo.Builder() + .numberOfStages(1) + .singleWorkerStageWithConstraints( + new MachineDefinition(1, 10, 10, 10, 2), + Lists.newArrayList(), + Lists.newArrayList()) + .build(), + MantisJobDurationType.Transient, + 0, + 1000L, + 1L, + new WorkerPorts(2, 3, 4, 5, 6), + Optional.of(SineFunctionJobProvider.class.getName()), + "user", + "111"); + } + + private static Throwable rootCause(Throwable throwable) { + Throwable current = throwable; + while (current.getCause() != null) { + current = current.getCause(); + } + return current; + } + + private static class BlockingRuntimeTask extends AbstractIdleService implements RuntimeTask { + private final WorkerId workerId; + private final CountDownLatch preparationStarted; + private final CountDownLatch releasePreparation; + private final CountDownLatch preparationFinished = new CountDownLatch(1); + private final AtomicBoolean started = new AtomicBoolean(); + + private BlockingRuntimeTask( + WorkerId workerId, + CountDownLatch preparationStarted, + CountDownLatch releasePreparation) { + this.workerId = workerId; + this.preparationStarted = preparationStarted; + this.releasePreparation = releasePreparation; + } + + @Override + public void initialize( + String executeStageRequestString, + String workerConfigurationString, + UserCodeClassLoader userCodeClassLoader) { + preparationStarted.countDown(); + try { + await(releasePreparation); + } finally { + preparationFinished.countDown(); + } + } + + @Override + public String getWorkerId() { + return workerId.getId(); + } + + @Override + protected void startUp() { + started.set(true); + } + + @Override + protected void shutDown() { + } + } + + private static final class CloseTrackingClassLoader extends ClassLoader implements Closeable { + private final AtomicBoolean closed = new AtomicBoolean(); + + private CloseTrackingClassLoader(ClassLoader parent) { + super(parent); + } + + @Override + public void close() { + closed.set(true); + } + } + + private static final class FailingStopRuntimeTask extends AbstractIdleService + implements RuntimeTask { + private final WorkerId workerId; + private final CountDownLatch taskStarted; + + private FailingStopRuntimeTask(WorkerId workerId, CountDownLatch taskStarted) { + this.workerId = workerId; + this.taskStarted = taskStarted; + } + + @Override + public void initialize( + String executeStageRequestString, + String workerConfigurationString, + UserCodeClassLoader userCodeClassLoader) { + } + + @Override + public String getWorkerId() { + return workerId.getId(); + } + + @Override + protected void startUp() { + taskStarted.countDown(); + } + + @Override + protected void shutDown() { + throw new RuntimeException("stop failed"); + } + } + + private static void await(CountDownLatch latch) { + try { + if (!latch.await(5, TimeUnit.SECONDS)) { + throw new IllegalStateException("timed out waiting for test latch"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } + private void setupLocalControl(int port) throws IOException { this.localApiServer = HttpServer.create(new InetSocketAddress("localhost", port), 0); this.localApiServer.createContext(