#!/usr/bin/env ruby
#  Phusion Passenger - http://www.modrails.com/
#  Copyright (C) 2008  Phusion
#
#  Phusion Passenger is a trademark of Hongli Lai & Ninh Bui.
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; version 2 of the License.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License along
#  with this program; if not, write to the Free Software Foundation, Inc.,
#  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 
$LOAD_PATH << File.expand_path(File.dirname(__FILE__) + "/../lib")
$LOAD_PATH << File.expand_path(File.dirname(__FILE__) + "/../ext")
require 'rubygems'
require 'optparse'
require 'socket'
require 'thread'
require 'passenger/platform_info'
require 'passenger/message_channel'
require 'passenger/utils'
 
include Passenger
include Passenger::Utils
include PlatformInfo
 
# A thread or a process, depending on the Ruby VM implementation.
class Subprocess
  attr_accessor :channel
  
  def initialize(name, &block)
    if RUBY_PLATFORM == "java"
      a, b = UNIXSocket.pair
      @thread = Thread.new do
        block.call(true, MessageChannel.new(b))
      end
      @channel = MessageChannel.new(a)
      @thread_channel = b
    else
      a, b = UNIXSocket.pair
      @pid = safe_fork(name) do
        a.close
        $0 = name
        Process.setsid
        block.call(false, MessageChannel.new(b))
      end
      b.close
      @channel = MessageChannel.new(a)
    end
  end
  
  def stop
    if RUBY_PLATFORM == "java"
      @thread.terminate
      @channel.close
      @thread_channel.close
    else
      Process.kill('SIGKILL', @pid) rescue nil
      Process.waitpid(@pid) rescue nil
      @channel.close
    end
  end
end
 
