Bouncer versión 2

Hice una nueva versión del bouncer que permite bouncear a varios targets desde un mismo host:

#!/usr/bin/env ruby

# == Synopsis
#
# bouncer.rb: Redirects TCP connections to distant machines. Handles simultaneously many connections.
#
# == Usage
#
# ruby redirect.rb [OPTION]
#
# -h, --help:
#   Show help
#
# --ip ip, -i ip:
#   Accept connections from ip (default 127.0.0.1)
#
# --port port, -p port:
#   Listen on port (default 12345)
#
# --target ip:port, -t ip:port
#   Connect to ip:port (default 127.0.0.1:23456)
#

require 'getoptlong'
require 'rdoc/usage'
require 'socket'

class Bouncer
 def self.for_host( source_ip, start_port, targets )
  rules = {}
  targets.each_with_index do
   |target, index|
   rules[ start_port + index ] = target
  end
  self.new( source_ip, rules )
 end
 def initialize( source_ip, port_target_map )
  @port_target_map = port_target_map
  @source_ip = source_ip

  @servers = @port_target_map.keys.map do
    |port|
    ss = TCPServer.new( "", port )
    ss.setsockopt( Socket::SOL_SOCKET, Socket::SO_REUSEADDR, 1 )
    ss
  end
    
  @descriptors = @servers.clone
  @next_step = {}
 end

 def handle_new_connection( ss )
  incoming = ss.accept
  if incoming.peeraddr[3] == @source_ip
   begin
    outgoing = TCPSocket.new( *@port_target_map[incoming.addr[1]] )
    @next_step[ incoming ] = outgoing
    @next_step[ outgoing ] = incoming

    @descriptors += [ incoming, outgoing ]
   rescue
    silent_close( incoming )
   end
  else
   silent_close( incoming )
  end
 end

 def silent_close( sock )
  begin
   sock.close
  rescue
   #do nothing intentionally
  end
 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|
   silent_close(s)
   @descriptors.delete(s)
   @next_step.delete(s)
  end
 end

 def run
  loop do
   connections = select( @descriptors )
   connections[0].each do
    |sock|
    if @servers.include? sock then
     handle_new_connection( sock )
    else
     begin
      sock.eof? ? finish_connection(sock) : propagate(sock)
     rescue
      finish_connection(sock)
     end
    end
   end
  end
 end
end

if $0 == __FILE__ then

 opts = GetoptLong.new(
   [ '--help', '-h', GetoptLong::NO_ARGUMENT ],
   [ '--ip', '-i', GetoptLong::REQUIRED_ARGUMENT ],
   [ '--port', '-p', GetoptLong::REQUIRED_ARGUMENT ],
   [ '--target', '-t', GetoptLong::REQUIRED_ARGUMENT ]
 )

 ip = '127.0.0.1'
 port = '12345'
 target = '127.0.0.1:23456'

 opts.each do
  |opt, arg|
  case opt
  when '--help'
   RDoc::usage
   exit
  when '--ip'
   ip = arg
  when '--port'
   port = arg
  when '--target'
   target = arg
  end
 end

 port = port.to_i
 target = target.split(":")

 trap("SIGINT") do
   exit
 end

 Bouncer.new(ip, port => [target[0], target[1].to_i]).run
end

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.

Load balancer minimalista en ruby

Update: Puse la versión más nueva en este post.

Ayer estuve inspirado y dediqué un poquito de mi tiempo a aprender un poco más sobre cómo manejar sockets. Hacía mucho tiempo que no programaba sockets en bajo nivel y nunca había necesitado usar select para manejar muchos al mismo tiempo. De hecho, la última vez que tuve que hacer algo parecido a un server que acepte muchas conexiones TCP lo hice en java 1.3, y en java 1.3 no hay select (agregaron algo parecido en java.nio, que apareció después).
Así que puse manos a la obra e hice un minicloncito de un load balancer en ruby.
Por supuesto que no hice todo. En particular, me comí el manejo de errores. Pero si hace round robin de conexiones y tiene menos de 100 líneas de código, lo que creo que lo hace un buen ejemplo de juguete sobre como manejar sockets en ruby.
Bueno, basta de cháchara, acá va el código completo:

#!/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
   @current_out = (@current_out + 1) % @targets.length
   @targets[@current_out]
 end

 def accept_new_connection
   out_descriptor = next_out

   incoming = @serverSocket.accept
   outgoing = TCPSocket.new(out_descriptor[0], out_descriptor[1])

   @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|
     s.close
     @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
           if sock.eof? then
             finish_connection(sock)
           else
             propagate(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

Para usarlo, copien todo el codigo a un archivo (como balancer.rb) y fijense en el cartel que explica los parámetros. El primer parámetro es el puerto donde escucha y los otros son pares host:port a los que tiene que conectarse alternativamente.

Bueno, eso es todo por ahora.

Happy hacking,
Aureliano.