#!/usr/bin/env ruby
#
#  equity 0.1 - simple queueing software load balancer
#
#  Usage: equity [-d] <listen-port> [<node-address>:]<node-port> ...
#
#  The -d option turns on debug mode. Equity will remain in the foreground and
#  print status messages.
#
#  Node addresses default to localhost.
#

require 'socket'
require 'equity/node'

# Prints a debug message if debugging is enabled.
def debug(socket, message)
  return unless $debugging
  print "[#{socket.object_id.to_s(16)}] " if socket
  puts message
end


# Process arguments and instantiate nodes.
$debugging = !!ARGV.delete('-d')

if ARGV.length < 2
  STDERR.puts 'equity 0.1'
  STDERR.puts 'Usage: equity [-d] <listen-port> [<node-address>:]<node-port> ...'
  STDERR.print "\n"
  STDERR.puts 'The -d option turns on debug mode. Equity will remain in the foreground and'
  STDERR.puts 'print status messages.'
  STDERR.print "\n"
  STDERR.puts 'Node addresses default to localhost.'
  exit(64) # EX_USAGE
end

$listen_port = ARGV.shift.to_i

$nodes = []
ARGV.each do |node_spec|
  address, port = node_spec.split(':', 2)
  if port.nil?
    port = address
    address = 'localhost'
  end
  $nodes << Equity::Node.new(address, port)
end

# Daemonize.
unless $debugging
  fork && exit
  Process.setsid
  trap 'SIGHUP', 'IGNORE'
  fork && exit
  Dir.chdir '/tmp'
  File.umask 0000
  ObjectSpace.each_object(IO) do |io|
    unless [STDIN, STDOUT, STDERR].include?(io)
      io.close rescue nil
    end
  end
  STDIN.reopen "/dev/null"
  STDOUT.reopen "/dev/null", "a"
  STDERR.reopen STDOUT
end

# Start up server.
$client_queue = []

$server = TCPServer.new(nil, $listen_port)
$server.listen(5)

# Set up SIGINT handler to print node counters.
trap 'SIGINT', Proc.new {
  debug(nil, "\nNode Counters")
  $nodes.each do |node|
    debug(nil, "#{node} - #{node.counter} connections")
  end
  exit
}

puts "Ready on port #{$listen_port}" if $debugging

# Loop forever.
while true
  # Wait for one or more sockets to be readable.
  sockets = [$server, $nodes.collect {|n| n.sockets}]
  sockets.flatten!
  selected = select(sockets, nil, nil, 5)
  if selected
    selected[0].each do |socket|
      if socket == $server
        # Incoming connection. Accept it and queue it.
        client = $server.accept
        $client_queue << client
        debug(client, 'connection accepted')
      else
        # Data received. Transfer it to the socket's mate.
        node = $nodes.find {|n| n.owns_socket?(socket)}
        next if node.nil?
        begin
          data = socket.recvfrom(65536)[0]
          raise Errno::EPIPE if data.empty?
          socket.mate.write(data)
        rescue Errno::EPIPE, Errno::ECONNRESET
          # Connection closed. Disconnect this node.
          node.disconnect
          debug(socket, 'disconnected')
        end
      end
    end
  end
  
  # Dequeue as many clients as possible, showing favor to nodes that have
  # handled the fewest clients.
  while !$client_queue.empty?
    client = $client_queue.shift
    found_a_node = false
    sorted_notes = $nodes.sort {|a,b| a.counter <=> b.counter}
    sorted_notes.each do |node|
      next if node.connected?
      begin
        node.connect(client)
      rescue
        debug(client, "connection to #{node} failed")
        next
      end
      debug(client, "assigned to #{node}")
      found_a_node = true
      break
    end
    unless found_a_node
      $client_queue.unshift(client)
      break
    end
  end
end
