agent_runner.rb 2.8KB

    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'