Ensure us to close socket on MemcacheUnexpectedCloseError

During network maintenance operations or during server reboot users could
face MemcacheUnexpectedCloseError.

This kind of exception will be raised down the stack. It's a design
choice [1]. Users should manage this kind of exception.

However, at some points the opened sockets remains more or less opened
after this kind of fail on the client side [2]. It will require n+1
calls to see them recreated when the connection will be possible. That's
could be problem with services that implement several cache logics in parallel
and where their clients are based on pymemcache.

These changes simply ensure to close the socket properly. The socket
will be recreated by the client if needed.

[1] https://github.com/pinterest/pymemcache/issues/307
[2] https://bugzilla.redhat.com/show_bug.cgi?id=1977711

2 files changed, 38 insertions, 4 deletions

diff --git a/pymemcache/client/base.py b/pymemcache/client/base.py
index abba063..64f65d7 100644
--- a/pymemcache/client/base.py
+++ b/pymemcache/client/base.py
@@ -993,7 +993,12 @@ class Client(object):
             except Exception as e:
                 raise ValueError("Unable to parse line %s: %s" % (line, e))
 
-        buf, value = _readvalue(self.sock, buf, int(size))
+        value = None
+        try:
+            buf, value = _readvalue(self.sock, buf, int(size))
+        except MemcacheUnexpectedCloseError:
+            self.close()
+            raise
         key = remapped_keys[key]
         value = self.serde.deserialize(key, value, int(flags))
 
@@ -1019,9 +1024,14 @@ class Client(object):
             self.sock.sendall(cmd)
 
             buf = b''
+            line = None
             result = {}
             while True:
-                buf, line = _readline(self.sock, buf)
+                try:
+                    buf, line = _readline(self.sock, buf)
+                except MemcacheUnexpectedCloseError:
+                    self.close()
+                    raise
                 self._raise_errors(line, name)
                 if line == b'END' or line == b'OK':
                     return result
@@ -1091,8 +1101,13 @@ class Client(object):
 
             results = {}
             buf = b''
+            line = None
             for key in keys:
-                buf, line = _readline(self.sock, buf)
+                try:
+                    buf, line = _readline(self.sock, buf)
+                except MemcacheUnexpectedCloseError:
+                    self.close()
+                    raise
                 self._raise_errors(line, name)
 
                 if line in VALID_STORE_RESULTS[name]:
@@ -1116,8 +1131,13 @@ class Client(object):
 
             results = []
             buf = b''
+            line = None
             for cmd in cmds:
-                buf, line = _readline(self.sock, buf)
+                try:
+                    buf, line = _readline(self.sock, buf)
+                except MemcacheUnexpectedCloseError:
+                    self.close()
+                    raise
                 self._raise_errors(line, cmd_name)
                 results.append(line)
             return results
diff --git a/pymemcache/test/test_client.py b/pymemcache/test/test_client.py
index c09f48a..118a229 100644
--- a/pymemcache/test/test_client.py
+++ b/pymemcache/test/test_client.py
@@ -820,6 +820,20 @@ class TestClient(ClientTestMixin, unittest.TestCase):
         with pytest.raises(MemcacheServerError):
             _set()
 
+    def test_closing_socket_on_unexpected_closed_error(self):
+        client = self.make_client([
+            b'VALUE ',
+            MemcacheUnexpectedCloseError("foo bar"),
+        ])
+
+        def _set():
+            client.set(b'key', b'value', noreply=False)
+
+        with pytest.raises(MemcacheUnexpectedCloseError):
+            _set()
+
+        assert client.sock is None
+
     def test_set_unknown_error(self):
         client = self.make_client([b'foobarbaz\r\n'])