class StressTester
  def start
    @options = parse_options
    load_hawler
    
    Thread.abort_on_exception = true
    if GC.respond_to?(:copy_on_write_friendly=)
      GC.copy_on_write_friendly = true
    end
    @terminal_height = ENV['LINES'] ? ENV['LINES'].to_i : 24
    @terminal_width = ENV['COLUMNS'] ? ENV['COLUMNS'].to_i : 80
    
    if Process.euid != 0
      puts "*** WARNING: This program might not be able to restart " <<
        "Apache because it's not running as root. Please run " <<
        "this tool as root."
      puts
      puts "Press Enter to continue..."
      begin
        STDIN.readline
      rescue Interrupt
        exit 1
      end
    end
    
    run_crawlers
  end
 
  def parse_options
    options = {
      :concurrency => 20,
      :depth => 20,
      :nice => true,
      :apache_restart_interval => 24 * 60,
      :app_restart_interval => 55
    }
    parser = OptionParser.new do |opts|
      opts.banner = "Usage: passenger-stress-test <hostname> <app_root> [options]\n\n" <<
        "Stress test the given (Passenger-powered) website by:\n" <<
        "  * crawling it with multiple concurrently running crawlers.\n" <<
        "  * gracefully restarting Apache at random times (please point the 'APXS2'\n" <<
        "    variable to your Apache's 'apxs' binary).\n" <<
        "  * restarting the target (Passenger-powered) application at random time.\n" <<
        "\n" <<
        "Example:\n" <<
        "  passenger-stress-test mywebsite.com /webapps/mywebsite\n" <<
        "\n"
    
      opts.separator "Options:"
      opts.on("-c", "--concurrency N", Integer,
        "Number of crawlers to start (default = #{options[:concurrency]})") do |v|
        options[:concurrency] = v
      end
      opts.on("-p", "--apache-restart-interval N", Integer,
        "Gracefully restart Apache after N minutes\n" <<
        (" " * 37) << "(default = #{options[:apache_restart_interval]})") do |v|
        options[:apache_restart_interval] = v
      end
      opts.on("-a", "--app-restart-interval N", Integer,
        "Restart the application after N minutes\n" <<
        (" " * 37) << "(default = #{options[:app_restart_interval]})") do |v|
        options[:app_restart_interval] = v
      end
      opts.on("-h", "--help", "Show this message") do
        puts opts
        exit
      end
    end
    parser.parse!
  
    options[:host] = ARGV[0]
    options[:app_root] = ARGV[1]
    if !options[:host] || !options[:app_root]
      puts parser
      exit 1
    end
    return options
  end
 
  def load_hawler
    begin
      require 'hawler'
    rescue LoadError
      STDERR.puts "This tool requires Hawler (http://tinyurl.com/ywgk6x). Please install it with:"
      STDERR.puts
      STDERR.puts "  gem install --source http://spoofed.org/files/hawler/ hawler"
      exit 1
    end
  end
 
  def run_crawlers
    @started = false
    @crawlers = []
    
    # Start crawler processes.
    GC.start if GC.copy_on_write_friendly?
    @options[:concurrency].times do |i|
      STDOUT.write("Starting crawler #{i + 1} of #{@options[:concurrency]}...\n")
      STDOUT.flush
      process = Subprocess.new("crawler #{i + 1}") do |is_thread, channel|
        if !is_thread && @options[:nice]
          system("renice 1 #{Process.pid} >/dev/null 2>/dev/null")
        end
        while true
          crawl!(i + 1, channel)
        end
      end
      @crawlers << {
        :id => i + 1,
        :process => process,
        :channel => process.channel,
        :mutex => Mutex.new,
        :current_uri => nil,
        :crawled => 0
      }
    end
    
    puts
    if RUBY_PLATFORM != "java"
      # 'sleep' b0rks when running in JRuby?
      sleep 1
    end
    begin
      $0 = "Passenger Crawler: control process"
      io_to_crawler = {}
      ios = []
      @crawlers.each do |crawler|
        io_to_crawler[crawler[:channel].io] = crawler
        ios << crawler[:channel].io
      end
      
      # Tell each crawler to start crawling.
      @crawlers.each do |crawler|
        crawler[:channel].write("start")
      end
      
      # Show progress periodically.
      @start_time = Time.now
      progress_reporter = Thread.new(&method(:report_progress))
      @next_apache_restart = Time.now + @options[:apache_restart_interval] * 60
      apache_restarter = Thread.new(&method(:restart_apache))
      @next_app_restart = Time.now + @options[:app_restart_interval] * 60
      app_restarter = Thread.new(&method(:restart_app))
      
      while true
        note_progress(ios, io_to_crawler)
      end
    rescue Interrupt
      trap('SIGINT') {}
      puts "Shutting down..."
      @done = true
      @crawlers.each do |crawler|
        STDOUT.write("Stopping crawler #{crawler[:id]} of #{@options[:concurrency]}...\r")
        STDOUT.flush
        crawler[:process].stop
      end
      progress_reporter.join if progress_reporter
      apache_restarter.join if apache_restarter
      app_restarter.join if app_restarter
      puts
    end
  end
  
  def note_progress(ios, io_to_crawler)
    select(ios)[0].each do |io|
      crawler = io_to_crawler[io]
      uri = crawler[:channel].read[0]
      crawler[:mutex].synchronize do
        crawler[:current_uri] = uri
        crawler[:crawled] += 1
      end
    end
  end
  
  def report_progress
    while !@done
      output = "\n" * @terminal_height
      output << "### Running for #{duration(Time.now.to_i - @start_time.to_i)}\n"
      @crawlers.each do |crawler|
        crawler[:mutex].synchronize do
          line = sprintf("Crawler %-2d: %-3d -> %s",
            crawler[:id],
            crawler[:crawled],
            crawler[:current_uri])
          output << sprintf("%-#{@terminal_width}s\n", line)
        end
      end
      output << "Next Apache restart: in #{duration(@next_apache_restart.to_i - Time.now.to_i)}\n"
      output << "Next app restart   : in #{duration(@next_app_restart.to_i - Time.now.to_i)}\n"
      STDOUT.write(output)
      sleep 0.5
    end
  end
  
  def restart_apache
    while !@done
      if Time.now > @next_apache_restart
        @next_apache_restart = Time.now + @options[:apache_restart_interval] * 60
        system("#{HTTPD} -k graceful")
      end
    end
  end
  
  def restart_app
    while !@done
      if Time.now > @next_app_restart
        @next_app_restart = Time.now + @options[:app_restart_interval] * 60
        system("touch #{@options[:app_root]}/tmp/restart.txt")
      end
    end
  end
  
  def duration(seconds)
    result = ""
    if seconds >= 60
      minutes = (seconds / 60)
      if minutes >= 60
        hours = minutes / 60
        minutes = minutes % 60
        if hours == 1
          result << "#{hours} hour "
        else
          result << "#{hours} hours "
        end
      end
      
      seconds = seconds % 60
      if minutes == 1
        result << "#{minutes} minute "
      else
        result << "#{minutes} minutes "
      end
    end
    result << "#{seconds} seconds"
    return result
  end
 
  def crawl!(id, channel)
    progress_reporter = lambda do |uri, referer, response|
      begin
        if !@started
          # At the beginning, wait until the control process
          # tells us to start.
          @started = true
          channel.read
        end
        channel.write(uri, referer, response)
      rescue
        if RUBY_PLATFORM == "java"
          Thread.current.terminate
        else
          Process.kill('SIGKILL', Process.pid)
        end
      end
    end
    crawler = Hawler.new(@options[:host], progress_reporter)
    if RUBY_PLATFORM == "java"
      trap('SIGINT') do
        raise Interrupt, "Interrupted"
      end
    end
    crawler.recurse = true
    crawler.depth = @options[:depth]
    crawler.start
  end
end
 
StressTester.new.start