base.rb in angelo-0.1.4

- old
+ new

@@ -8,11 +8,13 @@
     def_delegators :@responder, :content_type, :headers, :request
 
     @@addr = DEFAULT_ADDR
     @@port = DEFAULT_PORT
 
-    if ARGV.any?
+    @@ping_time = DEFAULT_PING_TIME
+
+    if ARGV.any? and not Kernel.const_defined?('RSpec')
       require 'optparse'
       OptionParser.new { |op|
         op.on('-p port',   'set the port (default is 4567)')      { |val| @@port = Integer(val) }
         op.on('-o addr',   "set the host (default is #{@@addr})") { |val| @@addr = val }
       }.parse!(ARGV.dup)
@@ -20,11 +22,11 @@
 
     attr_accessor :responder
 
     class << self
 
-      attr_accessor :app_file
+      attr_accessor :app_file, :server
 
       def inherited subclass
         subclass.app_file = caller(1).map {|l| l.split(/:(?=|in )/, 3)[0,1]}.flatten[0]
 
         def subclass.root
@@ -74,56 +76,114 @@
 
       def socket path, &block
         routes[:socket][path] = WebsocketResponder.new &block
       end
 
+      def on_pong &block
+        WebsocketResponder.on_pong = block
+      end
+
+      def async name, &block
+        Angelo::Server.define_action name, &block
+      end
+
       def websockets
-        @websockets ||= WebsocketsArray.new
+        @websockets ||= WebsocketsArray.new server
         @websockets.reject! &:closed?
         @websockets
       end
 
       def content_type type
         Responder.content_type type
       end
 
       def run addr = @@addr, port = @@port
         @server = Angelo::Server.new self, addr, port
+        @server.async.ping_websockets
         trap "INT" do
           @server.terminate if @server and @server.alive?
           exit
         end
         sleep
       end
 
     end
 
+    def async meth, *args
+      self.class.server.async.__send__ meth, *args
+    end
+
     def params
       @params ||= case request.method
                   when GET;  parse_query_string
                   when POST; parse_post_body
                   when PUT;  parse_post_body
                   end
     end
 
     def websockets; self.class.websockets; end
 
+    async :handle_websocket do |ws|
+      begin
+        while !ws.closed? do
+          ws.read
+        end
+      rescue IOError
+        websockets.remove_socket ws
+      end
+    end
+
+    async :ping_websockets do
+      every(@@ping_time) do
+        websockets.each do |ws|
+          ws.socket << ::WebSocket::Message.ping.to_data
+        end
+      end
+    end
+
     class WebsocketsArray < Array
+      include Celluloid::Logger
 
+      @@peeraddrs = {}
+      @@socket_context = {}
+
+      def initialize server, context = nil
+        @context, @server = context, server
+        super()
+      end
+
+      def << ws
+        @@socket_context[ws] = @context if @context
+        @@peeraddrs[ws] = ws.peeraddr
+        @server.async.handle_websocket ws
+        super ws
+      end
+
       def each &block
         super do |ws|
           begin
             yield ws
-          rescue Reel::SocketError => rse
-            warn "#{rse.class} - #{rse.message}"
-            delete ws
+          rescue Reel::SocketError
+            remove_socket ws
           end
         end
       end
 
+      def remove_socket ws
+        if c = @@socket_context[ws]
+          warn "removing socket from context ':#{c}' (#{@@peeraddrs[ws][2]})"
+          self[c].delete ws
+        else
+          warn "removing socket (#{@@peeraddrs[ws][2]})"
+          delete ws
+        end
+        @@peeraddrs.delete ws
+      end
+
       def [] context
+        raise ArgumentError.new "symbol required" unless Symbol === context
         @@websockets ||= {}
-        @@websockets[context] ||= self.class.new
+        @@websockets[context] ||= self.class.new @server, context
       end
 
     end
 
   end