From 9a77d1df995d41aa41233aa4e50c95b4144c7269 Mon Sep 17 00:00:00 2001
From: Chico Venancio <fvenancio@wikimedia.org>
Date: Mon, 12 Mar 2018 22:03:14 +0000
Subject: [PATCH] test cull_idle_servers url encoding

---
 images/paws-hub/Dockerfile           |   2 +
 images/paws-hub/cull_idle_servers.py | 151 +++++++++++++++++++++++++++
 2 files changed, 153 insertions(+)
 create mode 100755 images/paws-hub/cull_idle_servers.py

diff --git a/images/paws-hub/Dockerfile b/images/paws-hub/Dockerfile
index 99969caf..2b57a50e 100644
--- a/images/paws-hub/Dockerfile
+++ b/images/paws-hub/Dockerfile
@@ -9,8 +9,10 @@ RUN adduser --disabled-password \
     --home ${HOME} \
     --force-badname \
     ${NB_USER}
+COPY cull_idle_servers.py /usr/local/bin/cull_idle_servers.py
 
 RUN chown ${NB_USER}:${NB_USER} /srv/jupyterhub
 
 USER ${NB_USER}
+
 CMD ["jupyterhub", "--config", "/srv/jupyterhub_config.py"]
diff --git a/images/paws-hub/cull_idle_servers.py b/images/paws-hub/cull_idle_servers.py
new file mode 100755
index 00000000..db0e4658
--- /dev/null
+++ b/images/paws-hub/cull_idle_servers.py
@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+# Imported from https://github.com/jupyterhub/jupyterhub/blob/0.8.0rc1/examples/cull-idle/cull_idle_servers.py
+"""script to monitor and cull idle single-user servers
+
+Caveats:
+
+last_activity is not updated with high frequency,
+so cull timeout should be greater than the sum of:
+
+- single-user websocket ping interval (default: 30s)
+- JupyterHub.last_activity_interval (default: 5 minutes)
+
+You can run this as a service managed by JupyterHub with this in your config::
+
+
+    c.JupyterHub.services = [
+        {
+            'name': 'cull-idle',
+            'admin': True,
+            'command': 'python cull_idle_servers.py --timeout=3600'.split(),
+        }
+    ]
+
+Or run it manually by generating an API token and storing it in `JUPYTERHUB_API_TOKEN`:
+
+    export JUPYTERHUB_API_TOKEN=`jupyterhub token`
+    python cull_idle_servers.py [--timeout=900] [--url=http://127.0.0.1:8081/hub/api]
+"""
+
+import datetime
+import json
+import os
+
+try:
+    from urllib.parse import quote
+except ImportError:
+    from urllib import quote
+
+from dateutil.parser import parse as parse_date
+
+from tornado.gen import coroutine
+from tornado.log import app_log
+from tornado.httpclient import AsyncHTTPClient, HTTPRequest
+from tornado.ioloop import IOLoop, PeriodicCallback
+from tornado.options import define, options, parse_command_line
+
+
+@coroutine
+def cull_idle(url, api_token, timeout, cull_users=False):
+    """Shutdown idle single-user servers
+
+    If cull_users, inactive *users* will be deleted as well.
+    """
+    auth_header = {
+            'Authorization': 'token %s' % api_token
+        }
+    req = HTTPRequest(url=url + '/users',
+        headers=auth_header,
+    )
+    now = datetime.datetime.utcnow()
+    cull_limit = now - datetime.timedelta(seconds=timeout)
+    client = AsyncHTTPClient()
+    resp = yield client.fetch(req)
+    users = json.loads(resp.body.decode('utf8', 'replace'))
+    futures = []
+
+    @coroutine
+    def cull_one(user, last_activity):
+        """cull one user"""
+
+        # shutdown server first. Hub doesn't allow deleting users with running servers.
+        if user['server']:
+            app_log.info("Culling server for %s (inactive since %s)", user['name'], last_activity)
+            req = HTTPRequest(url=url + '/users/%s/server' % quote(user['name']),
+                method='DELETE',
+                headers=auth_header,
+            )
+            resp = yield client.fetch(req)
+            if resp.code == 202:
+                msg = "Server for {} is slow to stop.".format(user['name'])
+                if cull_users:
+                    app_log.warning(msg + " Not culling user yet.")
+                    # return here so we don't continue to cull the user
+                    # which will fail if the server is still trying to shutdown
+                    return
+                app_log.warning(msg)
+        if cull_users:
+            app_log.info("Culling user %s (inactive since %s)", user['name'], last_activity)
+            req = HTTPRequest(url=url + '/users/%s' % user['name'],
+                method='DELETE',
+                headers=auth_header,
+            )
+            yield client.fetch(req)
+
+    for user in users:
+        if not user['server'] and not cull_users:
+            # server not running and not culling users, nothing to do
+            continue
+        if not user['last_activity']:
+            continue
+        last_activity = parse_date(user['last_activity'])
+        if last_activity < cull_limit:
+            # user might be in a transition (e.g. starting or stopping)
+            # don't try to cull if this is happening
+            if user['pending']:
+                app_log.warning("Not culling user %s with pending %s", user['name'], user['pending'])
+                continue
+            futures.append((user['name'], cull_one(user, last_activity)))
+        else:
+            app_log.debug("Not culling %s (active since %s)", user['name'], last_activity)
+
+    for (name, f) in futures:
+        try:
+            yield f
+        except Exception:
+            app_log.exception("Error culling %s", name)
+        else:
+            app_log.debug("Finished culling %s", name)
+
+
+if __name__ == '__main__':
+    define('url', default=os.environ.get('JUPYTERHUB_API_URL'), help="The JupyterHub API URL")
+    define('timeout', default=600, help="The idle timeout (in seconds)")
+    define('cull_every', default=0, help="The interval (in seconds) for checking for idle servers to cull")
+    define('cull_users', default=False,
+        help="""Cull users in addition to servers.
+                This is for use in temporary-user cases such as tmpnb.""",
+    )
+
+    parse_command_line()
+    if not options.cull_every:
+        options.cull_every = options.timeout // 2
+    api_token = os.environ['JUPYTERHUB_API_TOKEN']
+
+    try:
+        AsyncHTTPClient.configure("tornado.curl_httpclient.CurlAsyncHTTPClient")
+    except ImportError as e:
+        app_log.warning("Could not load pycurl: %s\npycurl is recommended if you have a large number of users.", e)
+
+    loop = IOLoop.current()
+    cull = lambda : cull_idle(options.url, api_token, options.timeout, options.cull_users)
+    # schedule first cull immediately
+    # because PeriodicCallback doesn't start until the end of the first interval
+    loop.add_callback(cull)
+    # schedule periodic cull
+    pc = PeriodicCallback(cull, 1e3 * options.cull_every)
+    pc.start()
+    try:
+        loop.start()
+    except KeyboardInterrupt:
+        pass