Send RESET when putting a connection back into the pool

Author: robsdedude

neo4j/io/__init__.py

@@ -83,6 +83,7 @@
     ConfigurationError,
     DriverError,
     IncompleteCommit,
+    Neo4jError,
     ReadServiceUnavailable,
     ServiceUnavailable,
     SessionExpired,
@@ -121,6 +122,9 @@ class Bolt(abc.ABC):
 
     PROTOCOL_VERSION = None
 
+    # flag if connection needs RESET to go back to READY state
+    _is_reset = True
+
     # The socket
     in_use = False
 
@@ -144,7 +148,6 @@ def __init__(self, unresolved_address, sock, max_connection_lifetime, *, auth=No
         self.responses = deque()
         self._max_connection_lifetime = max_connection_lifetime
         self._creation_timestamp = perf_counter()
-        self._is_reset = True
         self.routing_context = routing_context
 
         # Determine the user agent
@@ -447,6 +450,10 @@ def rollback(self, **handlers):
         """ Appends a ROLLBACK message to the output queue."""
         pass
 
+    @property
+    def is_reset(self):
+        return self._is_reset
+
     @abc.abstractmethod
     def reset(self):
         """ Appends a RESET message to the outgoing queue, sends it and consumes
@@ -564,23 +571,26 @@ def _set_defunct(self, message, error=None, silent=False):
     def stale(self):
         return (self._stale
                 or (0 <= self._max_connection_lifetime
-                    <= perf_counter()- self._creation_timestamp))
+                    <= perf_counter() - self._creation_timestamp))
 
     _stale = False
 
     def set_stale(self):
         self._stale = True
 
+    @abc.abstractmethod
     def close(self):
         """ Close the connection.
         """
-        raise NotImplementedError
+        pass
 
+    @abc.abstractmethod
     def closed(self):
-        raise NotImplementedError
+        pass
 
+    @abc.abstractmethod
     def defunct(self):
-        raise NotImplementedError
+        pass
 
 
 class IOPool:
@@ -682,6 +692,13 @@ def release(self, *connections):
         """
         with self.lock:
             for connection in connections:
+                if not connection.is_reset:
+                    try:
+                        connection.reset()
+                    except (Neo4jError, DriverError, BoltError) as e:
+                        log.debug(
+                            "Reset on IOPool.release failed: {}".format(e)
+                        )
                 connection.in_use = False
             self.cond.notify_all()
 

neo4j/io/_bolt4.py

@@ -61,18 +61,6 @@ class Bolt4x0(Bolt):
 
     PROTOCOL_VERSION = Version(4, 0)
 
-    # The socket
-    in_use = False
-
-    # The socket
-    _closed = False
-
-    # The socket
-    _defunct = False
-
-    #: The pool of which this connection is a member
-    pool = None
-
     supports_multiple_results = True
 
     supports_multiple_databases = True

neo4j/work/transaction.py

@@ -160,7 +160,7 @@ def rollback(self):
             raise TransactionError("Transaction closed")
         metadata = {}
         try:
-            if not self._connection._is_reset:
+            if not self._connection.is_reset:
                 # DISCARD pending records then do a rollback.
                 self._consume_results()
                 self._connection.rollback(on_success=metadata.update)

testkitbackend/skipped_tests.json

@@ -65,8 +65,6 @@
     "Driver uses custom resolver for each connection, not only initial seeding",
   "stub.retry.TestRetryClustering.test_retry_ForbiddenOnReadOnlyDatabase_ChangingWriter":
     "Test makes assumptions about how verify_connectivity is implemented",
-  "stub.disconnected.SessionRunDisconnected.test_fail_on_reset":
-    "It is not reseting connection when putting back to pool",
   "stub.authorization.AuthorizationTests.test_should_retry_on_auth_expired_on_begin_using_tx_function":
     "Flaky: test requires the driver to contact servers in a specific order",
   "stub.authorization.AuthorizationTestsV3.test_should_retry_on_auth_expired_on_begin_using_tx_function":

tests/unit/io/test_direct.py

@@ -19,7 +19,10 @@
 # limitations under the License.
 
 
-from unittest import TestCase
+from unittest import (
+    mock,
+    TestCase,
+)
 import pytest
 from threading import (
     Condition,
@@ -277,6 +280,26 @@ def test_multithread(self):
             # The pool size is still 5, but all are free
             self.assert_pool_size(address, 0, 5, pool)
 
+    def test_reset_when_released(self):
+        def test(is_reset):
+            with mock.patch(__name__ + ".QuickConnection.is_reset",
+                            new_callable=mock.PropertyMock) as is_reset_mock:
+                with mock.patch(__name__ + ".QuickConnection.reset",
+                                new_callable=mock.MagicMock) as reset_mock:
+                    is_reset_mock.return_value = is_reset
+                    connection = self.pool._acquire(address, timeout=3)
+                    self.assertIsInstance(connection, QuickConnection)
+                    self.assertEqual(is_reset_mock.call_count, 0)
+                    self.assertEqual(reset_mock.call_count, 0)
+                    self.pool.release(connection)
+                    self.assertEqual(is_reset_mock.call_count, 1)
+                    self.assertEqual(reset_mock.call_count, int(not is_reset))
+
+        address = ("127.0.0.1", 7687)
+        for is_reset in (True, False):
+            with self.subTest():
+                test(is_reset)
+
 
 def acquire_release_conn(pool, address, acquired_counter, release_event):
     conn = pool._acquire(address, timeout=3)