WIP with some failures

Author: johnnyshields

lib/minigun/execution/executor.rb

@@ -232,7 +232,6 @@ def fork_for_item(item, stage, user_context, output_queue)
               stage,
               [capture_queue],
               {},
-              {},
               stage_stats: stage_stats
             )
 
@@ -279,7 +278,6 @@ def fork_for_item(item, stage, user_context, output_queue)
             stage,
             [capture_queue],
             {},
-            {},
             stage_stats: stage_stats
           )
           if stage.respond_to?(:block) && stage.block

lib/minigun/pipeline.rb

@@ -5,7 +5,7 @@ module Minigun
   # A Pipeline can be standalone or part of a multi-pipeline Task
   class Pipeline
     attr_reader :name, :config, :stages, :hooks, :dag, :output_queues, :stats,
-                :context, :stage_hooks, :stage_input_queues, :runtime_edges, :input_queues, :parent_pipeline, :task
+                :context, :stage_hooks, :runtime_edges, :input_queues, :parent_pipeline, :task
 
     def initialize(name, task, parent_pipeline, config = {}, stages: nil, hooks: nil, stage_hooks: nil, dag: nil, stats: nil)
       @name = name
@@ -283,8 +283,8 @@ def run_pipeline(_context)
       # Insert router stages for fan-out
       insert_router_stages_for_fan_out
 
-      # Create one input queue per stage (except producers)
-      @stage_input_queues = build_stage_input_queues
+      # Create one input queue per stage (except producers) and register with Task
+      build_stage_input_queues
       @produced_count = Concurrent::AtomicFixnum.new(0)
       @stage_threads = []
 
@@ -303,10 +303,8 @@ def run_pipeline(_context)
       @stage_threads.each(&:join)
     end
 
-    # Build one input queue per stage (except producers)
+    # Build one input queue per stage (except producers) and register with Task
     def build_stage_input_queues
-      queues = {}
-
       # Find entrance infrastructure (router or single entry stage)
       entry_stages = @stages.select do |s|
         s.run_mode != :autonomous && @dag.upstream(s).empty?
@@ -318,17 +316,13 @@ def build_stage_input_queues
 
         # Entrance router uses PipelineStage's queue
         if stage == @entrance_router && @input_queues && @input_queues[:input]
-          queue = @input_queues[:input]
-          queues[stage] = queue
-          register_queue(stage, queue)
+          register_queue(stage, @input_queues[:input])
           next
         end
 
         # Single entry stage uses PipelineStage's queue directly (no router)
         if entry_stages.size == 1 && stage == entry_stages.first && @input_queues && @input_queues[:input]
-          queue = @input_queues[:input]
-          queues[stage] = queue
-          register_queue(stage, queue)
+          register_queue(stage, @input_queues[:input])
           next
         end
 
@@ -339,11 +333,8 @@ def build_stage_input_queues
                 else
                   SizedQueue.new(size) # Bounded queue with backpressure
                 end
-        queues[stage] = queue
         register_queue(stage, queue)
       end
-
-      queues
     end
 
     # Insert RouterStage instances for fan-out patterns

lib/minigun/queue_wrappers.rb

@@ -40,10 +40,9 @@ def pop
 
   # Wrapper around stage output that routes to downstream queues
   class OutputQueue
-    def initialize(stage, downstream_queues, all_stage_queues, runtime_edges, stage_stats: nil)
+    def initialize(stage, downstream_queues, runtime_edges, stage_stats: nil)
       @stage = stage
       @downstream_queues = downstream_queues # Array of Queue objects
-      @all_stage_queues = all_stage_queues   # Hash keyed by Stage objects
       @runtime_edges = runtime_edges         # Track dynamic routing (keyed by Stage objects)
       @stage_stats = stage_stats             # Stats object for tracking (optional)
       @to_cache = {}                         # Memoization cache for .to() results
@@ -79,7 +78,6 @@ def to(target)
       @to_cache[target] = OutputQueue.new(
         @stage,
         [target_queue],
-        @all_stage_queues,
         @runtime_edges,
         stage_stats: @stage_stats
       )

lib/minigun/stage.rb

@@ -9,7 +9,6 @@ module Minigun
     :stage,
     :dag,
     :runtime_edges,
-    :stage_input_queues,
     :stage_stats,
     # Worker-specific (nil/empty for producers)
     :worker,
