trident 0.1.0

data/lib/trident/signal_handler.rb

@@ -0,0 +1,167 @@
+module Trident
+  class SignalHandler
+    include GemLogger::LoggerSupport
+
+    CHUNK_SIZE = (16 * 1024)
+    SIGNAL_QUEUE_MAX_SIZE = 5
+    MSG_STOP = 'STOP'
+
+    class << self
+
+      attr_accessor :instance
+
+      def start(signal_mappings, target)
+        raise "Already started, call stop if restart needed" if instance
+        logger.info "Starting signal handler"
+        self.instance = new(signal_mappings, target)
+        instance.start
+      end
+
+      def stop
+        raise "No signal handler started" unless instance
+        logger.info "Stopping signal handler"
+        instance.stop
+        self.instance = nil
+      end
+
+      def reset_for_fork
+        raise "No signal handler started" unless instance
+        instance.reset_for_fork
+        self.instance = nil
+      end
+
+    end
+
+    attr_reader :target, :signal_mappings, :signal_queue, :self_pipe, :original_signal_handlers
+
+    def initialize(signal_mappings, target)
+      @target = target
+      @signal_queue = []
+      @self_pipe = []
+      @original_signal_handlers = {}
+      self.signal_mappings = signal_mappings
+    end
+
+    def start
+      setup_self_pipe
+      setup_signal_handlers
+
+      logger.info "Main loop started"
+      loop do
+        signal_result = handle_signal_queue
+        break if signal_result == :break
+        msg = snooze if signal_queue.empty?
+        logger.debug "Main loop awake"
+        break if msg == MSG_STOP
+      end
+      logger.info "Main loop exited"
+    end
+
+    def stop
+      reset_signal_handlers
+      wakeup(MSG_STOP)
+    end
+
+    def reset_for_fork
+      @self_pipe = []
+      reset_signal_handlers
+    end
+
+    def wakeup(msg='.')
+      begin
+        # mutexes (and thus logging) not allowed within a trap context
+        # puts "Waking main loop"
+        self_pipe.last.write_nonblock(msg) # wakeup master process from select
+      rescue Errno::EAGAIN, Errno::EINTR
+        # pipe is full, master should wake up anyways
+        retry
+      end
+    end
+
+    def snooze
+      msg = ""
+      begin
+        logger.info "Snoozing main loop"
+        ready = IO.select([self_pipe.first], nil, nil, 1) or return
+        ready.first && ready.first.first or return
+        loop { msg << self_pipe.first.read_nonblock(CHUNK_SIZE) }
+      rescue Errno::EAGAIN, Errno::EINTR
+      end
+      msg
+    end
+
+    private
+
+    def signal_mappings=(mappings)
+      @signal_mappings = {}
+      mappings.each do |k, v|
+        k = "SIG#{k}" unless k =~ /^SIG/i
+        k = k.upcase
+
+        raise ArgumentError,
+              "Duplicate signal handler: #{k}" if @signal_mappings.has_key?(k)
+
+        @signal_mappings[k] = Array(v)
+      end
+
+      # Should always handle CHLD signals as they wakeup/drive the main
+      # loop on status changes from child processes
+      @signal_mappings = {"SIGCHLD" => ["update"]}.merge(@signal_mappings)
+    end
+
+    def setup_self_pipe
+      self_pipe.each { |io| io.close rescue nil }
+      self_pipe.replace(IO.pipe)
+      self_pipe.each { |io| io.fcntl(Fcntl::F_SETFD, Fcntl::FD_CLOEXEC) }
+    end
+
+    def setup_signal_handlers
+      logger.info "Installing signal handlers"
+      signal_mappings.each do |signal_name, actions|
+        raise ArgumentError,
+              "Target does not respond to action: #{actions}" unless actions.all? { |a| target.respond_to?(a) }
+
+        logger.info "Adding signal mapping: #{signal_name} -> #{actions.inspect}"
+        original_signal_handlers[signal_name] = trap_deferred(signal_name)
+      end
+    end
+
+    def reset_signal_handlers
+      original_signal_handlers.each do |signal_name, original_handler|
+        trap(signal_name, original_handler)
+      end
+      original_signal_handlers.clear
+    end
+
+    # defer a signal for later processing in #join (master process)
+    def trap_deferred(signal)
+      trap(signal) do |signal_number|
+        if signal_queue.size < SIGNAL_QUEUE_MAX_SIZE
+          # mutexes (and thus logging) not allowed within a trap context
+          # puts "Adding signal to queue: #{signal}"
+          signal_queue << signal
+          wakeup
+        else
+          $stderr.puts "Signal queue exceeded max size, ignoring #{signal}"
+        end
+      end
+    end
+
+    def handle_signal_queue
+      signal_result = nil
+      signal = signal_queue.shift
+      if signal
+        logger.info "Handling signal: #{signal}"
+        actions = signal_mappings[signal]
+        if actions
+          actions.each do |action|
+            logger.info "Sending to target: #{action}"
+            signal_result = target.send(action)
+          end
+        end
+      end
+      signal_result
+    end
+
+  end
+end

