takuma104のRuby/Rails日記

Ruby/EventMachine での非同期サーバを Generator を使って同期サーバっぽく書く [その２]

ruby, eventmachine | 03:28 |

昨日の続きで、さらに TCPServer での Thread-loop 内と全く同じメソッドが EventMachine でも使えるように改良してみました。

使い方は、

ruby hoge.rb -e

とか -e オプションをつけると、EventMachine を使用したサーバになり、-eをつけないと TCPServer を使用したサーバになります。

ポイントは、on_accept メソッドを両者で使い回しているところです。EventMachine のほうを socket に合わせる形にしている FilberIO クラスにほとんど秘密があります（この名前が微妙だな。なにかいい名前無いかな??）

ソースは以下です。

#!/usr/bin/env ruby -KU

require 'rubygems'
require 'eventmachine'
require 'socket'
require 'generator'
require 'optparse'

class FiberIO
  def initialize(connection)
    @recv_data = ''
    @connection = connection
    Generator.new do |g|
      @generator = g
      yield(self)
    end
  end
  
  def read(size)
    loop do
      d = pop(size)
      if d
        return d
      else
        remain = size - @recv_data.size
        @generator.yield(remain)
      end
    end
  end
  
  def write(buf)
    @connection.send_data buf
  end
  
  def push(data)
    @recv_data << data
    if @generator.next?
      @generator.next
    end
  end
  
  def close
    @connection.close_connection_after_writing
  end
  
  def peeraddr
    addr = Socket.unpack_sockaddr_in(@connection.get_peername)
    ["AF_INET", addr[0], addr[1], addr[1]]
  end
  
private
  def pop(size)
    if @recv_data.size >= size
      r = @recv_data[0..size-1]
      @recv_data = @recv_data[size..-1]
      r
    else
      nil
    end
  end
end

class GeneratorSampleConnection < EventMachine::Connection
  def initialize(arg)
    @proc = arg
  end
  
  def post_init
    @io = FiberIO.new(self) {|io| @proc.call(io)}
  end

  def receive_data(data)
    @io.push(data)
  end
end

def on_accept(sock)
  puts "Connected from %s:%d" % [sock.peeraddr[3],sock.peeraddr[1]]
  puts sock.read(10)
  puts "---"
  puts sock.read(20)
  puts "---END"
  sock.write ">> OK\r\n"
  sock.close
end

def loop_eventmachine(host, port)
  EventMachine::run do
    EventMachine.start_server(host, port, GeneratorSampleConnection, Proc.new do |io|
      on_accept(io)
    end)
    puts "Now accepting connections on address port #{port} (EventMachine)"
  end
end

def loop_tcpserver(port)
  gs = TCPServer.open(port)
  puts "Now accepting connections on address port #{port} (TCPServer)"
  loop do
    Thread.start(gs.accept) do |sock|
      on_accept(sock)
    end
  end
end

if $0 == __FILE__
  host = "127.0.0.1"
  options = {:p=>3793,:e=>false}
  
  opt = OptionParser.new
  opt.on('-p VAL') {|v| options[:p] = v.to_i }
  opt.on('-e') {|v| options[:e] = true }
  opt.parse!(ARGV)

  if ARGV.size != 0
    puts "-e      : use EventMachine (default=TCPServer)"
    puts "-p port : listening port (default=3793)"
  end
  
  if options[:e]
    loop_eventmachine(host, options[:p])
  else
    loop_tcpserver(options[:p])
  end
end