#!/usr/bin/env ruby
# RailsFork - Saving memory in Ruby on Rails by preforking.
#
# Copyright (c) 2007, Hongli Lai
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this list
#   of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above copyright notice, this list
#   of conditions and the following disclaimer in the documentation and/or other
#   materials provided with the distribution.
# * Neither the name of the Hongli nor the names of its contributors may be used to
#   endorse or promote products derived from this software without specific prior
#   written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
require 'optparse'
require 'ostruct'
require 'socket'

class Object
	def self._rails_fork_remove_constants(*syms)
		syms.each do |sym|
			if constants.include?(sym.to_s)
				remove_const(sym)
			end
		end
	end
end

# The main program class.
class RailsFork
	def start
		options = parse_arguments
		init_environment
		puts "Spawning #{options.count} Rails processes..."
		preload_rails
		fork_children(options)
		exit!
	end

private
	# Parse arguments. Exits if arguments are wrong.
	def parse_arguments
		options = OpenStruct.new
		options.socket_dir = "log"
		options.prefix = "fastcgi.socket-"

		opts = OptionParser.new(nil, 20, '  ') do |o|
			o.banner = "Usage: railsfork [options] <number of processes>"
			o.separator ""
			o.separator "Options:"

			o.on("-d", "--dir DIR", "Directory in which to put the socket files.", "The default is '#{options.socket_dir}'") do |value|
				if !File.directory?(value)
					raise StandardError.new("Directory '#{value}' does not exist.")
				end
				options.socket_dir = value
			end
			o.on("-p", "--prefix NAME", "A prefix for the socket files.", "The default is '#{options.prefix}'") do |value|
				options.prefix = value
			end
		end
		begin
			opts.parse!(ARGV)
			if ARGV.size < 1
				raise OptionParser::MissingArgument.new("number of processes")
			else
				options.count = ARGV[0].to_i
			end
			return options
		rescue Exception => e
			puts "*** ERROR: " + e
			puts
			puts opts.help
			exit 1
		end
	end

	# Make sure variables like RAILS_ROOT are correct.
	def init_environment
		if !ENV.include?('RAILS_ROOT')
			ENV['RAILS_ROOT'] = Dir.pwd
		end
		@dispatch_script = "#{ENV['RAILS_ROOT']}/public/dispatch.fcgi"
		if !File.file?(@dispatch_script)
			puts "*** ERROR: The file '#{@dispatch_script}' doesn't exist."
			puts
			print "RailsFork thinks that '#{ENV['RAILS_ROOT']}' is your Ruby on Rails root folder. "
			print "If that's incorrect, please set the RAILS_ROOT environment variable to your Ruby "
			print "on Rails application's root folder.\n"
			exit 1
		end
	end

	# Preload the Ruby on Rails libraries.
	def preload_rails
		require 'config/environment'
		require 'fcgi_handler'
		require 'ruby_version_check'
		require 'initializer'
		require 'dispatcher'

		Rails::Initializer.run
		require 'breakpoint' if defined?(BREAKPOINT_SERVER_PORT)
		require_dependency('application.rb')

		# Remove some constants so we don't get an annoying warning.
		Object._rails_fork_remove_constants(:RAILS_GEM_VERSION, :MAX_SESSION_TIME)
	end

	# Fork the child FastCGI processes, wait until they terminate, then clean up.
	def fork_children(options)
		@children = []
		options.count.times do |number|
			child = Child.new(options, @dispatch_script, number)
			@children.push(child)
		end

		trap("INT")  { kill_children }
		trap("TERM") { kill_children }
		trap("HUP")  { signal_children("HUP") }
		trap("USR1") { signal_children("USR1") }
		trap("USR2") { signal_children("USR2") }

		while !@children.empty? do
			@children.first.wait
			@children.pop
		end
	end

	def kill_children
		return if @exiting
		@exiting = true
		puts "Shutting down, killing #{@children.size} child processes..."
		@children.each do |child|
			child.kill
			puts "* Child #{child.number} killed."
		end
		exit!(2)
	end

	def signal_children(sig)
		puts "SIG#{sig} received, forwarded to child processes."
		@children.each do |child|
			child.signal(sig)
		end
	end
end

# Represents a forked FastCGI child process.
class RailsFork::Child
	attr_reader :number
	attr_reader :socket_file
	attr_reader :pid

	# Create a new FastCGI child process, which loads the given
	# dispatch script.
	def initialize(options, dispatch_script, number)
		@number = number
		@socket_file = "#{options.socket_dir}/#{options.prefix}#{number}"
		@pid = fork do
			# Note that we can't install signal handlers in the child process.
			# Rails's FastCGI handler installs its own handlers.

			if File.exist?(@socket_file)
				File.unlink(@socket_file)
			end
			server = UNIXServer.new(@socket_file)
			$stdin.reopen(server)
			puts "* Process #{number} listening on socket #{socket_file}"
			load dispatch_script
			exit!
		end
	end

	# Wait for the child to terminate and remove its socket file.
	def wait
		begin
			Process.wait(@pid)
		rescue Errno::ECHILD
		end
		File.unlink(@socket_file)
	end

	# Send SIGTERM to this child process and wait for it to terminate.
	def kill
		signal("TERM")
		wait
	end

	# Send a signal to this child process.
	def signal(sig)
		begin
			Process.kill(sig, @pid)
		rescue Errno::ESRCH
			# Do nothing, ignore errors.
		end
	end
end

RailsFork.new.start