Don't break Inbox when hydration or unpacking fails

Author: robsdedude

neo4j/_async/io/_common.py

@@ -19,7 +19,6 @@
 import asyncio
 import logging
 import socket
-from contextlib import contextmanager
 from struct import pack as struct_pack
 
 from ..._async_compat.util import AsyncUtil
@@ -45,7 +44,7 @@ def __init__(self, sock, on_error, unpacker_cls):
         self._unpacker = unpacker_cls(self._buffer)
         self._broken = False
 
-    async def pop(self, hydration_hooks):
+    async def _buffer_one_chunk(self):
         assert not self._broken
         try:
             chunk_size = 0
@@ -64,20 +63,23 @@ async def pop(self, hydration_hooks):
 
                 if chunk_size == 0:
                     # chunk_size was the end marker for the message
-                    size, tag = self._unpacker.unpack_structure_header()
-                    fields = [self._unpacker.unpack(hydration_hooks)
-                              for _ in range(size)]
-                    # Reset for new message
-                    self._unpacker.reset()
-                    return tag, fields
+                    return
 
         except (OSError, socket.timeout, SocketDeadlineExceeded) as error:
             self._broken = True
             await AsyncUtil.callback(self.on_error, error)
             raise
-        except Exception:
-            self._broken = True
-            raise
+
+    async def pop(self, hydration_hooks):
+        await self._buffer_one_chunk()
+        try:
+            size, tag = self._unpacker.unpack_structure_header()
+            fields = [self._unpacker.unpack(hydration_hooks)
+                      for _ in range(size)]
+            return tag, fields
+        finally:
+            # Reset for new message
+            self._unpacker.reset()
 
 
 class AsyncOutbox:

neo4j/_sync/io/_common.py

@@ -19,7 +19,6 @@
 import asyncio
 import logging
 import socket
-from contextlib import contextmanager
 from struct import pack as struct_pack
 
 from ..._async_compat.util import Util
@@ -45,7 +44,7 @@ def __init__(self, sock, on_error, unpacker_cls):
         self._unpacker = unpacker_cls(self._buffer)
         self._broken = False
 
-    def pop(self, hydration_hooks):
+    def _buffer_one_chunk(self):
         assert not self._broken
         try:
             chunk_size = 0
@@ -64,20 +63,23 @@ def pop(self, hydration_hooks):
 
                 if chunk_size == 0:
                     # chunk_size was the end marker for the message
-                    size, tag = self._unpacker.unpack_structure_header()
-                    fields = [self._unpacker.unpack(hydration_hooks)
-                              for _ in range(size)]
-                    # Reset for new message
-                    self._unpacker.reset()
-                    return tag, fields
+                    return
 
         except (OSError, socket.timeout, SocketDeadlineExceeded) as error:
             self._broken = True
             Util.callback(self.on_error, error)
             raise
-        except Exception:
-            self._broken = True
-            raise
+
+    def pop(self, hydration_hooks):
+        self._buffer_one_chunk()
+        try:
+            size, tag = self._unpacker.unpack_structure_header()
+            fields = [self._unpacker.unpack(hydration_hooks)
+                      for _ in range(size)]
+            return tag, fields
+        finally:
+            # Reset for new message
+            self._unpacker.reset()
 
 
 class Outbox: