Client API for server shutdown protocol command (#328)

2 files changed, 44 insertions, 0 deletions

diff --git a/pymemcache/client/base.py b/pymemcache/client/base.py
index a0816cf..0999ba7 100644
--- a/pymemcache/client/base.py
+++ b/pymemcache/client/base.py
@@ -817,6 +817,32 @@ class Client(object):
         self._misc_cmd([cmd], b'quit', True)
         self.close()
 
+    def shutdown(self, graceful=False):
+        """
+        The memcached "shutdown" command.
+
+        This will request shutdown and eventual termination of the server,
+        optionally preceded by a graceful stop of memcached's internal state
+        machine. Note that the server needs to have been started with the
+        shutdown protocol command enabled with the --enable-shutdown flag.
+
+        Args:
+          graceful: optional bool, True to request a graceful shutdown with
+                    SIGUSR1 (defaults to False, i.e. SIGINT shutdown).
+        """
+        cmd = b'shutdown'
+        if graceful:
+            cmd += b' graceful'
+        cmd += b'\r\n'
+
+        # The shutdown command raises a server-side error if the shutdown
+        # protocol command is not enabled. Otherwise, a successful shutdown
+        # is expected to close the remote end of the transport.
+        try:
+            self._misc_cmd([cmd], b'shutdown', False)
+        except MemcacheUnexpectedCloseError:
+            pass
+
     def _raise_errors(self, line, name):
         if line.startswith(b'ERROR'):
             raise MemcacheUnknownCommandError(name)
@@ -1249,6 +1275,10 @@ class PooledClient(object):
             finally:
                 self.client_pool.destroy(client)
 
+    def shutdown(self, graceful=False):
+        with self.client_pool.get_and_release(destroy_on_fail=True) as client:
+            client.shutdown(graceful)
+
     def __setitem__(self, key, value):
         self.set(key, value, noreply=True)
 
diff --git a/pymemcache/test/test_client.py b/pymemcache/test/test_client.py
index 93d327c..7479e48 100644
--- a/pymemcache/test/test_client.py
+++ b/pymemcache/test/test_client.py
@@ -31,6 +31,7 @@ from pymemcache.client.base import PooledClient, Client, normalize_server_spec
 from pymemcache.exceptions import (
     MemcacheClientError,
     MemcacheServerError,
+    MemcacheUnexpectedCloseError,
     MemcacheUnknownCommandError,
     MemcacheUnknownError,
     MemcacheIllegalInputError
@@ -688,6 +689,19 @@ class TestClient(ClientTestMixin, unittest.TestCase):
         assert result is None
         assert client.sock is None
 
+    def test_shutdown(self):
+        client = self.make_client([MemcacheUnexpectedCloseError('shutdown')])
+        result = client.shutdown()
+        assert result is None
+
+    def test_shutdown_disabled(self):
+        def _shutdown():
+            client = self.make_client([b'ERROR: shutdown not enabled\r\n'])
+            client.shutdown()
+
+        with pytest.raises(MemcacheUnknownCommandError):
+            _shutdown()
+
     def test_replace_stored(self):
         client = self.make_client([b'STORED\r\n'])
         result = client.replace(b'key', b'value', noreply=False)