websocket_celluloid_server.rb

require 'forwardable'
require 'websocket/driver'
require 'celluloid/io'
require 'active_support/core_ext/numeric/bytes'

class WebSocketHandler
  extend Forwardable

  def_delegators :@driver, :on, :parse, :env, :protocol

  def initialize(socket, options = {})
    @timers = []
    @socket = socket
    @driver = WebSocket::Driver.server(@socket, options)
    @driver.on(:close) { bye }
  end

  def start
    @driver.start
  rescue IOError, Errno::EPIPE, Errno::ECONNRESET
    bye
  end

  def write(msg = '', type = nil, code = nil)
    @driver.frame(msg, type, code)
  rescue IOError, Errno::EPIPE, Errno::ECONNRESET
    bye
  end

  def close
    @driver.close
  rescue IOError, Errno::EPIPE, Errno::ECONNRESET
    bye
  end

  def bye
    cancel_timers!
    @socket.close unless @socket.closed?
  end

  def every(interval, &block)
    @timers << Celluloid.every(interval, &block)
  end

  def after(interval, &block)
    @timers << Celluloid.after(interval, &block)
  end

  def cancel_timers!
    @timers.map(&:cancel)
  end
end

class Server
  include Celluloid::IO
  include Celluloid::Logger

  finalizer :shutdown

  def initialize
    @server = TCPServer.new('0.0.0.0', 8080)
    @server.setsockopt(Socket::IPPROTO_TCP, :TCP_NODELAY, 1)
    async.run
  end

  def run
    loop { async.handle_connection(@server.accept) }
  end

  def handle_connection(socket)
    websocket = WebSocketHandler.new(socket, max_length: 1.megabyte, protocols: ['wamp.2.json', 'wamp'])

    websocket.on(:error) {|e| error e}

    websocket.on(:connect) do
      if WebSocket::Driver.websocket? websocket.env
        websocket.start
      else
        websocket.close
      end
    end

    websocket.on(:open) do
      info websocket.instance_variable_get(:@driver).instance_variable_get(:@delegate).instance_variable_get(:@protocol)
      websocket.write(Time.now.to_s)
      websocket.every(1) { websocket.write(Time.now.to_s) }
    end

    websocket.on(:message) {|e| info e.data}

    loop do
      begin
        websocket.parse(socket.readpartial(16384))
      rescue IOError, Errno::EPIPE, Errno::ECONNRESET
        break
      end
    end
  ensure
    websocket.close rescue nil
  end

  def shutdown
    @server.close if @server
  end
end

Server.run