taski 0.5.0 → 0.7.1

data/lib/taski/execution/executor.rb

@@ -1,129 +1,204 @@
 # frozen_string_literal: true
 
 require "monitor"
-require "etc"
 
 module Taski
   module Execution
     # Producer-Consumer pattern executor for parallel task execution.
     #
-    # Architecture:
-    # - Main Thread: Manages all state, coordinates execution, handles events
-    # - Worker Threads: Execute tasks and send completion events
+    # Executor is the orchestrator that coordinates all execution components.
+    #
+    # == Architecture
+    #
+    #   Executor
+    #     ├── Scheduler: Dependency management and execution order
+    #     ├── WorkerPool: Thread management and task distribution
+    #     └── ExecutionContext: Observer notifications and output capture
+    #             └── Observers (e.g., TreeProgressDisplay)
+    #
+    # == Execution Flow
     #
-    # Communication Queues:
-    # - Execution Queue (Main -> Worker): Tasks ready to execute
+    # 1. Build dependency graph via Scheduler
+    # 2. Set up progress display via ExecutionContext
+    # 3. Start WorkerPool threads
+    # 4. Enqueue ready tasks (no dependencies) to WorkerPool
+    # 5. Run event loop:
+    #    - Pop completion events from workers
+    #    - Mark completed in Scheduler
+    #    - Enqueue newly ready tasks to WorkerPool
+    # 6. Shutdown WorkerPool when root task completes
+    # 7. Teardown progress display
+    #
+    # == Communication Queues
+    #
+    # - Execution Queue (Main -> Worker): Tasks ready to execute (via WorkerPool)
     # - Completion Queue (Worker -> Main): Events from workers
+    #
+    # == Thread Safety
+    #
+    # - Main Thread: Manages all state, coordinates execution, handles events
+    # - Worker Threads: Execute tasks and send completion events (via WorkerPool)
     class Executor
-      # Task execution states for the executor's internal tracking
-      STATE_PENDING = :pending
-      STATE_ENQUEUED = :enqueued
-      STATE_COMPLETED = :completed
-
       class << self
         # Execute a task and all its dependencies
         # @param root_task_class [Class] The root task class to execute
         # @param registry [Registry] The task registry
-        def execute(root_task_class, registry:)
-          new(registry: registry).execute(root_task_class)
+        ##
+        # Create a new Executor and run execution for the specified root task class.
+        # @param root_task_class [Class] The top-level task class to execute.
+        # @param registry [Taski::Registry] Registry providing task definitions and state.
+        # @param execution_context [ExecutionContext, nil] Optional execution context to use; when nil a default context is created.
+        # @return [Object] The result returned by the execution of the root task.
+        def execute(root_task_class, registry:, execution_context: nil)
+          new(registry: registry, execution_context: execution_context).execute(root_task_class)
+        end
+
+        # Execute clean for a task and all its dependencies (in reverse order)
+        # @param root_task_class [Class] The root task class to clean
+        # @param registry [Registry] The task registry
+        ##
+        # Runs reverse-order clean execution beginning at the given root task class.
+        # @param [Class] root_task_class - The root task class whose dependency graph will drive the clean run.
+        # @param [Object] registry - Task registry used to resolve and track tasks during execution.
+        # @param [ExecutionContext, nil] execution_context - Optional execution context for observers and output capture; if `nil`, a default context is created.
+        def execute_clean(root_task_class, registry:, execution_context: nil)
+          new(registry: registry, execution_context: execution_context).execute_clean(root_task_class)
         end
       end
 
-      def initialize(registry:, worker_count: nil)
+      ##
+      # Initialize an Executor and its internal coordination components.
+      # @param [Object] registry - Task registry used to look up task definitions and state.
+      # @param [Integer, nil] worker_count - Optional number of worker threads to use; when `nil`,
+      #   uses Taski.args_worker_count which retrieves the worker count from the runtime args.
+      # @param [Taski::Execution::ExecutionContext, nil] execution_context - Optional execution context for observers and output capture; when `nil` a default context (with progress observer and execution trigger) is created.
+      def initialize(registry:, worker_count: nil, execution_context: nil)
         @registry = registry
