/** * Collects all {@link ExecutionTask}s of this instance. * * @return the {@link ExecutionTask}s */ public Set<ExecutionTask> collectAllTasks() { Set<ExecutionTask> allTasks = new HashSet<>(); for (ExecutionStage stage : this.getStages()) { allTasks.addAll(stage.getAllTasks()); } return allTasks; }
/** * Clean up {@link ExecutionStage}-related state and re-create {@link StageActivator}s etc. from the {@link ExecutionPlan}/ * * @param executionPlan whose {@link ExecutionStage}s will be executed * @param optimizationContext contains additional optimization info for the {@code executionPlan} */ public void prepare(ExecutionPlan executionPlan, OptimizationContext optimizationContext) { this.allStages.clear(); this.activatedStageActivators.clear(); this.suspendedStages.clear(); // Remove obsolete StageActivators (after re-optimization). this.allStages.addAll(executionPlan.getStages()); new ArrayList<>(this.pendingStageActivators.keySet()).stream() .filter(stage -> !this.allStages.contains(stage)) .forEach(this.pendingStageActivators::remove); // Create StageActivators for all ExecutionStages. for (ExecutionStage stage : this.allStages) { // Avoid re-activating already executed ExecutionStages. if (this.completedStages.contains(stage)) continue; final StageActivator activator = this.getOrCreateActivator( stage, () -> this.determineInitialOptimizationContext(stage, optimizationContext) ); this.tryToActivate(activator); } }