data/lib/trident/utils.rb

@@ -0,0 +1,9 @@
+module Trident
+  module Utils
+
+    def procline(id, string)
+      $0 = "trident[#{id}]: #{string}"
+    end
+
+  end
+end

data/lib/trident/version.rb

@@ -0,0 +1,3 @@
+module Trident
+  VERSION = "0.1.0"
+end

data/test/fixtures/integration_project/config/trident.yml

@@ -0,0 +1,51 @@
+# a name for labelling pool processes
+application: app_name
+
+# loads handlers in the master process before forking workers
+prefork: true
+
+# Configure the behavior of the pool manager in response to signals
+# actions are methods on PoolManager and can be sequenced
+signals:
+  INT: stop_gracefully
+  TERM: stop_forcefully
+  USR1: [stop_forcefully, wait]
+
+# define and setup all the PoolHandlers
+handlers:
+  myhandler:
+    environment: |+
+      class MyWorker
+        def initialize(o)
+          @o = o
+        end
+        def start
+          loop { sleep 1 }
+        end
+      end
+    class: MyWorker
+    options:
+    signals:
+      default: USR2
+      stop_gracefully: INT
+      stop_forcefully: TERM
+
+# Configure the worker pools
+pools:
+  mypool1:
+    # number of worker processes
+    size: 3
+    # chooses a handler defined above for running process
+    handler: myhandler
+    # options passed to each handler's initializer (merged into handler options above)
+    options:
+      name: one
+  mypool2:
+    # number of worker processes
+    size: 2
+    # chooses a handler defined above for running process
+    handler: myhandler
+    # options passed to each handler's initializer (merged into handler options above)
+    options:
+      name: two
+

data/test/integration/trident_test.rb

