Load balancer minimalista en ruby (parte 2)

En el post anterior mostré un load balancer que implementé en ruby. Ahora le agregué algo de manejo de errores y algo de fail-over para que sea más estable.

Este es el código nuevo:

#!/usr/bin/env ruby
require 'socket'

def usage
  puts "Very simple balancer"
  puts "balancer.rb port host1:port1 [host2:port2 ....]"
  exit
end

def build_targets(argv)
  argv.map do
    |param|
    result = param.split(":")
    result[1] = result[1].to_i
    result
  end
end

class Balancer
  def initialize(port, *targets)
    @descriptors = []
    @next_step = {}
    
    @serverSocket = TCPServer.new( "", port )
    @serverSocket.setsockopt( Socket::SOL_SOCKET, Socket::SO_REUSEADDR, 1 )
    puts("Balancer started on port #{port}")
    
    @descriptors << @serverSocket
    
    @targets = targets
    @current_out = 0
  end
  
  def next_out
    out_socket = nil
    until out_socket
      begin
        @current_out = (@current_out + 1) % @targets.length
        descr = @targets[@current_out]
        out_socket = TCPSocket.new(descr[0], descr[1])
      rescue
        # do nothing on purpose
      end
    end
    out_socket
  end
  
  def accept_new_connection
    
    incoming = @serverSocket.accept
    outgoing = next_out
    
    @next_step[incoming] = outgoing
    @next_step[outgoing] = incoming
    
    @descriptors += [ incoming, outgoing ]
  end
  
  def propagate(sock)
    next_sock = @next_step[sock]
    next_sock.write(sock.read_nonblock(1000 * 1000))
  end
  
  def finish_connection(sock)
    next_sock = @next_step[sock]
    [sock, next_sock].each do
      |s|
      begin
        s.close
      rescue Object => e
        puts "Error closing socket: #{e.inspect}"
      end
      @descriptors.delete(s)
      @next_step.delete(s)
    end
  end
  
  def run
    loop do
      connections = select( @descriptors )
      if connections
        connections[0].each do
          |sock|
          if sock == @serverSocket then
            accept_new_connection
          else
            begin            
              if sock.eof? then
                finish_connection(sock)              
              else
                propagate(sock)
              end
            rescue
              finish_connection(sock)
            end
          end
        end
      end
    end
  end
end

trap("SIGINT") do
  exit
end

usage if ARGV.length < 2 
port = ARGV.shift.to_i
targets = build_targets(ARGV)

Balancer.new(port, *targets).run

Me parece que ahora pueden llegar a usarlo.

Happy hacking,
Aureliano.