X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/aa920f1659aa830861c196617fe664f35b5c12ef..6f1c0e1dd16f53ce4fb7e8785a130e5f588500e5:/services/api/app/middlewares/rack_socket.rb

diff --git a/services/api/app/middlewares/rack_socket.rb b/services/api/app/middlewares/rack_socket.rb
index 779c675030..8f82e585df 100644
--- a/services/api/app/middlewares/rack_socket.rb
+++ b/services/api/app/middlewares/rack_socket.rb
@@ -1,21 +1,41 @@
 require 'rack'
 require 'faye/websocket'
-require 'oj'
 require 'eventmachine'
 
+# A Rack middleware to handle inbound websocket connection requests and hand
+# them over to the faye websocket library.
 class RackSocket
 
   DEFAULT_ENDPOINT  = '/websocket'
 
+  # Stop EventMachine on signal, this should give it a chance to to unwind any
+  # open connections.
   def die_gracefully_on_signal
     Signal.trap("INT") { EM.stop }
     Signal.trap("TERM") { EM.stop }
   end
 
+  # Create a new RackSocket handler
+  # +app+  The next layer of the Rack stack.
+  #
+  # Accepts options:
+  # +:handler+ (Required) A class to handle new connections.  #initialize will
+  # call handler.new to create the actual handler instance object.  When a new
+  # websocket connection is established, #on_connect on the handler instance
+  # object will be called with the new connection.
+  #
+  # +:mount+ The HTTP request path that will be recognized for websocket
+  # connect requests, defaults to '/websocket'.
+  #
+  # +:websocket_only+  If true, the server will only handle websocket requests,
+  # and all other requests will result in an error.  If false, unhandled
+  # non-websocket requests will be passed along on to 'app' in the usual Rack
+  # way.
   def initialize(app = nil, options = nil)
     @app = app if app.respond_to?(:call)
     @options = [app, options].grep(Hash).first || {}
     @endpoint = @options[:mount] || DEFAULT_ENDPOINT
+    @websocket_only = @options[:websocket_only] || false
 
     # from https://gist.github.com/eatenbyagrue/1338545#file-eventmachine-rb
     if defined?(PhusionPassenger)
@@ -38,64 +58,33 @@ class RackSocket
       }
     end
 
-    @channel = EventMachine::Channel.new
-    @bgthread = nil
+    # Create actual handler instance object from handler class.
+    @handler = @options[:handler].new
   end
 
+  # Handle websocket connection request, or pass on to the next middleware
+  # supplied in +app+ initialize (unless +:websocket_only+ option is true, in
+  # which case return an error response.)
+  # +env+ the Rack environment with information about the request.
   def call env
     request = Rack::Request.new(env)
     if request.path_info == @endpoint and Faye::WebSocket.websocket?(env)
-      ws = Faye::WebSocket.new(env)
-
-      sub = @channel.subscribe do |msg|
-        puts "sending #{msg}"
-        ws.send({:message => "log"}.to_json)
+      if @handler.overloaded?
+        return [503, {"Content-Type" => "text/plain"}, ["Too many connections, try again later."]]
       end
 
-      ws.on :message do |event|
-        puts "got #{event.data}"
-        ws.send(event.data)
-      end
+      ws = Faye::WebSocket.new(env, nil, :ping => 30)
 
-      ws.on :close do |event|
-        p [:close, event.code, event.reason]
-        @channel.unsubscribe sub
-        ws = nil
-      end
-
-      unless @bgthread
-        @bgthread = true
-        Thread.new do
-          # from http://stackoverflow.com/questions/16405520/postgres-listen-notify-rails
-          ActiveRecord::Base.connection_pool.with_connection do |connection|
-            conn = connection.instance_variable_get(:@connection)
-            begin
-              conn.async_exec "LISTEN logs"
-              while true
-                conn.wait_for_notify do |channel, pid, payload|
-                  puts "Received a NOTIFY on channel #{channel}"
-                  puts "from PG backend #{pid}"
-                  puts "saying #{payload}"
-                  @channel.push true
-                end
-              end
-            ensure
-              # Don't want the connection to still be listening once we return
-              # it to the pool - could result in weird behavior for the next
-              # thread to check it out.
-              conn.async_exec "UNLISTEN *"
-            end
-          end
-        end
-      end
+      # Notify handler about new connection
+      @handler.on_connect ws
 
       # Return async Rack response
       ws.rack_response
-    else
+    elsif not @websocket_only
       @app.call env
+    else
+      [406, {"Content-Type" => "text/plain"}, ["Only websocket connections are permitted on this port."]]
     end
   end
 
 end
-
-