-        @worker_count = worker_count || default_worker_count
-        @execution_queue = Queue.new
         @completion_queue = Queue.new
-        @workers = []
 
-        # State managed by main thread only
-        @dependencies = {}
-        @task_states = {}
-        @completed_tasks = Set.new
+        # ExecutionContext for observer pattern and output capture
+        @execution_context = execution_context || create_default_execution_context
+
+        # Scheduler for dependency management
+        @scheduler = Scheduler.new
+
+        # Determine effective worker count: explicit param > args > default
+        # Store as instance variable for consistent use in both run and clean phases
+        @effective_worker_count = worker_count || Taski.args_worker_count
+
+        # WorkerPool for thread management
+        @worker_pool = WorkerPool.new(
+          registry: @registry,
+          worker_count: @effective_worker_count
+        ) { |task_class, wrapper| execute_task(task_class, wrapper) }
       end
 
       # Execute root task and all dependencies
-      # @param root_task_class [Class] The root task class to execute
+      ##
+      # Execute the task graph rooted at the given task class.
+      #
+      # Builds the dependency graph, starts progress reporting and worker threads,
+      # enqueues tasks that are ready (no unmet dependencies), and processes worker
+      # completion events until the root task finishes. After completion or abort,
+      # shuts down workers, stops progress reporting, and restores stdout capture if
+      # this executor configured it.
+      # @param root_task_class [Class] The root task class to execute.
       def execute(root_task_class)
         # Build dependency graph from static analysis
-        build_dependency_graph(root_task_class)
-
-        # Set up tree progress display with root task (before start)
-        setup_tree_progress(root_task_class)
+        @scheduler.build_dependency_graph(root_task_class)
 
-        # Start progress display automatically for tree progress
-        start_progress_display
+        with_display_lifecycle(root_task_class) do
+          # Start worker threads
+          @worker_pool.start
 
-        # Start worker threads
-        start_workers
+          # Enqueue tasks with no dependencies
+          enqueue_ready_tasks
 
-        # Enqueue tasks with no dependencies
-        enqueue_ready_tasks
+          # Main event loop - continues until root task completes
+          run_main_loop(root_task_class)
 
-        # Main event loop - continues until root task completes
-        run_main_loop(root_task_class)
-
-        # Shutdown workers
-        shutdown_workers
+          # Shutdown workers
+          @worker_pool.shutdown
+        end
 
-        # Stop progress display
-        stop_progress_display
+        # Raise aggregated errors if any tasks failed
+        raise_if_any_failures
       end
 
-      private
+      # Execute clean for root task and all dependencies (in reverse dependency order)
+      # Clean operations run in reverse: root task cleans first, then dependencies
+      ##
+      # Executes the clean workflow for the given root task in reverse dependency order.
+      # Sets up progress display and optional output capture, starts a dedicated clean worker pool,
+      # enqueues ready-to-clean tasks, processes completion events until all tasks are cleaned,
+      # then shuts down workers and tears down progress and output capture as needed.
+      # @param [Class] root_task_class - The root task class to clean
+      def execute_clean(root_task_class)
+        # Build reverse dependency graph for clean order
+        # This must happen first to ensure root task and all static dependencies are included
+        @scheduler.build_reverse_dependency_graph(root_task_class)
+
+        # Merge runtime dependencies (e.g., Section's dynamically selected implementations)
+        # This allows clean to include tasks that were selected at runtime during run phase
+        runtime_deps = @execution_context.runtime_dependencies
+        @scheduler.merge_runtime_dependencies(runtime_deps)
+
+        with_display_lifecycle(root_task_class) do
+          # Create a new worker pool for clean operations
+          # Uses the same worker count as the run phase
+          @clean_worker_pool = WorkerPool.new(
+            registry: @registry,
+            worker_count: @effective_worker_count
+          ) { |task_class, wrapper| execute_clean_task(task_class, wrapper) }
+
+          # Start worker threads
+          @clean_worker_pool.start
+
+          # Enqueue tasks ready for clean (no reverse dependencies)
+          enqueue_ready_clean_tasks
+
+          # Main event loop - continues until all tasks are cleaned
+          run_clean_main_loop(root_task_class)
+
+          # Shutdown workers
+          @clean_worker_pool.shutdown
+        end
 
