zhuzhurk commented on a change in pull request #12256:
URL: https://github.com/apache/flink/pull/12256#discussion_r432347996



##########
File path: 
flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocator.java
##########
@@ -96,52 +97,88 @@ public DefaultExecutionSlotAllocator(
 
                        LOG.debug("Allocate slot with id {} for execution {}", 
slotRequestId, executionVertexId);
 
-                       CompletableFuture<LogicalSlot> slotFuture = 
calculatePreferredLocations(
+                       final CompletableFuture<SlotProfile> slotProfileFuture 
= createSlotProfile(
+                               schedulingRequirements,
+                               Collections.emptySet(),
+                               
schedulingRequirements.getPhysicalSlotResourceProfile(),
+                               allPreviousAllocationIds);
+
+                       final CompletableFuture<LogicalSlot> slotFuture = 
slotProfileFuture.thenCompose(
+                               slotProfile ->
+                                       slotProviderStrategy.allocateSlot(
+                                               slotRequestId,
+                                               new ScheduledUnit(
+                                                       executionVertexId,
+                                                       slotSharingGroupId,
+                                                       
schedulingRequirements.getCoLocationConstraint()),
+                                               slotProfile));
+
+                       final SlotExecutionVertexAssignment 
slotExecutionVertexAssignment =
+                               createAndRegisterSlotExecutionVertexAssignment(
                                        executionVertexId,
-                                       
schedulingRequirements.getPreferredLocations(),
-                                       inputsLocationsRetriever,
-                                       Collections.emptySet()).thenCompose(
-                                                       
(Collection<TaskManagerLocation> preferredLocations) ->
-                                                               
slotProviderStrategy.allocateSlot(
-                                                                       
slotRequestId,
-                                                                       new 
ScheduledUnit(
-                                                                               
executionVertexId,
-                                                                               
slotSharingGroupId,
-                                                                               
schedulingRequirements.getCoLocationConstraint()),
-                                                                       
SlotProfile.priorAllocation(
-                                                                               
schedulingRequirements.getTaskResourceProfile(),
-                                                                               
schedulingRequirements.getPhysicalSlotResourceProfile(),
-                                                                               
preferredLocations,
-                                                                               
Collections.singletonList(schedulingRequirements.getPreviousAllocationId()),
-                                                                               
allPreviousAllocationIds)));
-
-                       SlotExecutionVertexAssignment 
slotExecutionVertexAssignment =
-                                       new 
SlotExecutionVertexAssignment(executionVertexId, slotFuture);
-                       // add to map first to avoid the future completed 
before added.
-                       pendingSlotAssignments.put(executionVertexId, 
slotExecutionVertexAssignment);
-
-                       slotFuture.whenComplete(
-                                       (ignored, throwable) -> {
-                                               
pendingSlotAssignments.remove(executionVertexId);
-                                               if (throwable != null) {
-                                                       
slotProviderStrategy.cancelSlotRequest(slotRequestId, slotSharingGroupId, 
throwable);
-                                               }
-                                       });
+                                       slotFuture,
+                                       slotRequestId,
+                                       slotSharingGroupId);
 
                        
slotExecutionVertexAssignments.add(slotExecutionVertexAssignment);
                }
 
                return slotExecutionVertexAssignments;
        }
 
-       private void 
validateSchedulingRequirements(Collection<ExecutionVertexSchedulingRequirements>
 schedulingRequirements) {
+       protected void 
validateSchedulingRequirements(Collection<ExecutionVertexSchedulingRequirements>
 schedulingRequirements) {
                schedulingRequirements.stream()
                        
.map(ExecutionVertexSchedulingRequirements::getExecutionVertexId)
                        .forEach(id -> checkState(
                                !pendingSlotAssignments.containsKey(id),
                                "BUG: vertex %s tries to allocate a slot when 
its previous slot request is still pending", id));
        }
 
+       protected CompletableFuture<SlotProfile> createSlotProfile(

Review comment:
       In the latest update, I removed `createSlotProfile()` because I find it 
does not reduce much code duplication but introduced a nested method with weird 
parameters to distinguish `DefaultExecutionSlotAllocator` and 
`OneSlotPerExecutionSlotAllocator `.
   
   I think the major part that would be good to extract out from 
`ExecutionSlotAllocator` is the preferred location calculation. We can 
introduce a `PreferredLocationsRetriever` which wraps 
`InputsLocationsRetriever`, `calculatePreferredLocations(...)` and 
`getPreferredLocationsBasedOnInputs(...)`. 
   I think we can open another task to do it since it is not very related  to 
this change. (I'm already working on it though.)
   WDYT?




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
[email protected]


Reply via email to