@@ -162,11 +161,11 @@ def create_input_queue(stage_ctx)
     def create_output_queue(stage_ctx)
       # DAG and queues now use Stage objects
       downstream = stage_ctx.dag.downstream(stage_ctx.stage)
-      downstream_queues = downstream.filter_map { |ds| stage_ctx.stage_input_queues[ds] }
+      task = stage_ctx.stage.task
+      downstream_queues = downstream.filter_map { |ds| task&.find_queue(ds) }
       OutputQueue.new(
         stage_ctx.stage,
         downstream_queues,
-        stage_ctx.stage_input_queues,
         stage_ctx.runtime_edges,
         stage_stats: stage_ctx.stage_stats
       )
@@ -177,11 +176,13 @@ def send_end_signals(stage_ctx)
       dag_downstream = stage_ctx.dag.downstream(stage_ctx.stage)
       dynamic_targets = stage_ctx.runtime_edges[stage_ctx.stage].to_a
       all_targets = (dag_downstream + dynamic_targets).uniq
+      task = stage_ctx.stage.task
 
       all_targets.each do |target|
-        next unless stage_ctx.stage_input_queues[target]
+        queue = task&.find_queue(target)
+        next unless queue
 
-        stage_ctx.stage_input_queues[target] << EndOfSource.new(stage_ctx.stage)
+        queue << EndOfSource.new(stage_ctx.stage)
       end
     end
 
@@ -388,15 +389,19 @@ def initialize(name, pipeline, targets, options = {})
 
     def send_end_signals(worker_ctx)
       # Broadcast EndOfSource to ALL router targets
+      task = worker_ctx.stage.task
       @targets.each do |target|
-        worker_ctx.stage_input_queues[target] << EndOfSource.new(worker_ctx.stage)
+        queue = task&.find_queue(target)
+        queue&.<< EndOfSource.new(worker_ctx.stage)
       end
     end
   end
 
   # Broadcast router - sends each item to ALL downstream stages
   class RouterBroadcastStage < RouterStage
     def run_stage(worker_ctx)
+      task = worker_ctx.stage.task
+
       loop do
         item = worker_ctx.input_queue.pop
 
@@ -410,7 +415,8 @@ def run_stage(worker_ctx)
 
         # Broadcast to all downstream stages (fan-out semantics)
         @targets.each do |target|
-          worker_ctx.stage_input_queues[target] << item
+          queue = task&.find_queue(target)
+          queue&.<< item
         end
       end
     ensure
@@ -421,7 +427,8 @@ def run_stage(worker_ctx)
   # Round-robin router - distributes items across downstream stages
   class RouterRoundRobinStage < RouterStage
     def run_stage(worker_ctx)
-      target_queues = @targets.map { |target| worker_ctx.stage_input_queues[target] }
+      task = worker_ctx.stage.task
+      target_queues = @targets.map { |target| task&.find_queue(target) }.compact
       round_robin_index = 0
 
       loop do

lib/minigun/worker.rb

@@ -63,9 +63,11 @@ def handle_disconnected_stage(stage_ctx) # rubocop:disable Naming/PredicateMetho
 
         # Send EndOfSource to all downstream stages so they don't deadlock
         # DAG and queues now use Stage objects
+        task = stage_ctx.stage.task
         downstream = stage_ctx.dag.downstream(stage_ctx.stage)
         downstream.each do |target|
-          stage_ctx.stage_input_queues[target] << EndOfSource.new(stage_ctx.stage)
+          queue = task&.find_queue(target)
+          queue&.<< EndOfSource.new(stage_ctx.stage)
         end
 
         log_debug 'Done'
@@ -77,7 +79,7 @@ def handle_disconnected_stage(stage_ctx) # rubocop:disable Naming/PredicateMetho
 
     def create_stage_context
       dag = @pipeline.dag
-      stage_input_queues = @pipeline.stage_input_queues
+      task = @pipeline.task
 
       # Calculate sources for workers (empty for autonomous stages)
       # DAG now uses Stage objects instead of names
@@ -109,10 +111,9 @@ def create_stage_context
         stage: @stage,
         dag: dag,
         runtime_edges: @pipeline.runtime_edges,
-        stage_input_queues: stage_input_queues,
         stage_stats: stage_stats,
         # Worker-specific (nil/empty for producers)
-        input_queue: stage_input_queues[@stage],
+        input_queue: task&.find_queue(@stage),
         sources_expected: sources_expected,
         sources_done: Set.new
       )

spec/minigun/pipeline_spec.rb

@@ -4,7 +4,7 @@
 
 RSpec.describe Minigun::Pipeline do
   let(:config) { { max_threads: 3, max_processes: 2 } }
-  let(:task) { instance_double(Minigun::Task, stage_registry: nil, register_stage_queue: nil) }
+  let(:task) { Minigun::Task.new }
   let(:pipeline) { described_class.new(:test_pipeline, task, nil, config) }
 
   describe '#initialize' do

spec/minigun/stage_spec.rb

@@ -4,7 +4,7 @@
 
 RSpec.describe Minigun::Stage do
   let(:mock_registry) { instance_double(Minigun::StageRegistry, register: nil) }
-  let(:mock_task) { instance_double(Minigun::Task, stage_registry: mock_registry) }
+  let(:mock_task) { instance_double(Minigun::Task, stage_registry: mock_registry, find_queue: nil) }
   let(:mock_pipeline) { instance_double(Minigun::Pipeline, name: 'test_pipeline', task: mock_task) }
 
   describe 'base class' do
@@ -31,7 +31,7 @@
 
 RSpec.describe Minigun::ProducerStage do
   let(:mock_registry) { instance_double(Minigun::StageRegistry, register: nil) }
-  let(:mock_task) { instance_double(Minigun::Task, stage_registry: mock_registry) }
+  let(:mock_task) { instance_double(Minigun::Task, stage_registry: mock_registry, find_queue: nil) }
   let(:mock_pipeline) { instance_double(Minigun::Pipeline, name: 'test_pipeline', task: mock_task) }
 
   describe 'producer behavior' do
@@ -60,7 +60,7 @@
 
 RSpec.describe Minigun::ConsumerStage do
   let(:mock_registry) { instance_double(Minigun::StageRegistry, register: nil) }
-  let(:mock_task) { instance_double(Minigun::Task, stage_registry: mock_registry) }
+  let(:mock_task) { instance_double(Minigun::Task, stage_registry: mock_registry, find_queue: nil) }
   let(:mock_pipeline) { instance_double(Minigun::Pipeline, name: 'test_pipeline', task: mock_task) }
 
   describe 'processor behavior' do
@@ -117,7 +117,7 @@
 
 RSpec.describe Minigun::AccumulatorStage do
   let(:mock_registry) { instance_double(Minigun::StageRegistry, register: nil) }
-  let(:mock_task) { instance_double(Minigun::Task, stage_registry: mock_registry) }
+  let(:mock_task) { instance_double(Minigun::Task, stage_registry: mock_registry, find_queue: nil) }
   let(:mock_pipeline) { instance_double(Minigun::Pipeline, name: 'test_pipeline', task: mock_task) }
 
   it 'is a special batching stage' do
@@ -128,7 +128,7 @@
 
 RSpec.describe 'Stage common behavior' do
   let(:mock_registry) { instance_double(Minigun::StageRegistry, register: nil) }
-  let(:mock_task) { instance_double(Minigun::Task, stage_registry: mock_registry) }
+  let(:mock_task) { instance_double(Minigun::Task, stage_registry: mock_registry, find_queue: nil) }
   let(:mock_pipeline) { instance_double(Minigun::Pipeline, name: 'test_pipeline', task: mock_task) }
   let(:stage) { Minigun::ConsumerStage.new(:test, mock_pipeline, proc { |x, _output| x * 2 }, { foo: 'bar' }) }
 

spec/unit/execution/worker_spec.rb

@@ -3,12 +3,15 @@
 require 'spec_helper'
 
 RSpec.describe Minigun::Worker do
+  let(:stage_registry) { instance_double(Minigun::StageRegistry, register: nil) }
+  let(:task) { instance_double(Minigun::Task, find_queue: nil, stage_registry: stage_registry) }
+
   let(:pipeline) do
     instance_double(
       Minigun::Pipeline,
       name: 'test_pipeline',
       dag: dag,
-      stage_input_queues: {},
+      task: task,
       runtime_edges: {},
       context: user_context,
       stats: stats,
@@ -22,7 +25,8 @@
       name: :test_stage,
       execution_context: nil,
       log_type: 'Worker',
-      run_mode: :streaming
+      run_mode: :streaming,
+      task: task
     )
   end
 
@@ -55,8 +59,7 @@
   describe '#start' do
     it 'starts a worker thread' do
       input_queue = Queue.new
-      allow(pipeline).to receive(:stage_input_queues)
-        .and_return({ test_stage: input_queue })
+      allow(task).to receive(:find_queue).with(stage).and_return(input_queue)
       allow(dag).to receive(:upstream).with(:test_stage).and_return([:upstream])
       allow(dag).to receive(:downstream).with(:test_stage).and_return([])
 
@@ -89,8 +92,7 @@
   describe '#join' do
     it 'waits for worker thread to complete' do
       input_queue = Queue.new
-      allow(pipeline).to receive(:stage_input_queues)
-        .and_return({ test_stage: input_queue })
+      allow(task).to receive(:find_queue).with(stage).and_return(input_queue)
 
       worker = described_class.new(pipeline, stage, config)
 
@@ -151,10 +153,10 @@
 
   describe 'disconnected stage handling' do
     it 'exits early if no upstream sources' do
-      allow(pipeline).to receive(:stage_input_queues)
-        .and_return({ test_stage: Queue.new })
-      allow(dag).to receive(:upstream).with(:test_stage).and_return([])
-      allow(dag).to receive(:downstream).with(:test_stage).and_return([])
+      input_queue = Queue.new
+      allow(task).to receive(:find_queue).with(stage).and_return(input_queue)
+      allow(dag).to receive(:upstream).with(stage).and_return([])
+      allow(dag).to receive(:downstream).with(stage).and_return([])
 
       worker = described_class.new(pipeline, stage, config)
 
@@ -167,14 +169,14 @@
     end
 
     it 'sends END signals to downstream if disconnected' do
-      allow(pipeline).to receive(:stage_input_queues)
-        .and_return({ test_stage: Queue.new, downstream: Queue.new })
-      allow(dag).to receive(:upstream).with(:test_stage).and_return([])
-      allow(dag).to receive(:downstream).with(:test_stage).and_return([:downstream])
-
+      input_queue = Queue.new
+      downstream_stage = double('downstream_stage', name: :downstream, task: task)
       downstream_queue = Queue.new
-      allow(pipeline).to receive(:stage_input_queues)
-        .and_return({ test_stage: Queue.new, downstream: downstream_queue })
+
+      allow(task).to receive(:find_queue).with(stage).and_return(input_queue)
+      allow(task).to receive(:find_queue).with(downstream_stage).and_return(downstream_queue)
+      allow(dag).to receive(:upstream).with(stage).and_return([])
+      allow(dag).to receive(:downstream).with(stage).and_return([downstream_stage])
 
       worker = described_class.new(pipeline, stage, config)
       worker.start
@@ -213,8 +215,9 @@
       target_a_queue = Queue.new
       target_b_queue = Queue.new
 
-      allow(pipeline).to receive(:stage_input_queues)
-        .and_return({ broadcast_router => input_queue, target_a_stage => target_a_queue, target_b_stage => target_b_queue })
+      allow(task).to receive(:find_queue).with(broadcast_router).and_return(input_queue)
+      allow(task).to receive(:find_queue).with(target_a_stage).and_return(target_a_queue)
+      allow(task).to receive(:find_queue).with(target_b_stage).and_return(target_b_queue)
       allow(dag).to receive(:upstream).with(broadcast_router).and_return([source_stage])
 
       # Put items and END signal
@@ -261,8 +264,7 @@
   describe 'logging' do
     it 'logs when starting' do
       input_queue = Queue.new
-      allow(pipeline).to receive(:stage_input_queues)
-        .and_return({ test_stage: input_queue })
+      allow(task).to receive(:find_queue).with(stage).and_return(input_queue)
 
       input_queue << Minigun::EndOfSource.new(:upstream)
 
@@ -277,8 +279,7 @@
 
     it 'logs when done' do
       input_queue = Queue.new
-      allow(pipeline).to receive(:stage_input_queues)
-        .and_return({ test_stage: input_queue })
+      allow(task).to receive(:find_queue).with(stage).and_return(input_queue)
       allow(dag).to receive(:upstream).with(:test_stage).and_return([:upstream])
       allow(dag).to receive(:downstream).with(:test_stage).and_return([])
 

spec/unit/queue_wrappers_spec.rb

@@ -49,7 +49,6 @@
     described_class.new(
       test_stage,
       downstream_queues,
-      all_stage_queues,
       runtime_edges,
       stage_stats: stage_stats
     )