-      def default_worker_count
-        Etc.nprocessors.clamp(2, 8)
+        # Raise aggregated errors if any clean tasks failed
+        raise_if_any_clean_failures
       end
 
-      # Build dependency graph by traversing from root task
-      # Populates @dependencies and @task_states
-      def build_dependency_graph(root_task_class)
-        # @type var queue: Array[singleton(Taski::Task)]
-        queue = [root_task_class]
-
-        while (task_class = queue.shift)
-          next if @task_states.key?(task_class)
-
-          deps = task_class.cached_dependencies
-          @dependencies[task_class] = deps.dup
-          @task_states[task_class] = STATE_PENDING
-
-          deps.each { |dep| queue << dep }
-        end
-      end
+      private
 
-      # Enqueue tasks that have all dependencies completed
+      # Enqueue all tasks that are ready to execute
       def enqueue_ready_tasks
-        @task_states.each_key do |task_class|
-          next unless @task_states[task_class] == STATE_PENDING
-          next unless ready_to_execute?(task_class)
-
+        @scheduler.next_ready_tasks.each do |task_class|
           enqueue_task(task_class)
         end
       end
 
-      # Check if a task is ready to execute
-      def ready_to_execute?(task_class)
-        task_deps = @dependencies[task_class] || Set.new
-        task_deps.subset?(@completed_tasks)
-      end
-
       # Enqueue a single task for execution
       def enqueue_task(task_class)
         return if @registry.abort_requested?
 
-        @task_states[task_class] = STATE_ENQUEUED
+        @scheduler.mark_enqueued(task_class)
 
         wrapper = get_or_create_wrapper(task_class)
-        return unless wrapper.mark_running
-
-        Taski.progress_display&.register_task(task_class)
-        Taski.progress_display&.update_task(task_class, state: :running)
+        unless wrapper.mark_running
+          # Task is either already running or completed in another context (e.g., parent Executor)
+          # Wait for the task to complete if it's running elsewhere
+          wrapper.wait_for_completion
+
+          # Now mark it as completed in the scheduler and enqueue newly ready tasks
+          @scheduler.mark_completed(task_class)
+          enqueue_ready_tasks
+          return
+        end
 
-        @execution_queue.push({task_class: task_class, wrapper: wrapper})
+        @execution_context.notify_task_registered(task_class)
+        @execution_context.notify_task_started(task_class)
 
-        debug_log("Enqueued: #{task_class}")
+        @worker_pool.enqueue(task_class, wrapper)
       end
 
       # Get or create a task wrapper via Registry
@@ -131,40 +206,16 @@ module Taski
         @registry.get_or_create(task_class) do
           task_instance = task_class.allocate
           task_instance.send(:initialize)
-          TaskWrapper.new(task_instance, registry: @registry)
-        end
-      end
-
-      # Start worker threads
-      def start_workers
-        @worker_count.times do
-          worker = Thread.new { worker_loop }
-          @workers << worker
-          @registry.register_thread(worker)
-        end
-      end
-
-      # Worker thread main loop
-      def worker_loop
-        loop do
-          work_item = @execution_queue.pop
-          break if work_item == :shutdown
-
-          task_class = work_item[:task_class]
-          wrapper = work_item[:wrapper]
-
-          debug_log("Worker executing: #{task_class}")
-
-          execute_task(task_class, wrapper)
+          TaskWrapper.new(task_instance, registry: @registry, execution_context: @execution_context)
         end
       end
 
-      # Execute a task and send completion event
+      # Execute a task and send completion event (called by WorkerPool)
       def execute_task(task_class, wrapper)
         return if @registry.abort_requested?
 
