From 27de44b0ecba3cc6fd452e5ae078d75327f833cb Mon Sep 17 00:00:00 2001
From: yuvipanda <yuvipanda@gmail.com>
Date: Thu, 27 Jul 2017 16:32:45 -0700
Subject: [PATCH] Add support for limiting the number of concurrent spawns

---
 jupyterhub/app.py           | 19 +++++++++++++++++--
 jupyterhub/handlers/base.py | 15 +++++++++++++++
 jupyterhub/objects.py       |  9 +++++++++
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/jupyterhub/app.py b/jupyterhub/app.py
index 6e083616..f16e1fcf 100644
--- a/jupyterhub/app.py
+++ b/jupyterhub/app.py
@@ -531,11 +531,11 @@ class JupyterHub(Application):
     @default('authenticator')
     def _authenticator_default(self):
         return self.authenticator_class(parent=self, db=self.db)
-    
+
     allow_named_servers = Bool(False,
         help="Allow named single-user servers per user"
     ).tag(config=True)
-    
+
     # class for spawning single-user servers
     spawner_class = Type(LocalProcessSpawner, Spawner,
         help="""The class to use for spawning single-user servers.
@@ -544,6 +544,20 @@ class JupyterHub(Application):
         """
     ).tag(config=True)
 
+    concurrent_spawn_limit = Integer(
+        None,
+        allow_none=True,
+        help="""
+        Maximum number of concurrent users that can be spawning at a time.
+
+        If more than this many users attempt to spawn at a time, they
+        enter an exponential delay with a timeout.
+
+        If set to `None`, no concurrent limit is enforced. If set to 0,
+        then no spawning is allowed.
+        """
+    ).tag(config=True)
+
     db_url = Unicode('sqlite:///jupyterhub.sqlite',
         help="url for the database. e.g. `sqlite:///jupyterhub.sqlite`"
     ).tag(config=True)
@@ -1245,6 +1259,7 @@ class JupyterHub(Application):
             statsd=self.statsd,
             allow_named_servers=self.allow_named_servers,
             oauth_provider=self.oauth_provider,
+            concurrent_spawn_limit=self.concurrent_spawn_limit,
         )
         # allow configured settings to have priority
         settings.update(self.tornado_settings)
diff --git a/jupyterhub/handlers/base.py b/jupyterhub/handlers/base.py
index 6c02c6fd..e34975b0 100644
--- a/jupyterhub/handlers/base.py
+++ b/jupyterhub/handlers/base.py
@@ -366,6 +366,20 @@ class BaseHandler(RequestHandler):
     def spawn_single_user(self, user, server_name='', options=None):
         if server_name in user.spawners and user.spawners[server_name]._spawn_pending:
             raise RuntimeError("Spawn already pending for: %s" % user.name)
+
+        if self.settings['concurrent_spawn_limit'] is not None and \
+           self.hub.pending_spawns > self.settings['concurrent_spawn_limit']:
+            # This will throw an error if we hit our timeout
+            self.log.info(
+                'More than %s pending spawns, throttling',
+                self.settings['concurrent_spawn_limit']
+            )
+            yield exponential_backoff(
+                lambda: self.hub.pending_spawns < self.settings['concurrent_spawn_limit'],
+                'Too many users are starting right now, try again later',
+            )
+
+        self.hub.pending_spawns += 1
         tic = IOLoop.current().time()
         user_server_name = user.name
         if server_name:
@@ -402,6 +416,7 @@ class BaseHandler(RequestHandler):
                 spawner.add_poll_callback(self.user_stopped, user)
             finally:
                 spawner._proxy_pending = False
+                self.hub.pending_spawns -= 1
 
         try:
             yield gen.with_timeout(timedelta(seconds=self.slow_spawn_timeout), f)
diff --git a/jupyterhub/objects.py b/jupyterhub/objects.py
index 15e2b420..184b3995 100644
--- a/jupyterhub/objects.py
+++ b/jupyterhub/objects.py
@@ -149,6 +149,15 @@ class Hub(Server):
 
     cookie_name = 'jupyter-hub-token'
 
+    pending_spawns = Integer(
+        0,
+        help="""
+        Number of users currently attempting to spawn.
+
+        Tracked by BaseHandler.spawn_single_user. It's in the Hub object
+        since this needs a central place for it to be kept track of.
+        """
+    )
     @property
     def server(self):
         warnings.warn("Hub.server is deprecated in JupyterHub 0.8. Access attributes on the Hub directly.",