StreamingScheduler implements Scheduler { @Override public void schedulePlan(final PhysicalPlan submittedPhysicalPlan, final int maxScheduleAttempt) { taskDispatcher.run(); planStateManager.updatePlan(submittedPhysicalPlan, maxScheduleAttempt); planStateManager.storeJSON("submitted"); final List<Stage> allStages = submittedPhysicalPlan.getStageDAG().getTopologicalSort(); final List<Task> allTasks = allStages.stream().flatMap(stageToSchedule -> { final List<StageEdge> stageIncomingEdges = submittedPhysicalPlan.getStageDAG().getIncomingEdgesOf(stageToSchedule.getId()); final List<StageEdge> stageOutgoingEdges = submittedPhysicalPlan.getStageDAG().getOutgoingEdgesOf(stageToSchedule.getId()); final List<Map<String, Readable>> vertexIdToReadables = stageToSchedule.getVertexIdToReadables(); final List<String> taskIdsToSchedule = planStateManager.getTaskAttemptsToSchedule(stageToSchedule.getId()); taskIdsToSchedule.forEach(taskId -> { final int index = RuntimeIdManager.getIndexFromTaskId(taskId); stageOutgoingEdges.forEach(outEdge -> pipeManagerMaster.onTaskScheduled(outEdge.getId(), index)); }); return taskIdsToSchedule.stream().map(taskId -> new Task( submittedPhysicalPlan.getPlanId(), taskId, stageToSchedule.getExecutionProperties(), stageToSchedule.getSerializedIRDAG(), stageIncomingEdges, stageOutgoingEdges, vertexIdToReadables.get(RuntimeIdManager.getIndexFromTaskId(taskId)))); }).collect(Collectors.toList()); pendingTaskCollectionPointer.setToOverwrite(allTasks); taskDispatcher.onNewPendingTaskCollectionAvailable(); } @Inject  StreamingScheduler(final TaskDispatcher taskDispatcher,
                     final PendingTaskCollectionPointer pendingTaskCollectionPointer,
                     final ExecutorRegistry executorRegistry,
                     final PlanStateManager planStateManager,
                     final PipeManagerMaster pipeManagerMaster); @Override void schedulePlan(final PhysicalPlan submittedPhysicalPlan,
                           final int maxScheduleAttempt); @Override void updatePlan(final PhysicalPlan newPhysicalPlan); @Override void onTaskStateReportFromExecutor(final String executorId,
                                            final String taskId,
                                            final int taskAttemptIndex,
                                            final TaskState.State newState,
                                            @Nullable final String vertexPutOnHold,
                                            final TaskState.RecoverableTaskFailureCause failureCause); @Override void onSpeculativeExecutionCheck(); @Override void onExecutorAdded(final ExecutorRepresenter executorRepresenter); @Override void onExecutorRemoved(final String executorId); @Override void terminate();  }
@Test(timeout = 10000) public void testScheduleEverything() throws Exception { final PhysicalPlan physicalPlan = TestPlanGenerator.generatePhysicalPlan(TestPlanGenerator.PlanType.TwoVerticesJoined, false); final int numOfTotalTasks = physicalPlan.getStageDAG().getVertices().size() * ATTEMPTS_PER_STAGE; scheduler.schedulePlan(physicalPlan, 1); verify(pendingTaskCollectionPointer).setToOverwrite(argThat(tasks -> tasks.size() == numOfTotalTasks)); }