-        begin
-          result = execute_task_run(wrapper)
+        with_task_context(task_class) do
+          result = wrapper.task.run
           wrapper.mark_completed(result)
           @completion_queue.push({task_class: task_class, wrapper: wrapper})
         rescue Taski::TaskAbortException => e
@@ -177,71 +228,291 @@ module Taski
         end
       end
 
-      # Execute task run method
-      # Note: Previously captured stdout for progress display, but this was removed
-      # due to thread-safety concerns with global $stdout mutation.
-      def execute_task_run(wrapper)
-        wrapper.task.run
-      end
-
       # Main thread event loop - continues until root task completes
       def run_main_loop(root_task_class)
-        until @completed_tasks.include?(root_task_class)
-          break if @registry.abort_requested? && no_running_tasks?
+        until @scheduler.completed?(root_task_class)
+          break if @registry.abort_requested? && !@scheduler.running_tasks?
 
           event = @completion_queue.pop
           handle_completion(event)
         end
       end
 
-      def no_running_tasks?
-        @task_states.values.none? { |state| state == STATE_ENQUEUED }
-      end
-
-      # Handle task completion event
+      ##
+      # Marks the given task as completed in the scheduler and enqueues any tasks that become ready as a result.
+      # @param [Hash] event - Completion event containing information about the finished task.
+      # @param [Class] event[:task_class] - The task class that completed.
       def handle_completion(event)
         task_class = event[:task_class]
 
         debug_log("Completed: #{task_class}")
 
-        @task_states[task_class] = STATE_COMPLETED
-        @completed_tasks.add(task_class)
+        @scheduler.mark_completed(task_class)
 
         # Enqueue newly ready tasks
         enqueue_ready_tasks
       end
 
-      # Shutdown worker threads
-      def shutdown_workers
-        @worker_count.times { @execution_queue.push(:shutdown) }
-        @workers.each(&:join)
+      # ========================================
+      # Clean Execution Methods
+      # ========================================
+
+      ##
+      # Enqueues all tasks that are currently ready to be cleaned.
+      def enqueue_ready_clean_tasks
+        @scheduler.next_ready_clean_tasks.each do |task_class|
+          enqueue_clean_task(task_class)
+        end
       end
 
-      def setup_tree_progress(root_task_class)
-        progress = Taski.progress_display
-        return unless progress.is_a?(TreeProgressDisplay)
+      ##
+      # Enqueues a single task for reverse-order (clean) execution.
+      # If execution has been aborted, does nothing. Marks the task as clean-enqueued,
+      # skips if the task is not registered or not eligible to run, notifies the
+      # execution context that cleaning has started, and schedules the task on the
+      # clean worker pool.
+      # @param [Class] task_class - The task class to enqueue for clean execution.
+      def enqueue_clean_task(task_class)
+        return if @registry.abort_requested?
 
-        progress.set_root_task(root_task_class)
+        @scheduler.mark_clean_enqueued(task_class)
+
+        wrapper = get_or_create_wrapper(task_class)
+        return unless wrapper.mark_clean_running
+
+        @execution_context.notify_clean_started(task_class)
+
+        @clean_worker_pool.enqueue(task_class, wrapper)
       end
 
