Lambda in main thread and Enumerator support

lib/parallel.rb

@@ -97,20 +97,26 @@ def wait
 
   class JobFactory
     def initialize(source, mutex)
-      @lambda = (source.respond_to?(:call) && source) || queue_wrapper(source)
-      @source = source.to_a unless @lambda # turn Range and other Enumerable-s into an Array
+      @lambda = enum_wrapper(source) || (source.respond_to?(:call) && source) || queue_wrapper(source)
+      @source = source.to_a unless @lambda # turn non-Enumerable-s into an Array
+      @runloop_queue = Thread::Queue.new if @lambda
       @mutex = mutex
       @index = -1
       @stopped = false
     end
 
-    def next
-      if producer?
+    def next(queue_for_thread = nil)
+      if @runloop_queue && queue_for_thread
+        return if @stopped
+        item = runloop_enq(queue_for_thread)
+        return if item == Stop
+        index = @index += 1
+      elsif producer?
         # - index and item stay in sync
         # - do not call lambda after it has returned Stop
         item, index = @mutex.synchronize do
           return if @stopped
-          item = @lambda.call
+          item = call_lambda
           @stopped = (item == Stop)
           return if @stopped
           [item, @index += 1]
@@ -123,6 +129,25 @@ def next
       [item, index]
     end
 
+    def runloop
+      return unless @runloop_queue
+
+      loop do
+        queue = @runloop_queue.pop
+        return if queue == Stop
+        item = call_lambda
+        queue.push(item)
+        break if item == Stop
+      end
+      @stopped = true
+      begin
+        while queue = @runloop_queue.pop(true)
+          queue.push(Stop) if queue != Stop # Unlock waiting threads.
+        end
+      rescue ThreadError # All threads are unlocked.
+      end
+    end
+
     def size
       if producer?
         Float::INFINITY
@@ -142,15 +167,33 @@ def unpack(data)
       producer? ? data : [@source[data], data]
     end
 
+    def stopper = @runloop_queue&.push(Stop)
+
     private
 
+    def call_lambda
+      @lambda.call
+    rescue StopIteration
+      Stop
+    end
+
+    def runloop_enq(queue_for_thread)
+      @runloop_queue.push(queue_for_thread)
+      queue_for_thread.pop # Wait until @lambda returns.
+    end
+
     def producer?
       @lambda
     end
 
     def queue_wrapper(array)
       array.respond_to?(:num_waiting) && array.respond_to?(:pop) && -> { array.pop(false) }
     end
+
+    def enum_wrapper(source)
+      # Convert what is inaccessible by the index
+      !source.respond_to?(:[]) && source.respond_to?(:next) && source.method(:next)
+    end
   end
 
   class UserInterruptHandler
@@ -211,13 +254,24 @@ def restore_interrupt(old, signal)
   class << self
     def in_threads(options = { count: 2 })
       threads = []
-      count, = extract_count_from_options(options)
+      count, options = extract_count_from_options(options)
+      finished_monitor = options[:runloop] && Queue.new(1..(count - 1)) # Insert values, one less in count than the number of threads.
+      stopper = options[:stopper]
 
       Thread.handle_interrupt(Exception => :never) do
         Thread.handle_interrupt(Exception => :immediate) do
           count.times do |i|
-            threads << Thread.new { yield(i) }
+            threads << Thread.new do
+              yield(i)
+            ensure
+              begin
+                finished_monitor&.pop(true) # This must be executed even if the worker thread is killed (by #work_in_processes).
+              rescue ThreadError # Queue#pop raises ThreadError when the queue is empty.
+                stopper&.call # Stop JobFactory#runloop
+              end
+            end
           end
+          options[:runloop]&.call # Invoke lambda in caller thread, and provide jobs to thread queue.
           threads.map(&:value)
         end
       ensure
@@ -431,10 +485,11 @@ def work_in_threads(job_factory, options, &block)
       results_mutex = Mutex.new # arrays are not thread-safe on jRuby
       exception = nil
 
