class Forkomatic

# File lib/forkomatic.rb, line 21
def initialize(args)
  params = {}
  if args.is_a?(String)
    # Load config from a file.
    params = load_config(args)
  elsif args.is_a?(Integer)
    # Given an integer, forkomatic will only run N runners 1 time.
    params[:wait_for_children] = true
    params[:max_children] = args
    params[:work_interval] = 0
    params[:max_iterations] = 1
  elsif args.is_a?(Hash)
    # Specify the parameters directly.
    params = args
  end
  t = params
  params.inject({}) {|t, (key, val)| t[key.to_sym] = val; t}
  self.jobs = []
  self.max_children = params[:max_children] || 1
  self.work_interval = params[:work_interval].nil? || params[:work_interval] < 0 ? 0 : params[:work_interval]
  self.max_iterations = params[:max_iterations]
  self.wait_for_children = params[:wait_for_children].nil? ? true : params[:wait_for_children]
end

# File lib/forkomatic.rb, line 130
def available
  # Initialize if need be.
  return @max_children if @jobs.nil? || @jobs.empty?
  # Reap children runners without waiting.
  reap_all
  @max_children - @jobs.length
end

# File lib/forkomatic.rb, line 101
def build_jobs(count)
  (1..count).each.collect {Forkomatic::Job.new}
end

# File lib/forkomatic.rb, line 139
def child_pids
  reap_all
  pids = []
  @jobs.each {|job| pids.push(job.pid) if job.pid}
  pids
end

# File lib/forkomatic.rb, line 46
def load_config(config_file)
  params = {}
  # Allowed options.
  options = ['max_children', 'work_interval', 'max_iterations', 'wait_for_children']
  begin
    # Try to read the config file, and store the values.
    data = File.open(config_file, "r").read.split(/\n/)
    data.each do |line|
      if line =~ /^\s*([a-zA-Z_]+)\s+([0-9]+)/
        config_item = $1
        config_value = $2
        # Make sure option is valid.
        if options.map(&:downcase).include?(config_item.downcase)
          params[config_item.to_sym] = config_value.to_i
        end
      end
    end
  rescue => e
    puts "Error loading config file: #{e.to_s}"
  end
  params
end

# File lib/forkomatic.rb, line 106
def reap(pid)
  return true if pid.nil?
  begin
    return true if Process.waitpid(pid, Process::WNOHANG)
  rescue Errno::ECHILD
    return true
  rescue => e
    puts "ERROR: #{e.to_s}"
  end
  return false
end

# File lib/forkomatic.rb, line 119
def reap_all
  finished = []
  @jobs.each do |job|
    if reap(job.pid)
      finished.push(job.pid)
    end
  end
  @jobs.delete_if {|job| finished.include?(job.pid)}
end

# File lib/forkomatic.rb, line 82
def run
  Signal.trap("INT")  { shutdown(true) }
  iteration = 0
  while (@max_iterations.nil? || iteration < @max_iterations) do
    iteration += 1
    current_jobs = build_jobs(available)
    current_jobs.each do |job|
      pid = Process.fork do
        job.work!
      end
      job.pid = pid
      @jobs.push(job)
    end
    sleep @work_interval if @work_interval > 0
  end
  Process.waitall if @wait_for_children
end

# File lib/forkomatic.rb, line 70
def shutdown(all)
  child_pids.each do |pid|
    begin
      Process.kill("TERM", pid)
    rescue => e
      puts e.to_s
    end
  end
  exit if all
end