-      def start_progress_display
-        progress = Taski.progress_display
-        return unless progress.is_a?(TreeProgressDisplay)
+      ##
+      # Executes the clean lifecycle for a task and emits a completion event.
+      #
+      # Runs the task's `clean` method, updates the provided wrapper with success or failure
+      # (which handles timing and observer notification), and pushes a completion event onto
+      # the executor's completion queue.
+      # This method respects an abort requested state from the registry (no-op if abort already requested)
+      # and triggers a registry abort when a `Taski::TaskAbortException` is raised.
+      # It also starts and stops per-task output capture when available and sets the thread-local
+      # `ExecutionContext.current` for the duration of the clean.
+      # @param [Class] task_class - The task class being cleaned.
+      # @param [Taski::Execution::TaskWrapper] wrapper - The wrapper instance for the task, used to record clean success or failure.
+      def execute_clean_task(task_class, wrapper)
+        return if @registry.abort_requested?
+
+        with_task_context(task_class) do
+          result = wrapper.task.clean
+          wrapper.mark_clean_completed(result)
+          @completion_queue.push({task_class: task_class, wrapper: wrapper, clean: true})
+        rescue Taski::TaskAbortException => e
+          @registry.request_abort!
+          wrapper.mark_clean_failed(e)
+          @completion_queue.push({task_class: task_class, wrapper: wrapper, error: e, clean: true})
+        rescue => e
+          wrapper.mark_clean_failed(e)
+          @completion_queue.push({task_class: task_class, wrapper: wrapper, error: e, clean: true})
+        end
+      end
+
+      ##
+      # Runs the main event loop that processes clean completion events until all tasks have been cleaned.
+      # Continuously pops events from the internal completion queue and delegates them to the clean completion handler,
+      # stopping early if an abort is requested and no clean tasks are running.
+      # @param [Class] root_task_class - The root task class that defines the overall clean lifecycle.
+      def run_clean_main_loop(root_task_class)
+        # Find all tasks in the dependency graph
+        # Continue until all tasks have been cleaned
+        until all_tasks_cleaned?
+          break if @registry.abort_requested? && !@scheduler.running_clean_tasks?
+
+          event = @completion_queue.pop
+          handle_clean_completion(event)
+        end
+      end
+
+      ##
+      # Processes a clean completion event and advances the cleaning workflow.
+      # Marks the completed task in the scheduler and enqueues any tasks that become ready to clean.
+      # @param [Hash] event - A completion event hash containing the `:task_class` key for the task that finished cleaning.
+      def handle_clean_completion(event)
+        task_class = event[:task_class]
 
-        progress.start
+        debug_log("Clean completed: #{task_class}")
+
+        @scheduler.mark_clean_completed(task_class)
+
+        # Enqueue newly ready clean tasks
+        enqueue_ready_clean_tasks
+      end
+
+      ##
+      # Determines whether all tasks have finished their clean phase.
+      # @return [Boolean] `true` if there are no ready-to-clean tasks and no running clean tasks, `false` otherwise.
+      def all_tasks_cleaned?
+        @scheduler.next_ready_clean_tasks.empty? && !@scheduler.running_clean_tasks?
+      end
+
+      # Notify observers about the root task
+      # @param root_task_class [Class] The root task class
+      # @return [void]
+      def setup_progress_display(root_task_class)
+        @execution_context.notify_set_root_task(root_task_class)
+      end
+
+      # Set up output capture if progress display is active and not already set up
+      # @return [Boolean] true if this executor set up the capture
+      def setup_output_capture_if_needed
+        return false unless Taski.progress_display
+        return false if @execution_context.output_capture_active?
+
+        @execution_context.setup_output_capture($stdout)
+        true
+      end
+
+      # Tear down output capture and restore original $stdout
+      # @return [void]
+      def teardown_output_capture
+        @execution_context.teardown_output_capture
+      end
+
+      def start_progress_display
+        @execution_context.notify_start
       end
 
       def stop_progress_display
+        @execution_context.notify_stop
+      end
+
+      # Execute a block with task-local context set up.
+      # Sets ExecutionContext.current, Taski.current_registry, and output capture.
+      # Cleans up all context in ensure block.
+      #
+      # @param task_class [Class] The task class being executed
+      # @yield The block to execute with context set up
+      def with_task_context(task_class)
+        output_capture = @execution_context.output_capture
+        output_capture&.start_capture(task_class)
+
+        ExecutionContext.current = @execution_context
+        Taski.set_current_registry(@registry)
+
+        yield
+      ensure
+        output_capture&.stop_capture
+        ExecutionContext.current = nil
+        Taski.clear_current_registry
+      end
+
+      # Execute a block with progress display and output capture lifecycle.
+      # Sets up progress display, output capture, starts display, then yields.
+      # Ensures proper cleanup even on interrupt.
+      #
+      # @param root_task_class [Class] The root task class
+      # @yield The block to execute
+      def with_display_lifecycle(root_task_class)
+        setup_progress_display(root_task_class)
+        should_teardown_capture = setup_output_capture_if_needed
+        start_progress_display
+
+        yield
+      ensure
+        stop_progress_display
+        @saved_output_capture = @execution_context.output_capture
+        teardown_output_capture if should_teardown_capture
+      end
+
+      def create_default_execution_context
+        context = ExecutionContext.new
         progress = Taski.progress_display
