Configurable connection buffers capacity, descriptive error message

Author: liutec

neo4j/bolt/connection.py

@@ -42,7 +42,7 @@
 from neo4j.packstream import Packer, Unpacker
 from neo4j.util import import_best as _import_best
 from time import clock
-from neo4j.config import default_config, INFINITE, TRUST_ON_FIRST_USE
+from neo4j.config import get_config, INFINITE, TRUST_ON_FIRST_USE
 
 ChunkedInputBuffer = _import_best("neo4j.bolt._io", "neo4j.bolt.io").ChunkedInputBuffer
 ChunkedOutputBuffer = _import_best("neo4j.bolt._io", "neo4j.bolt.io").ChunkedOutputBuffer
@@ -170,16 +170,19 @@ def __init__(self, address, sock, error_handler, **config):
         self.socket = sock
         self.error_handler = error_handler
         self.server = ServerInfo(SocketAddress.from_socket(sock))
-        self.input_buffer = ChunkedInputBuffer()
-        self.output_buffer = ChunkedOutputBuffer()
+        self.input_buffer = ChunkedInputBuffer(get_config(config, "input_buffer_capacity"))
+        self.output_buffer = ChunkedOutputBuffer(
+            get_config(config, "output_buffer_capacity"),
+            get_config(config, "output_buffer_max_chunk_size")
+        )
         self.packer = Packer(self.output_buffer)
         self.unpacker = Unpacker()
         self.responses = deque()
-        self._max_connection_lifetime = config.get("max_connection_lifetime", default_config["max_connection_lifetime"])
+        self._max_connection_lifetime = get_config(config, "max_connection_lifetime")
         self._creation_timestamp = clock()
 
         # Determine the user agent and ensure it is a Unicode value
-        user_agent = config.get("user_agent", default_config["user_agent"])
+        user_agent = get_config(config, "user_agent")
         if isinstance(user_agent, bytes):
             user_agent = user_agent.decode("UTF-8")
         self.user_agent = user_agent
@@ -328,14 +331,20 @@ def _fetch(self):
         return len(details), 1
 
     def _receive(self):
+        err_msg = None
         try:
             received = self.input_buffer.receive_message(self.socket, 8192)
+        except BufferError:
+            err_msg = "Overflow. Increase \"input_buffer_capacity\" for connection {!r}"
+            received = False
         except SocketError:
             received = False
         if not received:
             self._defunct = True
             self.close()
-            raise self.Error("Failed to read from defunct connection {!r}".format(self.server.address))
+            if not err_msg:
+                err_msg = "Failed to read from defunct connection {!r}"
+            raise self.Error(err_msg.format(self.server.address))
 
     def _unpack(self):
         unpacker = self.unpacker
@@ -405,8 +414,8 @@ def __init__(self, connector, connection_error_handler, **config):
         self.connections = {}
         self.lock = RLock()
         self.cond = Condition(self.lock)
-        self._max_connection_pool_size = config.get("max_connection_pool_size", default_config["max_connection_pool_size"])
-        self._connection_acquisition_timeout = config.get("connection_acquisition_timeout", default_config["connection_acquisition_timeout"])
+        self._max_connection_pool_size = get_config(config, "max_connection_pool_size")
+        self._connection_acquisition_timeout = get_config(config, "connection_acquisition_timeout")
 
     def __enter__(self):
         return self
@@ -544,10 +553,10 @@ def connect(address, ssl_context=None, error_handler=None, **config):
         else:
             raise ValueError("Unsupported address {!r}".format(address))
         t = s.gettimeout()
-        s.settimeout(config.get("connection_timeout", default_config["connection_timeout"]))
+        s.settimeout(get_config(config, "connection_timeout"))
         s.connect(address)
         s.settimeout(t)
-        s.setsockopt(SOL_SOCKET, SO_KEEPALIVE, 1 if config.get("keep_alive", default_config["keep_alive"]) else 0)
+        s.setsockopt(SOL_SOCKET, SO_KEEPALIVE, 1 if get_config(config, "keep_alive") else 0)
     except SocketTimeout:
         if s:
             try:
@@ -586,7 +595,7 @@ def connect(address, ssl_context=None, error_handler=None, **config):
                 s.close()
                 raise ProtocolError("When using a secure socket, the server should always "
                                     "provide a certificate")
-            trust = config.get("trust", default_config["trust"])
+            trust = get_config(config, "trust")
             if trust == TRUST_ON_FIRST_USE:
                 from neo4j.bolt.cert import PersonalCertificateStore
                 store = PersonalCertificateStore()

neo4j/config.py

@@ -64,7 +64,16 @@
     "connection_timeout": DEFAULT_CONNECTION_TIMEOUT,
     "keep_alive": True,
 
+    # Buffer capacity:
+    "input_buffer_capacity": 524288,
+    "output_buffer_capacity": 1048576,
+    "output_buffer_max_chunk_size": 16384,
+
     # Routing settings:
     "max_retry_time": DEFAULT_MAX_RETRY_TIME,
     "load_balancing_strategy": DEFAULT_LOAD_BALANCING_STRATEGY
 }
+
+def get_config(config, key, if_none=None):
+    return config.get(key, default_config.get(key)) or if_none
+

neo4j/v1/api.py

@@ -28,7 +28,7 @@
 from neo4j.exceptions import ProtocolError, ServiceUnavailable
 from neo4j.compat import urlparse
 from neo4j.exceptions import CypherError, TransientError
-from neo4j.config import default_config
+from neo4j.config import get_config
 
 from .exceptions import DriverError, SessionError, SessionExpired, TransactionError
 
@@ -144,7 +144,7 @@ class Driver(object):
 
     def __init__(self, pool, **config):
         self._pool = pool
-        self._max_retry_time = config.get("max_retry_time", default_config["max_retry_time"])
+        self._max_retry_time = get_config(config, "max_retry_time")
 
     def __del__(self):
         self.close()
@@ -227,23 +227,18 @@ class Session(object):
     # :class:`.Transaction` should be carried out.
     _bookmarks = ()
 
-    # Default maximum time to keep retrying failed transactions.
-    _max_retry_time = default_config["max_retry_time"]
-
     _closed = False
 
     def __init__(self, acquirer, access_mode, **parameters):
         self._acquirer = acquirer
         self._default_access_mode = access_mode
-        for key, value in parameters.items():
-            if key == "bookmark":
-                self._bookmarks = [value] if value else []
-            elif key == "bookmarks":
-                self._bookmarks = value or []
-            elif key == "max_retry_time":
-                self._max_retry_time = value
-            else:
-                pass  # for compatibility
+        # Default maximum time to keep retrying failed transactions.
+        self._max_retry_time = get_config(parameters, "max_retry_time")
+        if "bookmark" in parameters:
+            value = parameters["bookmark"]
+            self._bookmarks = [value] if value else []
+        elif "bookmarks" in parameters:
+            self._bookmarks = parameters["bookmarks"] or []
 
     def __del__(self):
         try:

neo4j/v1/routing.py

@@ -31,7 +31,7 @@
 from neo4j.v1.exceptions import SessionExpired
 from neo4j.v1.security import SecurityPlan
 from neo4j.v1.session import BoltSession
-from neo4j.config import default_config, LOAD_BALANCING_STRATEGY_LEAST_CONNECTED, LOAD_BALANCING_STRATEGY_ROUND_ROBIN
+from neo4j.config import get_config, LOAD_BALANCING_STRATEGY_LEAST_CONNECTED, LOAD_BALANCING_STRATEGY_ROUND_ROBIN
 
 
 class OrderedSet(MutableSet):
@@ -163,7 +163,7 @@ class LoadBalancingStrategy(object):
 
     @classmethod
     def build(cls, connection_pool, **config):
-        load_balancing_strategy = config.get("load_balancing_strategy", default_config["load_balancing_strategy"])
+        load_balancing_strategy = get_config(config, "load_balancing_strategy")
         if load_balancing_strategy == LOAD_BALANCING_STRATEGY_LEAST_CONNECTED:
             return LeastConnectedLoadBalancingStrategy(connection_pool)
         elif load_balancing_strategy == LOAD_BALANCING_STRATEGY_ROUND_ROBIN:

neo4j/v1/security.py

@@ -21,7 +21,7 @@
 from warnings import warn
 
 from neo4j.compat.ssl import SSL_AVAILABLE, SSLContext, PROTOCOL_SSLv23, OP_NO_SSLv2, CERT_REQUIRED
-from neo4j.config import default_config, TRUST_ALL_CERTIFICATES, TRUST_CUSTOM_CA_SIGNED_CERTIFICATES, TRUST_ON_FIRST_USE, TRUST_SIGNED_CERTIFICATES, TRUST_SYSTEM_CA_SIGNED_CERTIFICATES
+from neo4j.config import get_config, TRUST_ALL_CERTIFICATES, TRUST_CUSTOM_CA_SIGNED_CERTIFICATES, TRUST_ON_FIRST_USE, TRUST_SIGNED_CERTIFICATES, TRUST_SYSTEM_CA_SIGNED_CERTIFICATES
 
 ENCRYPTION_OFF = 0
 ENCRYPTION_ON = 1
@@ -49,10 +49,10 @@ class SecurityPlan(object):
 
     @classmethod
     def build(cls, **config):
-        encrypted = config.get("encrypted", default_config["encrypted"])
+        encrypted = get_config(config, "encrypted")
         if encrypted is None:
             encrypted = _encryption_default()
-        trust = config.get("trust", default_config["trust"])
+        trust = get_config(config, "trust")
         if encrypted:
             if not SSL_AVAILABLE:
                 raise RuntimeError("Bolt over TLS is only available in Python 2.7.9+ and "