-      in_threads(options) do |worker_num|
+      in_threads(options.merge(runloop: job_factory.method(:runloop), stopper: job_factory.method(:stopper))) do |worker_num|
+        queue_for_thread = Thread::Queue.new
         self.worker_number = worker_num
         # as long as there are more jobs, work on one of them
-        while !exception && (set = job_factory.next)
+        while !exception && (set = job_factory.next(queue_for_thread))
           begin
             item, index = set
             result = with_instrumentation item, index, options do
@@ -523,15 +578,16 @@ def work_in_processes(job_factory, options, &blk)
       exception = nil
 
       UserInterruptHandler.kill_on_ctrl_c(workers.map(&:pid), options) do
-        in_threads(options) do |i|
+        in_threads(options.merge(runloop: job_factory.method(:runloop), stopper: job_factory.method(:stopper))) do |i|
           worker = workers[i]
           worker.thread = Thread.current
+          queue_for_thread = Thread::Queue.new
           worked = false
 
           begin
             loop do
               break if exception
-              item, index = job_factory.next
+              item, index = job_factory.next(queue_for_thread)
               break unless index
 
               if options[:isolation]

spec/cases/infinite_sequence.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+# Reproduction case based on GitHub Issue #211
+# Original code provided by @cyclotron3k in the issue
+
+require 'prime'
+require './spec/cases/helper'
+
+private_key = 12344567899
+
+results = []
+
+[{ in_threads: 2 }, { in_threads: 0 }].each do |options|
+  primes = Prime.to_enum
+  Parallel.each(primes, options) do |prime|
+    if private_key % prime == 0
+      results << prime.to_s
+      raise Parallel::Break
+    end
+  end
+end
+
+print results.join(',')

spec/cases/lambda_call_same_thread.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+require './spec/cases/helper'
+
+runner_thread = nil
+all = [3, 2, 1]
+my_proc = proc {
+  runner_thread ||= Thread.current
+  if Thread.current != runner_thread
+    raise "proc is called in different thread!"
+  end
+
+  all.pop || Parallel::Stop
+}
+
+class Callback
+  def self.call(x)
+    $stdout.sync = true
+    "ITEM-#{x}"
+  end
+end
+puts(Parallel.map(my_proc, in_threads: 2) { |(i, _id)| Callback.call i })

spec/cases/lambda_can_stop_by_exception.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+require './spec/cases/helper'
+
+def generate_proc
+  count = 0
+  proc {
+    raise StopIteration if 3 <= count
+    count += 1
+  }
+end
+
+class Callback
+  def self.call(x)
+    $stdout.sync = true
+    "ITEM-#{x}"
+  end
+end
+
+[{ in_processes: 2 }, { in_threads: 2 }, { in_threads: 0 }].each do |options|
+  puts(Parallel.map(generate_proc, options) { |(i, _id)| Callback.call i })
+end

spec/parallel_spec.rb

@@ -722,6 +722,18 @@ def cpus
     end
   end
 
+  it "can process infinite sequence enumerator" do
+    ruby("spec/cases/infinite_sequence.rb").split(',').should == ['139'] * 2
+  end
+
+  it "can be finished by lambda raising StopIteration" do
+    ruby("spec/cases/lambda_can_stop_by_exception.rb").should == "ITEM-1\nITEM-2\nITEM-3\n" * 3
+  end
+
+  it "must call lambda in same thread" do
+    ruby("spec/cases/lambda_call_same_thread.rb").should == "ITEM-1\nITEM-2\nITEM-3\n"
+  end
+
   it "fails when running with a prefilled queue without stop since there are no threads to fill it" do
     error = (RUBY_VERSION >= "2.0.0" ? "No live threads left. Deadlock?" : "deadlock detected (fatal)")
     ruby("spec/cases/fatal_queue.rb 2>&1").should include error