huginn/lib/agent_runner.rb

require 'cgi'
require 'json'
require 'rufus-scheduler'
require 'pp'
require 'twitter'

class AgentRunner
  @@agents = []

  def initialize(options = {})
    @workers = {}
    @signal_queue = []
    @options = options
    @options[:only] = [@options[:only]].flatten if @options[:only]
    @options[:except] = [@options[:except]].flatten if @options[:except]
    @mutex = Mutex.new
    @scheduler = Rufus::Scheduler.new(frequency: ENV['SCHEDULER_FREQUENCY'].presence || 0.3)

    @scheduler.every 5 do
      restart_dead_workers if @running
    end

    @scheduler.every 60 do
      run_workers if @running
    end

    set_traps
  end

  def stop
    puts "Stopping AgentRunner..." unless Rails.env.test?
    @running = false
    @workers.each_pair do |_, w| w.stop! end
    @scheduler.stop
  end

  def run
    @running = true
    run_workers

    while @running
      if signal = @signal_queue.shift
        handle_signal(signal)
      end
      sleep 0.25
    end
    @scheduler.join
  end

  def set_traps
    %w(INT TERM QUIT).each do |signal|
      Signal.trap(signal) { @signal_queue << signal }
    end
  end

  def self.register(agent)
    @@agents << agent unless @@agents.include?(agent)
  end

  def self.with_connection
    ActiveRecord::Base.connection_pool.with_connection do
      yield
    end
  end

  private

  def run_workers
    workers             = load_workers
    new_worker_ids      = workers.keys
    current_worker_ids  = @workers.keys

    (current_worker_ids - new_worker_ids).each do |outdated_worker_id|
      puts "Killing #{outdated_worker_id}" unless Rails.env.test?
      @workers[outdated_worker_id].stop!
      @workers.delete(outdated_worker_id)
    end

    (new_worker_ids - current_worker_ids).each do |new_worker_id|
      puts "Starting #{new_worker_id}" unless Rails.env.test?
      @workers[new_worker_id] = workers[new_worker_id]
      @workers[new_worker_id].setup!(@scheduler, @mutex)
      @workers[new_worker_id].run!
    end
  end

  def load_workers
    workers = {}
    @@agents.each do |klass|
      next if @options[:only] && !@options[:only].include?(klass)
      next if @options[:except] && @options[:except].include?(klass)

      AgentRunner.with_connection do
        (klass.setup_worker || [])
      end.each do |agent_worker|
        workers[agent_worker.id] = agent_worker
      end
    end
    workers
  end

  def restart_dead_workers
    @workers.each_pair do |id, worker|
      if !worker.restarting && worker.thread && !worker.thread.alive?
        puts "Restarting #{id.to_s}" unless Rails.env.test?
        @workers[id].run!
      end
    end
  end

  def handle_signal(signal)
    case signal
    when 'INT', 'TERM', 'QUIT'
      stop
    end
  end
end

require 'agents/twitter_stream_agent'
require 'agents/jabber_agent'
require 'agents/local_file_agent'
require 'huginn_scheduler'
require 'delayed_job_worker'