From d812affef022b86b9c58a8110d60f5b13c3faa50 Mon Sep 17 00:00:00 2001
From: Yun Xu <yunx@zillowgroup.com>
Date: Sun, 25 Jun 2017 00:51:14 -0700
Subject: [PATCH] add graceful_shutdown_timeout to gunicorn worker

---
 sanic/app.py    |  3 ++-
 sanic/config.py |  1 +
 sanic/server.py |  8 ++++++++
 sanic/worker.py | 18 +++++++++++++++++-
 4 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/sanic/app.py b/sanic/app.py
index ff680d9c..5b071200 100644
--- a/sanic/app.py
+++ b/sanic/app.py
@@ -701,7 +701,8 @@ class Sanic:
             'backlog': backlog,
             'has_log': has_log,
             'websocket_max_size': self.config.WEBSOCKET_MAX_SIZE,
-            'websocket_max_queue': self.config.WEBSOCKET_MAX_QUEUE
+            'websocket_max_queue': self.config.WEBSOCKET_MAX_QUEUE,
+            'graceful_shutdown_timeout': self.config.GRACEFUL_SHUTDOWN_TIMEOUT
         }
 
         # -------------------------------------------- #
diff --git a/sanic/config.py b/sanic/config.py
index e3563bc1..075bc47d 100644
--- a/sanic/config.py
+++ b/sanic/config.py
@@ -127,6 +127,7 @@ class Config(dict):
         self.KEEP_ALIVE = keep_alive
         self.WEBSOCKET_MAX_SIZE = 2 ** 20  # 1 megabytes
         self.WEBSOCKET_MAX_QUEUE = 32
+        self.GRACEFUL_SHUTDOWN_TIMEOUT = 15.0 # 15 sec
 
         if load_env:
             self.load_environment_vars()
diff --git a/sanic/server.py b/sanic/server.py
index 11c2452c..7ceef655 100644
--- a/sanic/server.py
+++ b/sanic/server.py
@@ -363,6 +363,14 @@ class HttpProtocol(asyncio.Protocol):
             return True
         return False
 
+    def close(self):
+        """
+        Force close the connection.
+        """
+        if self.transport is not None:
+            self.transport.close()
+            self.transport = None
+
 
 def update_current_time(loop):
     """Cache the current time, since it is needed at the end of every
diff --git a/sanic/worker.py b/sanic/worker.py
index 876354ce..3adef65f 100644
--- a/sanic/worker.py
+++ b/sanic/worker.py
@@ -91,8 +91,24 @@ class GunicornWorker(base.Worker):
             for conn in self.connections:
                 conn.close_if_idle()
 
-            while self.connections:
+            # gracefully shutdown timeout
+            start_shutdown = 0
+            graceful_shutdown_timeout = self.cfg.graceful_timeout
+            while self.connections and (start_shutdown < graceful_shutdown_timeout):
                 await asyncio.sleep(0.1)
+                start_shutdown = start_shutdown + 0.1
+
+            # Force close non-idle connection after waiting for
+            # graceful_shutdown_timeout
+            coros = []
+            for conn in self.connections:
+                if hasattr(conn, "websocket") and conn.websocket:
+                    coros.append(conn.websocket.close_connection(force=True))
+                else:
+                    conn.close()
+            _shutdown = asyncio.gather(*coros, loop=self.loop)
+            await _shutdown
+
 
     async def _run(self):
         for sock in self.sockets: