This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/dev/pr-12948-4ae9604aaf876c64890cce88cc202a755bc36a64 in repository https://gitbox.apache.org/repos/asf/seatunnel.git
commit 7f87688505cfc915de0ed4e87a4eda6fcee9f2cf Author: Daniel <[email protected]> AuthorDate: Mon Oct 26 15:06:29 2026 +0000 [Fix][Zeta] Harden slot allocation recovery state (#12948) Co-authored-by: zhiweiniu <[email protected]> --- .../engine/server/CoordinatorService.java | 11 +- .../engine/server/dag/physical/PhysicalVertex.java | 9 +- .../engine/server/dag/physical/ResourceUtils.java | 8 +- .../engine/server/dag/physical/SubPlan.java | 15 +- .../server/resourcemanager/ResourceManager.java | 7 +- .../resourcemanager/ResourceRequestHandler.java | 135 ++++-- .../server/CoordinatorServiceStatusTest.java | 50 +++ .../physical/PhysicalVertexDeployStateTest.java | 52 +++ .../server/dag/physical/ResourceUtilsTest.java | 123 ++++++ .../server/dag/physical/SubPlanDeploymentTest.java | 136 ++++++ .../ResourceRequestHandlerTest.java | 485 +++++++++++++++++++++ 11 files changed, 981 insertions(+), 50 deletions(-) diff --git a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/CoordinatorService.java b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/CoordinatorService.java index aea9e45bfa..b4919814a2 100644 --- a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/CoordinatorService.java +++ b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/CoordinatorService.java @@ -1637,16 +1637,21 @@ public class CoordinatorService { } JobMaster runningJobMaster = runningJobMasterMap.get(jobId); if (runningJobMaster == null) { - JobHistoryService.JobState jobDetailState = jobHistoryService.getJobDetailState(jobId); - return null == jobDetailState ? JobStatus.UNKNOWABLE : jobDetailState.getJobStatus(); + return resolveHistoryStatus(jobHistoryService.getJobDetailState(jobId)); } JobStatus jobStatus = runningJobMaster.getJobStatus(); if (jobStatus == null) { - return jobHistoryService.getFinishedJobStateImap().get(jobId).getJobStatus(); + return resolveHistoryStatus(jobHistoryService.getFinishedJobStateImap().get(jobId)); } return jobStatus; } + static JobStatus resolveHistoryStatus(JobHistoryService.JobState jobState) { + return jobState == null || jobState.getJobStatus() == null + ? JobStatus.UNKNOWABLE + : jobState.getJobStatus(); + } + public boolean shouldShowAsRunningJob(long jobId) { JobStatus status = getJobStatus(jobId); return status != null && !status.isEndState(); diff --git a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/dag/physical/PhysicalVertex.java b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/dag/physical/PhysicalVertex.java index c2a73ef9c1..abe21a8ec8 100644 --- a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/dag/physical/PhysicalVertex.java +++ b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/dag/physical/PhysicalVertex.java @@ -322,7 +322,14 @@ public class PhysicalVertex { TaskGroupImmutableInformation taskGroupImmutableInformation = getTaskGroupImmutableInformation(); TaskDeployState state = taskGroupConsumer.apply(taskGroupImmutableInformation); - updateTaskState(ExecutionState.RUNNING); + return finishDeployment(state); + } + + @VisibleForTesting + TaskDeployState finishDeployment(TaskDeployState state) { + if (state.isSuccess()) { + updateTaskState(ExecutionState.RUNNING); + } return state; } diff --git a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/dag/physical/ResourceUtils.java b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/dag/physical/ResourceUtils.java index 3e2e1eac17..1d6ac8a473 100644 --- a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/dag/physical/ResourceUtils.java +++ b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/dag/physical/ResourceUtils.java @@ -32,6 +32,7 @@ import lombok.NonNull; import java.util.HashMap; import java.util.Map; +import java.util.concurrent.CancellationException; import java.util.concurrent.CompletionException; public class ResourceUtils { @@ -52,8 +53,11 @@ public class ResourceUtils { futures.forEach( (key, value) -> { try { - slotProfiles.put(key, value == null ? null : value.join()); - } catch (CompletionException e) { + SlotProfile slotProfile = value == null ? null : value.join(); + if (slotProfile != null) { + slotProfiles.put(key, slotProfile); + } + } catch (CompletionException | CancellationException e) { LOGGER.warning("Failed to join future for task group location: " + key, e); } }); diff --git a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/dag/physical/SubPlan.java b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/dag/physical/SubPlan.java index 3347ea4c33..be94e2de0e 100644 --- a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/dag/physical/SubPlan.java +++ b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/dag/physical/SubPlan.java @@ -787,7 +787,9 @@ public class SubPlan { task.makeTaskGroupDeploy(); } }); - updatePipelineState(PipelineStatus.RUNNING); + if (areAllVerticesRunningOrFinished(coordinatorVertexList, physicalVertexList)) { + updatePipelineState(PipelineStatus.RUNNING); + } break; case RUNNING: break; @@ -834,4 +836,15 @@ public class SubPlan { errorByPhysicalVertex.compareAndSet(null, ExceptionUtils.getMessage(e)); updatePipelineState(PipelineStatus.FAILING); } + + static boolean areAllVerticesRunningOrFinished( + List<PhysicalVertex> coordinators, List<PhysicalVertex> tasks) { + return coordinators.stream().allMatch(SubPlan::isRunningOrFinished) + && tasks.stream().allMatch(SubPlan::isRunningOrFinished); + } + + private static boolean isRunningOrFinished(PhysicalVertex vertex) { + return ExecutionState.RUNNING.equals(vertex.getExecutionState()) + || ExecutionState.FINISHED.equals(vertex.getExecutionState()); + } } diff --git a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/resourcemanager/ResourceManager.java b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/resourcemanager/ResourceManager.java index 3cfc165518..d2103d928f 100644 --- a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/resourcemanager/ResourceManager.java +++ b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/resourcemanager/ResourceManager.java @@ -73,9 +73,10 @@ public interface ResourceManager { * Requests the slots needed by a job's task groups from matching registered workers. * * <p>Completion requires all requested assignments. The existing request handler attempts to - * release partial allocations when allocation fails; failure completion is not a guarantee that - * every rollback RPC has already completed. Successful callers own the returned allocations and - * must release them when their task groups no longer need them. + * release partial allocations when allocation fails. Failure completion waits for those + * rollback attempts; an exhausted rollback failure is attached to the allocation failure. + * Successful callers own the returned allocations and must release them when their task groups + * no longer need them. * * @param jobId job that will own the granted slots * @param resourceProfile nonempty list of requested slot resource requirements diff --git a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/resourcemanager/ResourceRequestHandler.java b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/resourcemanager/ResourceRequestHandler.java index 550922f58b..4f8ca55e89 100644 --- a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/resourcemanager/ResourceRequestHandler.java +++ b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/resourcemanager/ResourceRequestHandler.java @@ -19,6 +19,7 @@ package org.apache.seatunnel.engine.server.resourcemanager; import org.apache.seatunnel.shade.com.google.common.annotations.VisibleForTesting; +import org.apache.seatunnel.engine.common.Constant; import org.apache.seatunnel.engine.common.config.server.AllocateStrategy; import org.apache.seatunnel.engine.common.runtime.DeployType; import org.apache.seatunnel.engine.common.utils.concurrent.CompletableFuture; @@ -45,6 +46,7 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.stream.Collectors; import static com.hazelcast.jet.impl.util.ExceptionUtil.withTryCatch; @@ -74,6 +76,8 @@ public class ResourceRequestHandler { private final SlotAllocationStrategy slotAllocationStrategy; + private final AtomicBoolean completionStarted = new AtomicBoolean(false); + public ResourceRequestHandler( long jobId, List<ResourceProfile> resourceProfile, @@ -92,12 +96,16 @@ public class ResourceRequestHandler { } public CompletableFuture<List<SlotProfile>> request(Map<String, String> tags) { - requestSlotWithRetry(resourceProfile, MAX_RETRY_TIMES, tags); + List<Integer> requestIndexes = new ArrayList<>(); + for (int i = 0; i < resourceProfile.size(); i++) { + requestIndexes.add(i); + } + requestSlotWithRetry(requestIndexes, MAX_RETRY_TIMES, tags); return completableFuture; } private CompletableFuture<SlotAndWorkerProfile> requestSlotWithRetry( - List<ResourceProfile> request, int retryTimes, Map<String, String> tags) { + List<Integer> requestIndexes, int retryTimes, Map<String, String> tags) { if (retryTimes <= 0) { LOGGER.fine("can't apply resource request with retry times: " + MAX_RETRY_TIMES); return CompletableFuture.supplyAsync( @@ -107,7 +115,8 @@ public class ResourceRequestHandler { + MAX_RETRY_TIMES); }); } - List<CompletableFuture<SlotAndWorkerProfile>> allRequestFuture = requestSlots(request); + List<CompletableFuture<SlotAndWorkerProfile>> allRequestFuture = + requestSlots(requestIndexes); // all resource preCheck done, also had sent request to worker return getAllOfFuture(allRequestFuture) .whenComplete( @@ -116,14 +125,15 @@ public class ResourceRequestHandler { (unused, error) -> { if (error != null) { completeRequestWithException(toExternalException(error)); + return; } else { - List<ResourceProfile> needRequestResource = - stillNeedRequestResource(); - if (!needRequestResource.isEmpty()) { + List<Integer> needRequestIndexes = + stillNeedRequestIndexes(); + if (!needRequestIndexes.isEmpty()) { Exception requestSlotWithRetryError = null; try { requestSlotWithRetry( - needRequestResource, + needRequestIndexes, retryTimes - 1, tags) .get(); @@ -162,14 +172,14 @@ public class ResourceRequestHandler { })); } - private List<ResourceProfile> stillNeedRequestResource() { - List<ResourceProfile> needRequestResource = new ArrayList<>(); + private List<Integer> stillNeedRequestIndexes() { + List<Integer> needRequestIndexes = new ArrayList<>(); for (int i = 0; i < resourceProfile.size(); i++) { if (!resultSlotProfiles.containsKey(i)) { - needRequestResource.add(resourceProfile.get(i)); + needRequestIndexes.add(i); } } - return needRequestResource; + return needRequestIndexes; } private int findNullIndexInResultSlotProfiles() { @@ -182,42 +192,46 @@ public class ResourceRequestHandler { } private List<CompletableFuture<SlotAndWorkerProfile>> requestSlots( - List<ResourceProfile> requestProfile) { + List<Integer> requestIndexes) { List<CompletableFuture<SlotAndWorkerProfile>> allRequestFuture = new ArrayList<>(); - for (int i = 0; i < requestProfile.size(); i++) { - ResourceProfile r = requestProfile.get(i); + for (int index : requestIndexes) { + ResourceProfile r = resourceProfile.get(index); Optional<WorkerProfile> workerProfile = preCheckWorkerResource(r); if (workerProfile.isPresent()) { // request slot to member CompletableFuture<SlotAndWorkerProfile> internalCompletableFuture = - singleResourceRequestToMember(i, r, workerProfile.get()); + singleResourceRequestToMember(index, r, workerProfile.get()); allRequestFuture.add(internalCompletableFuture); } else { - // if no worker can provide the resource, we should return a failed future + // Leave the original index unassigned so retry or dynamic-worker provisioning can + // handle it after all in-flight slot requests finish. LOGGER.fine("pre check worker resource failed, can't apply resource request: " + r); - allRequestFuture.add( - CompletableFuture.supplyAsync( - () -> { - throw new NoEnoughResourceException( - "can't apply resource request: " + r); - })); } } return allRequestFuture; } private void completeRequestWithException(Throwable e) { - if (completableFuture.completeExceptionally(e)) { - releaseAllResourceInternal(); + if (!completionStarted.compareAndSet(false, true)) { + return; } + releaseAllResourceInternal() + .whenComplete( + (ignored, releaseError) -> { + if (releaseError != null) { + e.addSuppressed(releaseError); + } + completableFuture.completeExceptionally(e); + }); } private void addSlotToCacheMap(int index, SlotProfile slotProfile) { // null value means the slot request failed, no suitable slot found if (null != slotProfile) { resultSlotProfiles.put(index, slotProfile); - if (resultSlotProfiles.size() == resourceProfile.size()) { + if (resultSlotProfiles.size() == resourceProfile.size() + && completionStarted.compareAndSet(false, true)) { List<SlotProfile> value = new ArrayList<>(); for (int i = 0; i < resultSlotProfiles.size(); i++) { value.add(resultSlotProfiles.get(i)); @@ -303,7 +317,7 @@ public class ResourceRequestHandler { private void applyByDynamicWorker(Map<String, String> tags) { List<ResourceProfile> needApplyResource = new ArrayList<>(); List<Integer> needApplyIndex = new ArrayList<>(); - for (int i = 0; i < resultSlotProfiles.size(); i++) { + for (int i = 0; i < resourceProfile.size(); i++) { if (!resultSlotProfiles.containsKey(i)) { needApplyResource.add(resourceProfile.get(i)); needApplyIndex.add(i); @@ -326,7 +340,7 @@ public class ResourceRequestHandler { })); } - private void releaseAllResourceInternal() { + private CompletableFuture<Void> releaseAllResourceInternal() { int appliedCount = resultSlotProfiles.size(); int requiredCount = resourceProfile.size(); @@ -345,19 +359,60 @@ public class ResourceRequestHandler { + "releasing slots: [%s], remaining: %d slots not assigned", jobId, requiredCount, appliedCount, slotIds, requiredCount - appliedCount)); - new ArrayList<>(resultSlotProfiles.keySet()) - .forEach( - index -> { - SlotProfile profile = resultSlotProfiles.remove(index); - if (profile == null) { - return; - } - LOGGER.fine( - String.format( - "Releasing slot %d for job %d from worker %s", - profile.getSlotID(), jobId, profile.getWorker())); - resourceManager.releaseResource(jobId, profile); - }); + List<SlotProfile> acquiredSlots = + resultSlotProfiles.values().stream() + .filter(Objects::nonNull) + .collect(Collectors.toList()); + acquiredSlots.forEach( + profile -> + LOGGER.fine( + String.format( + "Releasing slot %d for job %d from worker %s", + profile.getSlotID(), jobId, profile.getWorker()))); + List<CompletableFuture<Void>> releaseFutures = + acquiredSlots.stream() + .map( + profile -> + releaseResourceWithRetry( + profile, Constant.OPERATION_RETRY_TIME)) + .collect(Collectors.toList()); + return getAllOfFuture(releaseFutures); + } + + private CompletableFuture<Void> releaseResourceWithRetry( + SlotProfile slotProfile, int remainingAttempts) { + CompletableFuture<Void> result = new CompletableFuture<>(); + CompletableFuture<Void> releaseFuture; + try { + releaseFuture = resourceManager.releaseResource(jobId, slotProfile); + } catch (Exception e) { + releaseFuture = new CompletableFuture<>(); + releaseFuture.completeExceptionally(e); + } + releaseFuture.whenComplete( + (ignored, error) -> { + if (error == null) { + result.complete(null); + } else if (remainingAttempts > 1) { + LOGGER.warning( + String.format( + "Release slot %d for job %d failed, retrying (%d attempts left)", + slotProfile.getSlotID(), jobId, remainingAttempts - 1), + error); + releaseResourceWithRetry(slotProfile, remainingAttempts - 1) + .whenComplete( + (retryIgnored, retryError) -> { + if (retryError == null) { + result.complete(null); + } else { + result.completeExceptionally(retryError); + } + }); + } else { + result.completeExceptionally(error); + } + }); + return result; } private NoEnoughResourceException toExternalException(Throwable error) { diff --git a/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/CoordinatorServiceStatusTest.java b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/CoordinatorServiceStatusTest.java new file mode 100644 index 0000000000..6de93914ea --- /dev/null +++ b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/CoordinatorServiceStatusTest.java @@ -0,0 +1,50 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.seatunnel.engine.server; + +import org.apache.seatunnel.engine.common.job.JobStatus; +import org.apache.seatunnel.engine.server.master.JobHistoryService; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class CoordinatorServiceStatusTest { + + @Test + void missingHistoryRecordIsUnknowable() { + assertEquals(JobStatus.UNKNOWABLE, CoordinatorService.resolveHistoryStatus(null)); + } + + @Test + void missingStatusInsideHistoryRecordIsUnknowable() { + JobHistoryService.JobState jobState = + new JobHistoryService.JobState(1L, "job", null, 1L, null, null, null, null); + + assertEquals(JobStatus.UNKNOWABLE, CoordinatorService.resolveHistoryStatus(jobState)); + } + + @Test + void knownHistoryStatusIsPreserved() { + JobHistoryService.JobState jobState = + new JobHistoryService.JobState( + 1L, "job", JobStatus.FINISHED, 1L, null, null, null, null); + + assertEquals(JobStatus.FINISHED, CoordinatorService.resolveHistoryStatus(jobState)); + } +} diff --git a/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/dag/physical/PhysicalVertexDeployStateTest.java b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/dag/physical/PhysicalVertexDeployStateTest.java new file mode 100644 index 0000000000..6aedf95af3 --- /dev/null +++ b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/dag/physical/PhysicalVertexDeployStateTest.java @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.seatunnel.engine.server.dag.physical; + +import org.apache.seatunnel.engine.server.execution.ExecutionState; +import org.apache.seatunnel.engine.server.execution.TaskDeployState; + +import org.junit.jupiter.api.Test; + +import static org.mockito.Mockito.CALLS_REAL_METHODS; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; + +class PhysicalVertexDeployStateTest { + + @Test + void failedDeploymentMustNotMarkTaskRunning() { + PhysicalVertex vertex = mock(PhysicalVertex.class, CALLS_REAL_METHODS); + doNothing().when(vertex).updateTaskState(ExecutionState.RUNNING); + + vertex.finishDeployment(TaskDeployState.failed(new RuntimeException("deploy failed"))); + + verify(vertex, never()).updateTaskState(ExecutionState.RUNNING); + } + + @Test + void successfulDeploymentMarksTaskRunning() { + PhysicalVertex vertex = mock(PhysicalVertex.class, CALLS_REAL_METHODS); + doNothing().when(vertex).updateTaskState(ExecutionState.RUNNING); + + vertex.finishDeployment(TaskDeployState.success()); + + verify(vertex).updateTaskState(ExecutionState.RUNNING); + } +} diff --git a/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/dag/physical/ResourceUtilsTest.java b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/dag/physical/ResourceUtilsTest.java new file mode 100644 index 0000000000..348e465e53 --- /dev/null +++ b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/dag/physical/ResourceUtilsTest.java @@ -0,0 +1,123 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.seatunnel.engine.server.dag.physical; + +import org.apache.seatunnel.engine.common.utils.concurrent.CompletableFuture; +import org.apache.seatunnel.engine.server.execution.TaskGroupLocation; +import org.apache.seatunnel.engine.server.master.JobMaster; +import org.apache.seatunnel.engine.server.resourcemanager.NoEnoughResourceException; +import org.apache.seatunnel.engine.server.resourcemanager.resource.SlotProfile; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class ResourceUtilsTest { + + @Test + void nullSlotIsRejectedEvenWhenEveryTaskHasAnEntry() { + JobMaster jobMaster = mock(JobMaster.class); + PhysicalPlan physicalPlan = mock(PhysicalPlan.class); + SubPlan subPlan = mock(SubPlan.class); + PhysicalVertex vertex = mock(PhysicalVertex.class); + TaskGroupLocation taskLocation = new TaskGroupLocation(1L, 2, 3L); + Map<TaskGroupLocation, CompletableFuture<SlotProfile>> preApplied = new HashMap<>(); + preApplied.put(taskLocation, null); + + configure(jobMaster, physicalPlan, subPlan, Collections.singletonList(vertex), preApplied); + when(vertex.getTaskGroupLocation()).thenReturn(taskLocation); + + assertThrows( + NoEnoughResourceException.class, + () -> ResourceUtils.applyResourceForPipeline(jobMaster, subPlan)); + } + + @Test + void exceptionalFutureRetainsSuccessfulSlotsForCleanup() { + assertFailedFutureRetainsSuccessfulSlots(false); + } + + @Test + void canceledFutureRetainsSuccessfulSlotsForCleanup() { + assertFailedFutureRetainsSuccessfulSlots(true); + } + + private static void assertFailedFutureRetainsSuccessfulSlots(boolean cancelFuture) { + JobMaster jobMaster = mock(JobMaster.class); + PhysicalPlan physicalPlan = mock(PhysicalPlan.class); + SubPlan subPlan = mock(SubPlan.class); + PhysicalVertex successfulVertex = mock(PhysicalVertex.class); + PhysicalVertex failedVertex = mock(PhysicalVertex.class); + TaskGroupLocation successfulLocation = new TaskGroupLocation(1L, 2, 3L); + TaskGroupLocation failedLocation = new TaskGroupLocation(1L, 2, 4L); + PipelineLocation pipelineLocation = new PipelineLocation(1L, 2); + SlotProfile successfulSlot = mock(SlotProfile.class); + CompletableFuture<SlotProfile> failedFuture = new CompletableFuture<>(); + if (cancelFuture) { + failedFuture.cancel(false); + } else { + failedFuture.completeExceptionally(new RuntimeException("slot request failed")); + } + Map<TaskGroupLocation, CompletableFuture<SlotProfile>> preApplied = new HashMap<>(); + preApplied.put(successfulLocation, CompletableFuture.completedFuture(successfulSlot)); + preApplied.put(failedLocation, failedFuture); + + configure( + jobMaster, + physicalPlan, + subPlan, + Arrays.asList(successfulVertex, failedVertex), + preApplied); + when(subPlan.getPipelineLocation()).thenReturn(pipelineLocation); + when(successfulVertex.getTaskGroupLocation()).thenReturn(successfulLocation); + when(failedVertex.getTaskGroupLocation()).thenReturn(failedLocation); + + assertThrows( + NoEnoughResourceException.class, + () -> ResourceUtils.applyResourceForPipeline(jobMaster, subPlan)); + + verify(jobMaster) + .setOwnedSlotProfiles( + eq(pipelineLocation), + eq(Collections.singletonMap(successfulLocation, successfulSlot))); + } + + private static void configure( + JobMaster jobMaster, + PhysicalPlan physicalPlan, + SubPlan subPlan, + List<PhysicalVertex> vertices, + Map<TaskGroupLocation, CompletableFuture<SlotProfile>> preApplied) { + when(jobMaster.getPhysicalPlan()).thenReturn(physicalPlan); + when(physicalPlan.getPreApplyResourceFutures()).thenReturn(preApplied); + when(subPlan.getCoordinatorVertexList()).thenReturn(Collections.emptyList()); + when(subPlan.getPhysicalVertexList()).thenReturn(vertices); + when(subPlan.getJobMaster()).thenReturn(jobMaster); + when(subPlan.getPipelineLocation()).thenReturn(new PipelineLocation(1L, 2)); + } +} diff --git a/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/dag/physical/SubPlanDeploymentTest.java b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/dag/physical/SubPlanDeploymentTest.java new file mode 100644 index 0000000000..6a29a763ba --- /dev/null +++ b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/dag/physical/SubPlanDeploymentTest.java @@ -0,0 +1,136 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.seatunnel.engine.server.dag.physical; + +import org.apache.seatunnel.engine.common.config.JobConfig; +import org.apache.seatunnel.engine.core.job.JobImmutableInformation; +import org.apache.seatunnel.engine.core.job.PipelineStatus; +import org.apache.seatunnel.engine.server.execution.ExecutionState; + +import org.junit.jupiter.api.Test; + +import com.hazelcast.map.IMap; + +import java.util.Collections; +import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +class SubPlanDeploymentTest { + + @Test + void pipelineIsNotRunningWhenAnyTaskDeploymentFailed() { + assertFalse( + SubPlan.areAllVerticesRunningOrFinished( + Collections.singletonList(vertexIn(ExecutionState.RUNNING)), + Collections.singletonList(vertexIn(ExecutionState.FAILED)))); + } + + @Test + void pipelineIsNotRunningWhenCoordinatorDeploymentFailed() { + assertFalse( + SubPlan.areAllVerticesRunningOrFinished( + Collections.singletonList(vertexIn(ExecutionState.FAILED)), + Collections.singletonList(vertexIn(ExecutionState.RUNNING)))); + } + + @Test + void pipelineIsRunningOnlyAfterEveryVertexIsRunning() { + assertTrue( + SubPlan.areAllVerticesRunningOrFinished( + Collections.singletonList(vertexIn(ExecutionState.RUNNING)), + Collections.singletonList(vertexIn(ExecutionState.RUNNING)))); + } + + @Test + void finishedVertexDuringDeploymentStillPromotesPipelineToRunning() { + AtomicReference<PipelineStatus> state = new AtomicReference<>(PipelineStatus.DEPLOYING); + SubPlan subPlan = + subPlanIn( + state, + Collections.singletonList(vertexIn(ExecutionState.RUNNING)), + Collections.singletonList(vertexIn(ExecutionState.FINISHED))); + + subPlan.startSubPlanStateProcess(); + + assertTrue(PipelineStatus.RUNNING.equals(state.get())); + assertTrue(PipelineStatus.RUNNING.equals(subPlan.getPipelineState())); + } + + @Test + void failedVertexDuringDeploymentDoesNotPromotePipelineToRunning() { + AtomicReference<PipelineStatus> state = new AtomicReference<>(PipelineStatus.DEPLOYING); + SubPlan subPlan = + subPlanIn( + state, + Collections.singletonList(vertexIn(ExecutionState.FAILED)), + Collections.singletonList(vertexIn(ExecutionState.RUNNING))); + + subPlan.startSubPlanStateProcess(); + + assertTrue(PipelineStatus.DEPLOYING.equals(state.get())); + assertTrue(PipelineStatus.DEPLOYING.equals(subPlan.getPipelineState())); + } + + @SuppressWarnings("unchecked") + private static SubPlan subPlanIn( + AtomicReference<PipelineStatus> state, + List<PhysicalVertex> tasks, + List<PhysicalVertex> coordinators) { + JobConfig jobConfig = new JobConfig(); + jobConfig.setName("deployment-state"); + JobImmutableInformation jobInformation = mock(JobImmutableInformation.class); + when(jobInformation.getJobId()).thenReturn(1L); + when(jobInformation.getJobConfig()).thenReturn(jobConfig); + IMap<Object, Object> states = mock(IMap.class); + when(states.get(any())).thenAnswer(ignored -> state.get()); + doAnswer( + invocation -> { + state.set(invocation.getArgument(1)); + return null; + }) + .when(states) + .set(any(), any()); + IMap<Object, Long[]> timestamps = mock(IMap.class); + when(timestamps.get(any())).thenReturn(new Long[PipelineStatus.values().length]); + return new SubPlan( + 1, + 1, + 0L, + tasks, + coordinators, + jobInformation, + mock(ExecutorService.class), + states, + timestamps, + Collections.emptyMap()); + } + + private static PhysicalVertex vertexIn(ExecutionState state) { + PhysicalVertex vertex = mock(PhysicalVertex.class); + when(vertex.getExecutionState()).thenReturn(state); + return vertex; + } +} diff --git a/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/resourcemanager/ResourceRequestHandlerTest.java b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/resourcemanager/ResourceRequestHandlerTest.java new file mode 100644 index 0000000000..3f232a0270 --- /dev/null +++ b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/resourcemanager/ResourceRequestHandlerTest.java @@ -0,0 +1,485 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.seatunnel.engine.server.resourcemanager; + +import org.apache.seatunnel.engine.common.Constant; +import org.apache.seatunnel.engine.common.config.EngineConfig; +import org.apache.seatunnel.engine.common.utils.concurrent.CompletableFuture; +import org.apache.seatunnel.engine.server.resourcemanager.allocation.strategy.RandomStrategy; +import org.apache.seatunnel.engine.server.resourcemanager.opeartion.RequestSlotOperation; +import org.apache.seatunnel.engine.server.resourcemanager.resource.CPU; +import org.apache.seatunnel.engine.server.resourcemanager.resource.Memory; +import org.apache.seatunnel.engine.server.resourcemanager.resource.ResourceProfile; +import org.apache.seatunnel.engine.server.resourcemanager.resource.SlotProfile; +import org.apache.seatunnel.engine.server.resourcemanager.worker.WorkerProfile; +import org.apache.seatunnel.engine.server.service.slot.SlotAndWorkerProfile; + +import org.junit.jupiter.api.Test; + +import com.hazelcast.cluster.Address; +import com.hazelcast.spi.impl.operationservice.Operation; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class ResourceRequestHandlerTest { + + @Test + void partialAllocationFailureWaitsForRollbackBeforeCompletingRequest() throws Exception { + Address workerAddress = new Address("localhost", 5801); + SlotProfile successfulSlot = slot(workerAddress, 1); + CompletableFuture<Void> releaseFuture = new CompletableFuture<>(); + TestResourceManager resourceManager = + new TestResourceManager(successfulSlot, releaseFuture, 0); + + CompletableFuture<List<SlotProfile>> request = + requestTwoSlots(workerAddress, resourceManager); + + assertFalse(request.isDone(), "request must wait until the acquired slot is released"); + assertEquals(Collections.singletonList(successfulSlot), resourceManager.releasedSlots); + + releaseFuture.complete(null); + + assertThrows(CompletionException.class, request::join); + assertEquals(1, resourceManager.releaseAttempts.get()); + } + + @Test + void partialAllocationFailureRetriesRollback() throws Exception { + Address workerAddress = new Address("localhost", 5801); + TestResourceManager resourceManager = + new TestResourceManager( + slot(workerAddress, 1), CompletableFuture.completedFuture(null), 1); + + CompletableFuture<List<SlotProfile>> request = + requestTwoSlots(workerAddress, resourceManager); + + assertThrows(CompletionException.class, request::join); + assertEquals(2, resourceManager.releaseAttempts.get()); + } + + @Test + void rollbackFailureIsAttachedAfterRetriesAreExhausted() throws Exception { + Address workerAddress = new Address("localhost", 5801); + TestResourceManager resourceManager = + new TestResourceManager( + slot(workerAddress, 1), + CompletableFuture.completedFuture(null), + Constant.OPERATION_RETRY_TIME); + + CompletionException exception = + assertThrows( + CompletionException.class, + () -> requestTwoSlots(workerAddress, resourceManager).join()); + + assertEquals(Constant.OPERATION_RETRY_TIME, resourceManager.releaseAttempts.get()); + assertTrue(hasSuppressedFailure(exception)); + } + + @Test + void synchronousRollbackFailureIsRetried() throws Exception { + Address workerAddress = new Address("localhost", 5801); + SynchronousReleaseResourceManager resourceManager = + new SynchronousReleaseResourceManager(slot(workerAddress, 1)); + + CompletableFuture<List<SlotProfile>> request = + requestTwoSlots(workerAddress, resourceManager); + + assertThrows(CompletionException.class, request::join); + assertEquals(2, resourceManager.synchronousReleaseAttempts.get()); + } + + @Test + void rollbackRetriesOnlyTheSlotWhoseReleaseFailed() throws Exception { + Address workerAddress = new Address("localhost", 5801); + SlotProfile first = slot(workerAddress, 1); + SlotProfile second = slot(workerAddress, 2); + SelectiveReleaseResourceManager resourceManager = + new SelectiveReleaseResourceManager(first, second); + + CompletableFuture<List<SlotProfile>> request = + new ResourceRequestHandler( + 103L, + Arrays.asList( + new ResourceProfile(), + new ResourceProfile(), + new ResourceProfile()), + workers(workerAddress, first, second, slot(workerAddress, 3)), + resourceManager, + new RandomStrategy()) + .request(Collections.emptyMap()); + + assertThrows(CompletionException.class, request::join); + assertEquals(1, resourceManager.releaseAttempts.get(1).get()); + assertEquals(2, resourceManager.releaseAttempts.get(2).get()); + } + + @Test + void retryPreservesOriginalRequestIndexes() throws Exception { + Address workerAddress = new Address("localhost", 5801); + List<SlotProfile> slots = + Arrays.asList( + slot(workerAddress, 1), slot(workerAddress, 2), slot(workerAddress, 3)); + IndexedRetryResourceManager resourceManager = new IndexedRetryResourceManager(slots); + + List<SlotProfile> result = + new ResourceRequestHandler( + 101L, + Arrays.asList( + new ResourceProfile(), + new ResourceProfile(), + new ResourceProfile()), + workers(workerAddress, slots.toArray(new SlotProfile[0])), + resourceManager, + new RandomStrategy()) + .request(Collections.emptyMap()) + .join(); + + assertEquals(slots, result); + assertEquals(4, resourceManager.requestCount.get()); + } + + @Test + void dynamicWorkerReceivesEveryMissingResourceIndex() throws Exception { + Address workerAddress = new Address("localhost", 5801); + ResourceProfile available = resource(1, 1); + ResourceProfile missing = resource(2, 2); + SlotProfile initiallyAllocated = new SlotProfile(workerAddress, 1, available, "initial"); + SlotProfile dynamicallyAllocated = new SlotProfile(workerAddress, 2, missing, "dynamic"); + DynamicResourceManager resourceManager = + new DynamicResourceManager(initiallyAllocated, dynamicallyAllocated); + WorkerProfile worker = + new WorkerProfile( + workerAddress, + missing, + available, + true, + new SlotProfile[0], + new SlotProfile[] {initiallyAllocated}, + Collections.emptyMap()); + ConcurrentMap<Address, WorkerProfile> workers = new ConcurrentHashMap<>(); + workers.put(workerAddress, worker); + + List<SlotProfile> result = + new ResourceRequestHandler( + 102L, + Arrays.asList(available, missing), + workers, + resourceManager, + new RandomStrategy()) + .request(Collections.emptyMap()) + .join(); + + assertEquals(Arrays.asList(initiallyAllocated, dynamicallyAllocated), result); + assertEquals(1, resourceManager.requestedDynamicResources.size()); + assertSame(missing, resourceManager.requestedDynamicResources.get(0)); + } + + private static CompletableFuture<List<SlotProfile>> requestTwoSlots( + Address workerAddress, TestResourceManager resourceManager) { + return new ResourceRequestHandler( + 100L, + Arrays.asList(new ResourceProfile(), new ResourceProfile()), + workers(workerAddress), + resourceManager, + new RandomStrategy()) + .request(Collections.emptyMap()); + } + + private static boolean hasSuppressedFailure(Throwable throwable) { + Throwable current = throwable; + while (current != null) { + if (current.getSuppressed().length > 0) { + return true; + } + current = current.getCause(); + } + return false; + } + + private static ConcurrentMap<Address, WorkerProfile> workers(Address address) { + return workers(address, slot(address, 1), slot(address, 2)); + } + + private static ConcurrentMap<Address, WorkerProfile> workers( + Address address, SlotProfile... slots) { + WorkerProfile worker = + new WorkerProfile( + address, + new ResourceProfile(), + new ResourceProfile(), + false, + new SlotProfile[0], + slots, + Collections.emptyMap()); + ConcurrentMap<Address, WorkerProfile> workers = new ConcurrentHashMap<>(); + workers.put(address, worker); + return workers; + } + + private static SlotProfile slot(Address address, int id) { + return new SlotProfile(address, id, new ResourceProfile(), "sequence-" + id); + } + + private static ResourceProfile resource(int cpu, long memory) { + return new ResourceProfile(CPU.of(cpu), Memory.of(memory)); + } + + private static class IndexedRetryResourceManager extends AbstractResourceManager { + private final List<SlotProfile> slots; + private final AtomicInteger requestCount = new AtomicInteger(); + + private IndexedRetryResourceManager(List<SlotProfile> slots) { + super(null, new EngineConfig()); + this.slots = slots; + } + + @Override + protected <E> CompletableFuture<E> sendToMember(Operation operation, Address address) { + int request = requestCount.getAndIncrement(); + SlotProfile slotProfile; + if (request == 0) { + slotProfile = slots.get(0); + } else if (request == 1) { + slotProfile = null; + } else if (request == 2) { + slotProfile = slots.get(2); + } else if (request == 3) { + slotProfile = slots.get(1); + } else { + CompletableFuture<E> failed = new CompletableFuture<>(); + failed.completeExceptionally(new NoEnoughResourceException("unexpected retry")); + return failed; + } + WorkerProfile worker = workerProfile(address, slots, slotProfile); + return completed(new SlotAndWorkerProfile(worker, slotProfile)); + } + + @Override + public CompletableFuture<Void> releaseResource(long jobId, SlotProfile slotProfile) { + return CompletableFuture.completedFuture(null); + } + + @Override + public void heartbeat(WorkerProfile workerProfile) {} + } + + private static class DynamicResourceManager extends AbstractResourceManager { + private final SlotProfile initiallyAllocated; + private final SlotProfile dynamicallyAllocated; + private final AtomicInteger requestCount = new AtomicInteger(); + private List<ResourceProfile> requestedDynamicResources = Collections.emptyList(); + + private DynamicResourceManager( + SlotProfile initiallyAllocated, SlotProfile dynamicallyAllocated) { + super(null, new EngineConfig()); + this.initiallyAllocated = initiallyAllocated; + this.dynamicallyAllocated = dynamicallyAllocated; + } + + @Override + protected boolean supportDynamicWorker() { + return true; + } + + @Override + protected void findNewWorker( + List<ResourceProfile> resourceProfiles, Map<String, String> tagFilter) { + requestedDynamicResources = new ArrayList<>(resourceProfiles); + } + + @Override + public CompletableFuture<List<SlotProfile>> applyResources( + long jobId, List<ResourceProfile> resourceProfiles, Map<String, String> tagFilter) { + if (resourceProfiles.isEmpty()) { + CompletableFuture<List<SlotProfile>> failed = new CompletableFuture<>(); + failed.completeExceptionally(new NoEnoughResourceException("empty request")); + return failed; + } + return CompletableFuture.completedFuture( + Collections.singletonList(dynamicallyAllocated)); + } + + @Override + protected <E> CompletableFuture<E> sendToMember(Operation operation, Address address) { + SlotProfile slotProfile = + requestCount.getAndIncrement() == 0 ? initiallyAllocated : null; + WorkerProfile worker = + workerProfile( + address, + Arrays.asList(initiallyAllocated, dynamicallyAllocated), + slotProfile); + return completed(new SlotAndWorkerProfile(worker, slotProfile)); + } + + @Override + public CompletableFuture<Void> releaseResource(long jobId, SlotProfile slotProfile) { + return CompletableFuture.completedFuture(null); + } + + @Override + public void heartbeat(WorkerProfile workerProfile) {} + } + + private static WorkerProfile workerProfile( + Address address, List<SlotProfile> slots, SlotProfile assignedSlot) { + return new WorkerProfile( + address, + new ResourceProfile(), + new ResourceProfile(), + false, + assignedSlot == null ? new SlotProfile[0] : new SlotProfile[] {assignedSlot}, + slots.toArray(new SlotProfile[0]), + Collections.emptyMap()); + } + + @SuppressWarnings("unchecked") + private static <E> CompletableFuture<E> completed(Object value) { + return (CompletableFuture<E>) CompletableFuture.completedFuture(value); + } + + private static class TestResourceManager extends AbstractResourceManager { + private final AtomicInteger requestCount = new AtomicInteger(); + private final SlotProfile successfulSlot; + private final CompletableFuture<Void> releaseFuture; + private final int releaseFailuresBeforeSuccess; + private final AtomicInteger releaseAttempts = new AtomicInteger(); + private final List<SlotProfile> releasedSlots = new ArrayList<>(); + + private TestResourceManager( + SlotProfile successfulSlot, + CompletableFuture<Void> releaseFuture, + int releaseFailuresBeforeSuccess) { + super(null, new EngineConfig()); + this.successfulSlot = successfulSlot; + this.releaseFuture = releaseFuture; + this.releaseFailuresBeforeSuccess = releaseFailuresBeforeSuccess; + } + + @Override + protected <E> CompletableFuture<E> sendToMember(Operation operation, Address address) { + if (!(operation instanceof RequestSlotOperation)) { + throw new UnsupportedOperationException(); + } + if (requestCount.getAndIncrement() == 0) { + WorkerProfile worker = + new WorkerProfile( + address, + new ResourceProfile(), + new ResourceProfile(), + false, + new SlotProfile[] {successfulSlot}, + new SlotProfile[0], + Collections.emptyMap()); + return ResourceRequestHandlerTest.completed( + new SlotAndWorkerProfile(worker, successfulSlot)); + } + CompletableFuture<E> failed = new CompletableFuture<>(); + failed.completeExceptionally(new NoEnoughResourceException("test failure")); + return failed; + } + + @Override + public CompletableFuture<Void> releaseResource(long jobId, SlotProfile slotProfile) { + releasedSlots.add(slotProfile); + if (releaseAttempts.getAndIncrement() < releaseFailuresBeforeSuccess) { + CompletableFuture<Void> failed = new CompletableFuture<>(); + failed.completeExceptionally(new RuntimeException("transient release failure")); + return failed; + } + return releaseFuture; + } + + @Override + public void heartbeat(WorkerProfile workerProfile) { + // Keep the focused worker snapshot stable while the request and rollback complete. + } + } + + private static class SynchronousReleaseResourceManager extends TestResourceManager { + private final AtomicInteger synchronousReleaseAttempts = new AtomicInteger(); + + private SynchronousReleaseResourceManager(SlotProfile successfulSlot) { + super(successfulSlot, CompletableFuture.completedFuture(null), 0); + } + + @Override + public CompletableFuture<Void> releaseResource(long jobId, SlotProfile slotProfile) { + if (synchronousReleaseAttempts.getAndIncrement() == 0) { + throw new IllegalStateException("coordinator is closing"); + } + return CompletableFuture.completedFuture(null); + } + } + + private static class SelectiveReleaseResourceManager extends AbstractResourceManager { + private final List<SlotProfile> successfulSlots; + private final AtomicInteger requestCount = new AtomicInteger(); + private final ConcurrentMap<Integer, AtomicInteger> releaseAttempts = + new ConcurrentHashMap<>(); + + private SelectiveReleaseResourceManager(SlotProfile first, SlotProfile second) { + super(null, new EngineConfig()); + this.successfulSlots = Arrays.asList(first, second); + } + + @Override + protected <E> CompletableFuture<E> sendToMember(Operation operation, Address address) { + int request = requestCount.getAndIncrement(); + if (request < successfulSlots.size()) { + SlotProfile slotProfile = successfulSlots.get(request); + WorkerProfile worker = workerProfile(address, successfulSlots, slotProfile); + return completed(new SlotAndWorkerProfile(worker, slotProfile)); + } + CompletableFuture<E> failed = new CompletableFuture<>(); + failed.completeExceptionally(new NoEnoughResourceException("test failure")); + return failed; + } + + @Override + public CompletableFuture<Void> releaseResource(long jobId, SlotProfile slotProfile) { + int attempt = + releaseAttempts + .computeIfAbsent( + slotProfile.getSlotID(), ignored -> new AtomicInteger()) + .incrementAndGet(); + if (slotProfile.getSlotID() == 2 && attempt == 1) { + CompletableFuture<Void> failed = new CompletableFuture<>(); + failed.completeExceptionally(new IllegalStateException("transient failure")); + return failed; + } + return CompletableFuture.completedFuture(null); + } + + @Override + public void heartbeat(WorkerProfile workerProfile) {} + } +}
