From f5bd38d4b168540ddb0e33db9e35a51ec5d60c17 Mon Sep 17 00:00:00 2001
From: Peter Amstutz <peter.amstutz@curoverse.com>
Date: Fri, 17 Jun 2016 16:41:03 -0400
Subject: [PATCH] 9427: Add connection limit.

---
 services/api/app/middlewares/rack_socket.rb |  4 ++++
 services/api/lib/eventbus.rb                | 19 +++++++++++++++++--
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/services/api/app/middlewares/rack_socket.rb b/services/api/app/middlewares/rack_socket.rb
index 19350c40e6..8f82e585df 100644
--- a/services/api/app/middlewares/rack_socket.rb
+++ b/services/api/app/middlewares/rack_socket.rb
@@ -69,6 +69,10 @@ class RackSocket
   def call env
     request = Rack::Request.new(env)
     if request.path_info == @endpoint and Faye::WebSocket.websocket?(env)
+      if @handler.overloaded?
+        return [503, {"Content-Type" => "text/plain"}, ["Too many connections, try again later."]]
+      end
+
       ws = Faye::WebSocket.new(env, nil, :ping => 30)
 
       # Notify handler about new connection
diff --git a/services/api/lib/eventbus.rb b/services/api/lib/eventbus.rb
index c72ad9089f..0c5df827d0 100644
--- a/services/api/lib/eventbus.rb
+++ b/services/api/lib/eventbus.rb
@@ -75,6 +75,7 @@ class EventBus
     @channel = EventMachine::Channel.new
     @mtx = Mutex.new
     @bgthread = false
+    @connection_count = 0
   end
 
   # Push out any pending events to the connection +ws+
@@ -221,10 +222,17 @@ class EventBus
 
   # Constant maximum number of filters, to avoid silly huge database queries.
   MAX_FILTERS = 16
+  MAX_NOTIFY_BACKLOG = 1000
+  MAX_CONNECTIONS = 500
+
+  def overloaded?
+    @mtx.synchronize do
+      @connection_count >= MAX_CONNECTIONS
+    end
+  end
 
   # Called by RackSocket when a new websocket connection has been established.
   def on_connect ws
-
     # Disconnect if no valid API token.
     # current_user is included from CurrentApiClient
     if not current_user
@@ -241,6 +249,10 @@ class EventBus
     ws.queue = Queue.new
     ws.frame_mtx = Mutex.new
 
+    @mtx.synchronize do
+      @connection_count += 1
+    end
+
     # Subscribe to internal postgres notifications through @channel and
     # forward them to the thread associated with the connection.
     sub = @channel.subscribe do |msg|
@@ -274,7 +286,7 @@ class EventBus
       # Loop and react to socket events.
       loop do
         eventType, msg = ws.queue.pop
-        if ws.queue.length > 1000
+        if ws.queue.length > MAX_NOTIFY_BACKLOG
           ws.send ({status: 500, message: 'Notify backlog too long'}.to_json)
           ws.close
         else
@@ -287,6 +299,9 @@ class EventBus
           end
         end
       end
+      @mtx.synchronize do
+        @connection_count -= 1
+      end
     end
 
     # Start up thread to monitor the Postgres database, if none exists already.
-- 
2.30.2