@@ -0,0 +1,105 @@
+require_relative '../test_helper'
+
+class Trident::TridentTest < MiniTest::Should::TestCase
+
+  setup do
+    @project_root = File.expand_path('../../fixtures/integration_project', __FILE__)
+    @cli = "#{File.expand_path('../../..', __FILE__)}/bin/trident"
+  end
+
+  def wait_for(io, pattern)
+    timeout(5) do
+      loop do
+        line = io.readline
+        puts line if ENV['DEBUG']
+        break if line =~ pattern
+      end
+    end
+  end
+
+  def process_list
+    processes = {}
+    lines = `ps -e -opid,command`.lines.grep(/trident\[/)
+    lines.each do |line|
+      pieces = line.split
+      pid = pieces[0].to_i
+      next if pid == Process.pid
+      command = pieces[1..-1].join(' ')
+      processes[pid] = command
+    end
+    processes
+  end
+
+  def parse_manager(manager_str)
+    pools = {}
+    manager_str.scan(/(\w+)\[([0-9, ]+)\]/) do |pool, pids|
+      pids = pids.split(", ").collect(&:to_i)
+      pools[pool] = pids
+    end
+    pools
+  end
+
+  context "basic usage" do
+
+    should "start and stop pools" do
+      cmd = "#{@cli} --verbose --config #{@project_root}/config/trident.yml"
+      io = IO.popen(cmd, :err=>[:child, :out])
+
+      wait_for(io, /<pool-mypool1> Pool started with 3 workers/)
+      wait_for(io, /<pool-mypool2> Pool started with 2 workers/)
+
+      processes = process_list
+      assert_equal 6, processes.size
+      manager = processes[io.pid]
+      pools = parse_manager(manager)
+      pools.each do |pool, pids|
+        pids.each do |pid|
+          assert processes[pid], "no worker process"
+          assert_match /trident[pool-#{pool}-worker]/, processes[pid], "worker process not in right pool"
+        end
+      end
+
+      Process.kill("USR1", io.pid)
+
+      wait_for(io, /<pool-mypool1> Pool stopped/)
+      wait_for(io, /<pool-mypool2> Pool stopped/)
+      wait_for(io, /Main loop exited/)
+
+      Process.wait(io.pid)
+      assert_empty process_list
+    end
+
+  end
+
+  context "worker maintenance" do
+
+    should "restart failed workers" do
+      cmd = "#{@cli} --verbose --config #{@project_root}/config/trident.yml"
+      io = IO.popen(cmd, :err=>[:child, :out])
+
+      wait_for(io, /<pool-mypool1> Pool started with 3 workers/)
+      wait_for(io, /<pool-mypool2> Pool started with 2 workers/)
+
+      processes = process_list
+      assert_equal 6, processes.size
+      manager = processes[io.pid]
+      pools = parse_manager(manager)
+      children = pools['mypool1']
+      child = children.delete_at(1)
+      Process.kill("KILL", child)
+
+      wait_for(io, /<pool-mypool1> Spawned worker \d+, worker count now at 3/)
+      processes = process_list
+      assert_equal 6, processes.size
+      manager = processes[io.pid]
+      pools = parse_manager(manager)
+      assert_equal 3, pools['mypool1'].size
+      assert children.all? {|c| pools['mypool1'].include?(c) }
+
+      Process.kill("USR1", io.pid)
+      Process.wait(io.pid)
+      assert_empty process_list
+    end
+
+  end
+end

data/test/test_helper.rb

@@ -0,0 +1,144 @@
+require 'rubygems'
+
+if ENV['CI']
+  require 'coveralls'
+  Coveralls.wear!
+end
+
+require 'bundler'
+begin
+  Bundler.setup(:default, :development)
+rescue Bundler::BundlerError => e
+  $stderr.puts e.message
+  $stderr.puts "Run `bundle install` to install missing gems"
+  exit e.status_code
+end
+
+require 'minitest/autorun'
+require 'minitest/should'
+require "minitest/reporters"
+require "mocha/setup"
+require 'timeout'
+require 'tempfile'
+
+reporter = ENV['REPORTER']
+reporter = case reporter
+  when 'none' then nil
+  when 'spec' then MiniTest::Reporters::SpecReporter.new
+  when 'progress' then MiniTest::Reporters::ProgressReporter.new
+  else MiniTest::Reporters::DefaultReporter.new
+end
+MiniTest::Reporters.use!(reporter) if reporter
+
+require 'trident'
+include Trident
+
+GemLogger.default_logger = Logger.new("/dev/null")
+
+class ForkChild
+
+  attr_reader :pid, :read_from_child
+
+  def initialize
+    @read_from_child, @write_from_child = IO.pipe
+
+    @pid = fork do
+      @read_from_child.close
+      result = yield
+      Marshal.dump(result, @write_from_child)
+      exit!(0) # skips exit handlers.
+    end
+
+    @write_from_child.close
+  end
+
+  def wait(time=5)
+    timeout(time) do
+      result = @read_from_child.read
+      begin
+        Process.wait(@pid)
+      rescue Errno::ECHILD
+      end
+      raise "child failed" if result.empty?
+      Marshal.load(result)
+    end
+  end
+
+end
+
+class FileCounter
+
+  def initialize(file=Tempfile.new('file_counter').path)
+    @file = file
+  end
+
+  def increment
+    File.open(@file, File::RDWR|File::CREAT, 0644)  do |f|
+      f.flock(File::LOCK_EX)
+      value = f.read.to_i + 1
+      f.rewind
+      f.write("#{value}\n")
+      f.flush
+      f.truncate(f.pos)
+    end
+  end
+
+  def read
+    # read the counter using read lock
+    File.open(@file, "r")  do |f|
+      f.flock(File::LOCK_SH)
+      f.read.to_i
+    end
+  end
+
+end
+
+class MiniTest::Should::TestCase
+  ORIGINAL_PROCLINE = $0
+
+  setup do
+    $0 = ORIGINAL_PROCLINE
+  end
+end
+
+# Allow triggering single tests when running from rubymine
+# reopen the installed runner so we don't step on runner customizations
+class << MiniTest::Unit.runner
+  # Rubymine sends --name=/\Atest\: <context> should <should>\./
+  # Minitest runs each context as a suite
+  # Minitest filters methods by matching against: <suite>#test_0001_<should>
+  # Nested contexts are separted by spaces in rubymine, but ::s in minitest
+  
+  def _run_suites(suites, type)
+    if options[:filter]
+      if options[:filter] =~ /\/\\Atest\\: (.*) should (.*)\\\.\//
+        context_filter = $1
+        should_filter = $2
+        should_filter.strip!
+        should_filter.gsub!(" ", "_")
+        should_filter.gsub!(/\W/, "")
+        context_filter = context_filter.gsub(" ", "((::)| )")
+        options[:filter] = "/\\A#{context_filter}(Test)?#test(_\\d+)?_should_#{should_filter}\\Z/"
+      end
+    end
+    
+    super
+  end
+  
+  # Prevent "Empty test suite" verbosity when running in rubymine
+  def _run_suite(suite, type)
+    
+    filter = options[:filter] || '/./'
+    filter = Regexp.new $1 if filter =~ /\/(.*)\//    
+    all_test_methods = suite.send "#{type}_methods"
+    filtered_test_methods = all_test_methods.find_all { |m|
+      filter === m || filter === "#{suite}##{m}"
+    }
+    
+    if filtered_test_methods.size > 0    
+      super
+    else
+      [0, 0]
+    end
+  end
+end