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.