bartender for mruby

bartender.rb

# -*- coding: utf-8 -*-
module Bartender
  class App
    class FDMap
      def initialize; @map = {}; end
      def []=(fd, v)
        @map[fd.to_i] = [fd, v]
      end
      def [](fd)
        _, v = @map[fd.to_i]
        v
      end
      def keys
        @map.values.collect {|pair| pair[0]}
      end
      def delete(fd)
        @map.delete(fd.to_i)
      end
    end

    def initialize
      @input = FDMap.new
      @output = FDMap.new
      @running = false
    end

    def run
      @running = true
      while @running
        step
        break if empty?
      end
    end

    def stop
      @running = false
    end

    def empty?
      @input.empty? && @output.empty?
    end
    
    def step(timeout=nil)
      r, w = IO.select(@input.keys, @output.keys, [], timeout)
      r.each {|fd| @input[fd].call }
      w.each {|fd| @output[fd].call }
    end

    def event_map(event)
      case event
      when :read
        @input
      when :write
        @output
      else
        raise 'invalid event'
      end
    end

    def []=(event, fd, callback)
      return delete(event, fd) unless callback
      event_map(event)[fd] = callback
    end

    def delete(event, fd)
      event_map(event).delete(fd)
    end

    def select_io(event, fd)
      it = Fiber.current
      self[event, fd] = Proc.new { it.resume }
      Fiber.yield
    ensure
      self.delete(event, fd)
    end

    def select_readable(fd); select_io(:read, fd); end
    def select_writable(fd); select_io(:write, fd); end

    def _read(fd, sz)
      while true
        it = fd.read_nonblock(sz)
        return it unless it == :wait_readable
        select_readable(fd)
      end
    end

    def _write(fd, buf)
      while true
        it = fd.write_nonblock(buf)
        return it unless it == :wait_writable
        select_writable(fd)
      end
    end
  end

  @app = App.new

  def primary; @app; end

  module_function :primary

  class Writer
    def initialize(bartender, fd)
      @bartender = bartender
      @fd = fd
      @pool = []
    end

    def write(buf, buffered=false)
      push(buf)
      flush unless buffered
    end

    def flush
      until @pool.empty?
        len = @bartender._write(@fd, @pool[0])
        pop(len)
      end
    end

    private
    def push(string)
      return if string.bytesize == 0
      @pool << string
    end
    
    def pop(size)
      return if size < 0 
      raise if @pool[0].bytesize < size
      
      if @pool[0].bytesize == size
        @pool.shift
      else
        unless @pool[0].encoding == Encoding::BINARY
          @pool[0] = @pool[0].dup.force_encoding(Encoding::BINARY)
        end
        @pool[0].slice!(0...size)
      end
    end
  end

  class Reader
    def initialize(bartender, fd)
      @bartender = bartender
      @buf = ''
      @fd = fd
    end

    def read(n)
      while @buf.bytesize < n
        chunk = @bartender._read(@fd, n)
        break if chunk.nil? || chunk.empty?
        @buf += chunk
      end
      @buf.slice!(0, n)
    end

    def read_until(sep="\r\n", chunk_size=8192)
      until (index = @buf.index(sep))
        @buf += @bartender._read(@fd, chunk_size)
      end
      @buf.slice!(0, index+sep.bytesize)
    end

    def readln
      read_until("\n")
    end
  end

  class Server
    def initialize(bartender, addr_or_port, port=nil, &blk)
      if port
        address = addr_or_port
      else
        address, port = nil, addr_or_port
      end
      @bartender = bartender
      create_listeners(address, port).each do |soc|
        @bartender[:read, soc] = Proc.new do
          client = soc.accept
          on_accept(client)
        end
      end
      @blk = blk
    end

    def create_listeners(address, port)
      unless port
        raise ArgumentError, "must specify port"
      end
      sockets = [TCPServer.new(address, port)]
      return sockets
    end

    def on_accept(client)
      Fiber.new do
        @blk.yield(client)
      end.resume
    end
  end
end

if __FILE__ == $0
  bar = Bartender.primary
  Bartender::Server.new(bar, 54321) do |c|
    puts "hello #{c.to_i}"
    reader = Bartender::Reader.new(bar, c)
    writer = Bartender::Writer.new(bar, c)
    while line = reader.readln rescue nil
      puts line
      writer.write(line)
    end
    puts "bye #{c.to_i}"
  end
  while true 
    bar.step
  end
end