#!/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 'phusion_passenger/platform_info'
require 'phusion_passenger/message_channel'
require 'phusion_passenger/utils'

include PhusionPassenger
include PhusionPassenger::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