-        return unless progress.is_a?(TreeProgressDisplay)
+        context.add_observer(progress) if progress
+
+        # Set execution trigger to break circular dependency with TaskWrapper
+        context.execution_trigger = ->(task_class, registry) do
+          Executor.execute(task_class, registry: registry, execution_context: context)
+        end
 
-        progress.stop
+        context
       end
 
       def debug_log(message)
         return unless ENV["TASKI_DEBUG"]
         puts "[Executor] #{message}"
       end
+
+      # Raise error(s) if any tasks failed during execution
+      # TaskAbortException: raised directly (abort takes priority)
+      # All other errors: raises AggregateError containing all failures
+      def raise_if_any_failures
+        raise_if_any_failures_from(
+          @registry.failed_wrappers,
+          error_accessor: ->(w) { w.error }
+        )
+      end
+
+      # Raise error(s) if any tasks failed during clean execution
+      # TaskAbortException: raised directly (abort takes priority)
+      # All other errors: raises AggregateError containing all failures
+      def raise_if_any_clean_failures
+        raise_if_any_failures_from(
+          @registry.failed_clean_wrappers,
+          error_accessor: ->(w) { w.clean_error }
+        )
+      end
+
+      # Generic method to raise errors from failed wrappers
+      # @param failed_wrappers [Array<TaskWrapper>] Failed wrappers
+      # @param error_accessor [Proc] Lambda to extract error from wrapper
+      def raise_if_any_failures_from(failed_wrappers, error_accessor:)
+        return if failed_wrappers.empty?
+
+        # TaskAbortException takes priority - raise the first one directly
+        abort_wrapper = failed_wrappers.find { |w| error_accessor.call(w).is_a?(TaskAbortException) }
+        raise error_accessor.call(abort_wrapper) if abort_wrapper
+
+        # Flatten nested AggregateErrors and deduplicate by original error object_id
+        failures = flatten_failures_from(failed_wrappers, error_accessor: error_accessor)
+        unique_failures = failures.uniq { |f| error_identity(f.error) }
+
+        raise AggregateError.new(unique_failures)
+      end
+
+      # Flatten AggregateErrors into individual TaskFailure objects
+      # Wraps original errors with task-specific Error class for rescue matching
+      # @param failed_wrappers [Array<TaskWrapper>] Failed wrappers
+      # @param error_accessor [Proc] Lambda to extract error from wrapper
+      def flatten_failures_from(failed_wrappers, error_accessor:)
+        output_capture = @saved_output_capture
+
+        failed_wrappers.flat_map do |wrapper|
+          error = error_accessor.call(wrapper)
+          case error
+          when AggregateError
+            error.errors
+          else
+            wrapped_error = wrap_with_task_error(wrapper.task.class, error)
+            output_lines = output_capture&.recent_lines_for(wrapper.task.class) || []
+            [TaskFailure.new(task_class: wrapper.task.class, error: wrapped_error, output_lines: output_lines)]
+          end
+        end
+      end
+
+      # Wraps an error with the task-specific Error class
+      # @param task_class [Class] The task class
+      # @param error [Exception] The original error
+      # @return [TaskError] The wrapped error
+      def wrap_with_task_error(task_class, error)
+        # Don't double-wrap if already a TaskError
+        return error if error.is_a?(TaskError)
+
+        error_class = task_class.const_get(:Error)
+        error_class.new(error, task_class: task_class)
+      end
+
+      # Returns a unique identifier for an error, used for deduplication
+      # For TaskError, uses the wrapped cause's object_id
+      def error_identity(error)
+        error.is_a?(TaskError) ? error.cause&.object_id || error.object_id : error.object_id
+      end
     end
   end
 end