Merge pull request #333 from 4383/force-closing-socket

2 files changed, 46 insertions, 12 deletions

diff --git a/pymemcache/client/base.py b/pymemcache/client/base.py
index ad9ed8c..64f65d7 100644
--- a/pymemcache/client/base.py
+++ b/pymemcache/client/base.py
@@ -43,6 +43,13 @@ SOCKET_KEEPALIVE_SUPPORTED_SYSTEM = {
     'Linux',
 }
 
+STORE_RESULTS_VALUE = {
+    b'STORED': True,
+    b'NOT_STORED': False,
+    b'NOT_FOUND':  None,
+    b'EXISTS': False
+}
+
 
 # Some of the values returned by the "stats" command
 # need mapping into native Python types
@@ -986,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))
 
@@ -1012,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
@@ -1084,19 +1101,17 @@ 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]:
-                    if line == b'STORED':
-                        results[key] = True
-                    if line == b'NOT_STORED':
-                        results[key] = False
-                    if line == b'NOT_FOUND':
-                        results[key] = None
-                    if line == b'EXISTS':
-                        results[key] = False
+                    results[key] = STORE_RESULTS_VALUE[line]
                 else:
                     raise MemcacheUnknownError(line[:32])
             return results
@@ -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'])