(improvement) perf: remove copies on the read path (us improvements - x1.5-3.7 speedup!)

cassandra/bytesio.pxd

@@ -17,4 +17,5 @@ cdef class BytesIOReader:
     cdef char *buf_ptr
     cdef Py_ssize_t pos
     cdef Py_ssize_t size
+    cdef Py_ssize_t _initial_offset
     cdef char *read(self, Py_ssize_t n = ?) except NULL

cassandra/bytesio.pyx

@@ -16,12 +16,18 @@ cdef class BytesIOReader:
     """
     This class provides efficient support for reading bytes from a 'bytes' buffer,
     by returning char * values directly without allocating intermediate objects.
+
+    An optional offset allows reading from the middle of an existing buffer,
+    avoiding a copy when only a suffix of the bytes is needed.
     """
 
-    def __init__(self, bytes buf):
+    def __init__(self, bytes buf, Py_ssize_t offset=0):
+        if offset < 0 or offset > len(buf):
+            raise ValueError("offset %d out of range for buffer of length %d" % (offset, len(buf)))
         self.buf = buf
-        self.size = len(buf)
-        self.buf_ptr = self.buf
+        self._initial_offset = offset
+        self.size = len(buf) - offset
+        self.buf_ptr = <char*>self.buf + offset
 
     cdef char *read(self, Py_ssize_t n = -1) except NULL:
         """Read at most size bytes from the file

cassandra/cluster.py

@@ -926,6 +926,17 @@ def default_retry_policy(self, policy):
     To try with your own workload, set ``sockopts = [(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)]``
     """
 
+    in_buffer_size = None
+    """
+    An optional override (in bytes) for the per-connection socket receive
+    buffer size (see :attr:`.Connection.in_buffer_size`, which defaults to
+    16 KiB). A larger value reduces the number of ``recv()`` calls needed to
+    read large result sets, at the cost of additional steady-state memory
+    per connection -- a cost that is multiplied by the number of connections
+    in the pool, so it may matter at scale. Leave unset to use the
+    connection class's default.
+    """
+
     max_schema_agreement_wait = 10
     """
     The maximum duration (in seconds) that the driver will wait for schema
@@ -1233,6 +1244,7 @@ def __init__(self,
                  connection_class=None,
                  ssl_options=None,
                  sockopts=None,
+                 in_buffer_size=None,
                  cql_version=None,
                  protocol_version=_NOT_SET,
                  executor_threads=2,
@@ -1518,6 +1530,7 @@ def __init__(self,
         self.ssl_options = ssl_options
         self.ssl_context = ssl_context
         self.sockopts = sockopts
+        self.in_buffer_size = in_buffer_size
         self.cql_version = cql_version
         self.max_schema_agreement_wait = max_schema_agreement_wait
         self.control_connection_timeout = control_connection_timeout
@@ -1762,6 +1775,7 @@ def _make_connection_kwargs(self, endpoint, kwargs_dict):
         kwargs_dict.setdefault('port', self.port)
         kwargs_dict.setdefault('compression', self.compression)
         kwargs_dict.setdefault('sockopts', self.sockopts)
+        kwargs_dict.setdefault('in_buffer_size', self.in_buffer_size)
         kwargs_dict.setdefault('ssl_options', self.ssl_options)
         kwargs_dict.setdefault('ssl_context', self.ssl_context)
         kwargs_dict.setdefault('cql_version', self.cql_version)

cassandra/connection.py

@@ -746,14 +746,32 @@ def readable_io_bytes(self):
     def readable_cql_frame_bytes(self):
         return self.cql_frame_buffer.tell()
 
+    @staticmethod
+    def _reset_buffer(buf):
+        """
+        Reset a BytesIO buffer by discarding consumed data.
+
+        Uses ``getbuffer()[pos:]`` (a zero-copy memoryview slice) instead of
+        ``.read()`` which would first allocate an intermediate ``bytes`` object.
+        The ``BytesIO()`` constructor still copies the data into its own backing
+        store either way, so the net saving is one temporary ``bytes`` allocation
+        on the hot receive path.
+        """
+        pos = buf.tell()
+        view = buf.getbuffer()
+        try:
+            new_buf = io.BytesIO(view[pos:])
+        finally:
+            del view  # release memoryview before any buffer mutation
+        new_buf.seek(0, 2)  # 2 == SEEK_END
+        return new_buf
+
     def reset_io_buffer(self):
-        self._io_buffer = io.BytesIO(self._io_buffer.read())
-        self._io_buffer.seek(0, 2)  # 2 == SEEK_END
+        self._io_buffer = self._reset_buffer(self._io_buffer)
 
     def reset_cql_frame_buffer(self):
         if self.is_checksumming_enabled:
-            self._cql_frame_buffer = io.BytesIO(self._cql_frame_buffer.read())
-            self._cql_frame_buffer.seek(0, 2)  # 2 == SEEK_END
+            self._cql_frame_buffer = self._reset_buffer(self._cql_frame_buffer)
         else:
             self.reset_io_buffer()
 
@@ -787,7 +805,13 @@ class Connection(object):
 
     CALLBACK_ERR_THREAD_THRESHOLD = 100
 
-    in_buffer_size = 4096
+    # 16 KiB recv buffer reduces the number of syscalls when reading
+    # large result sets, at a modest per-connection memory cost (this is
+    # multiplied by the number of connections in the pool, so it can matter
+    # at scale). Override via the ``in_buffer_size`` kwarg, or globally via
+    # :attr:`.Cluster.in_buffer_size`, if the extra per-connection memory is
+    # a concern or if a different value benchmarks better for your workload.
+    in_buffer_size = 16384
     out_buffer_size = 4096
 
     cql_version = None
@@ -880,7 +904,8 @@ def __init__(self, host='127.0.0.1', port=9042, authenticator=None,
                  cql_version=None, protocol_version=ProtocolVersion.MAX_SUPPORTED, is_control_connection=False,
                  user_type_map=None, connect_timeout=None, allow_beta_protocol_version=False, no_compact=False,
                  ssl_context=None, owning_pool=None, shard_id=None, total_shards=None,
-                 on_orphaned_stream_released=None, application_info: Optional[ApplicationInfoBase] = None):
+                 on_orphaned_stream_released=None, application_info: Optional[ApplicationInfoBase] = None,
+                 in_buffer_size=None):
         # TODO next major rename host to endpoint and remove port kwarg.
         self.endpoint = host if isinstance(host, EndPoint) else DefaultEndPoint(host, port)
 
@@ -889,6 +914,8 @@ def __init__(self, host='127.0.0.1', port=9042, authenticator=None,
         self.ssl_context = ssl_context
         self.sockopts = sockopts
         self.compression = compression
+        if in_buffer_size is not None:
+            self.in_buffer_size = in_buffer_size
         self.cql_version = cql_version
         self.protocol_version = protocol_version
         self.is_control_connection = is_control_connection
@@ -1313,19 +1340,23 @@ def control_conn_disposed(self):
 
     @defunct_on_error
     def _read_frame_header(self):
-        buf = self._io_buffer.cql_frame_buffer.getvalue()
-        pos = len(buf)
+        cql_buf = self._io_buffer.cql_frame_buffer
+        pos = cql_buf.tell()
         if pos:
-            version = buf[0] & PROTOCOL_VERSION_MASK
-            if version not in ProtocolVersion.SUPPORTED_VERSIONS:
-                raise ProtocolError("This version of the driver does not support protocol version %d" % version)
-            # this frame header struct is everything after the version byte
-            header_size = frame_header_v3.size + 1
-            if pos >= header_size:
-                flags, stream, op, body_len = frame_header_v3.unpack_from(buf, 1)
-                if body_len < 0:
-                    raise ProtocolError("Received negative body length: %r" % body_len)
-                self._current_frame = _Frame(version, flags, stream, op, header_size, body_len + header_size)
+            buf = cql_buf.getbuffer()
+            try:
+                version = buf[0] & PROTOCOL_VERSION_MASK
+                if version not in ProtocolVersion.SUPPORTED_VERSIONS:
+                    raise ProtocolError("This version of the driver does not support protocol version %d" % version)
+                # this frame header struct is everything after the version byte
+                header_size = frame_header_v3.size + 1
+                if pos >= header_size:
+                    flags, stream, op, body_len = frame_header_v3.unpack_from(buf, 1)
+                    if body_len < 0:
+                        raise ProtocolError("Received negative body length: %r" % body_len)
+                    self._current_frame = _Frame(version, flags, stream, op, header_size, body_len + header_size)
+            finally:
+                del buf  # release memoryview before any buffer mutation
         return pos
 
     @defunct_on_error
@@ -1379,8 +1410,16 @@ def process_io_buffer(self):
                 return
             else:
                 frame = self._current_frame
-                self._io_buffer.cql_frame_buffer.seek(frame.body_offset)
-                msg = self._io_buffer.cql_frame_buffer.read(frame.end_pos - frame.body_offset)
+                # Use memoryview to avoid intermediate allocation, then
+                # convert to bytes.  Explicitly release the memoryview
+                # before any buffer mutation (seek / reset).
+                cql_buf = self._io_buffer.cql_frame_buffer
+                buf = cql_buf.getbuffer()
+                try:
+                    msg = bytes(buf[frame.body_offset:frame.end_pos])
+                finally:
+                    del buf  # release memoryview before buffer mutation
+                cql_buf.seek(frame.end_pos)
                 self.process_msg(frame, msg)
                 self._io_buffer.reset_cql_frame_buffer()
                 self._current_frame = None

cassandra/protocol.py

@@ -56,6 +56,54 @@ class NotSupportedError(Exception):
 class InternalError(Exception):
     pass
 
+
+class BytesReader:
+    """
+    Lightweight reader for bytes data without BytesIO overhead.
+    Provides the same read() interface but operates directly on a
+    bytes object, avoiding internal buffer copies.
+
+    Unlike io.BytesIO.read(n), read(n) raises EOFError when fewer than
+    n bytes remain.  This is intentional: protocol parsing should fail
+    fast on truncated or malformed frames rather than silently returning
+    short data.
+
+    A plain, already-immutable ``bytes`` object is stored as-is (no copy):
+    downstream zero-copy consumers (e.g. the Cython row parser) keep a
+    reference to it via ``remaining_buffer()``, which is only safe because
+    ``bytes`` can never be mutated in place. Any other buffer-like input
+    (``memoryview``, ``bytearray``, ...) is materialized into a fresh
+    ``bytes`` copy up front, so a caller mutating or recycling its original
+    buffer afterwards can never corrupt or dangle a reader that's still
+    holding onto (or handing out slices of) this data.
+    """
+    __slots__ = ('_data', '_pos', '_size')
+
+    def __init__(self, data):
+        # Only a plain `bytes` object is guaranteed immutable and safe to
+        # alias without copying; anything else (memoryview, bytearray, ...)
+        # is materialized up front.
+        self._data = data if isinstance(data, bytes) else bytes(data)
+        self._pos = 0
+        self._size = len(self._data)
+
+    def read(self, n=-1):
+        if n < 0:
+            result = self._data[self._pos:]
+            self._pos = self._size
+        else:
+            end = self._pos + n
+            if end > self._size:
+                raise EOFError("Cannot read past the end of the buffer")
+            result = self._data[self._pos:end]
+            self._pos = end
+        return result
+
+    def remaining_buffer(self):
+        """Return (underlying_bytes, current_position) for zero-copy handoff."""
+        return self._data, self._pos
+
+
 ColumnMetadata = namedtuple("ColumnMetadata", ['keyspace_name', 'table_name', 'name', 'type'])
 
 HEADER_DIRECTION_TO_CLIENT = 0x80
@@ -1207,32 +1255,53 @@ def decode_message(cls, protocol_version, protocol_features, user_type_map, stre
             body = decompressor(body)
             flags ^= COMPRESSED_FLAG
 
-        body = io.BytesIO(body)
-        if flags & TRACING_FLAG:
-            trace_id = UUID(bytes=body.read(16))
-            flags ^= TRACING_FLAG
-        else:
-            trace_id = None
-
-        if flags & WARNING_FLAG:
-            warnings = read_stringlist(body)
-            flags ^= WARNING_FLAG
-        else:
-            warnings = None
+        # Use lightweight BytesReader instead of io.BytesIO to avoid buffer copy.
+        #
+        # Unlike io.BytesIO, BytesReader.read(n) raises EOFError (rather than
+        # silently returning a short read) when the body is truncated. EOFError
+        # is not part of the driver's public decode-error vocabulary, so it is
+        # caught here -- at the decode_message() API boundary -- and translated
+        # into the canonical ProtocolError, the same exception type already
+        # used for other malformed-frame conditions (see _read_frame_header in
+        # connection.py). This keeps BytesReader itself generic/reusable while
+        # ensuring callers of this public entry point only ever see driver
+        # exception types for decode failures.
+        body = BytesReader(body)
+        try:
+            if flags & TRACING_FLAG:
+                trace_id = UUID(bytes=body.read(16))
+                flags ^= TRACING_FLAG
+            else:
+                trace_id = None
 
-        if flags & CUSTOM_PAYLOAD_FLAG:
-            custom_payload = read_bytesmap(body)
-            flags ^= CUSTOM_PAYLOAD_FLAG
-        else:
-            custom_payload = None
+            if flags & WARNING_FLAG:
+                warnings = read_stringlist(body)
+                flags ^= WARNING_FLAG
+            else:
+                warnings = None
 
-        flags &= USE_BETA_MASK  # will only be set if we asserted it in connection estabishment
+            if flags & CUSTOM_PAYLOAD_FLAG:
+                custom_payload = read_bytesmap(body)
+                flags ^= CUSTOM_PAYLOAD_FLAG
+            else:
+                custom_payload = None
 
-        if flags:
-            log.warning("Unknown protocol flags set: %02x. May cause problems.", flags)
+            flags &= USE_BETA_MASK  # will only be set if we asserted it in connection estabishment
 
-        msg_class = cls.message_types_by_opcode[opcode]
-        msg = msg_class.recv_body(body, protocol_version, protocol_features, user_type_map, result_metadata, cls.column_encryption_policy)
+            if flags:
+                log.warning("Unknown protocol flags set: %02x. May cause problems.", flags)
+
+            msg_class = cls.message_types_by_opcode[opcode]
+            msg = msg_class.recv_body(body, protocol_version, protocol_features, user_type_map, result_metadata, cls.column_encryption_policy)
+        except EOFError as exc:
+            # Local import to avoid a circular import at module load time:
+            # cassandra.connection imports from cassandra.protocol, so
+            # cassandra.protocol cannot import cassandra.connection at
+            # module scope.
+            from cassandra.connection import ProtocolError
+            raise ProtocolError(
+                "Ran out of data decoding %r message body (opcode %r): %s" %
+                (cls.message_types_by_opcode.get(opcode, opcode), opcode, exc)) from exc
         msg.stream_id = stream_id
         msg.trace_id = trace_id
         msg.custom_payload = custom_payload

cassandra/row_parser.pyx

@@ -35,13 +35,19 @@ def make_recv_results_rows(ColumnParser colparser):
         desc = ParseDesc(self.column_names, self.column_types, column_encryption_policy,
                         [ColDesc(md[0], md[1], md[2]) for md in column_metadata],
                         make_deserializers(self.column_types), protocol_version)
-        reader = BytesIOReader(f.read())
+        # Zero-copy handoff: reuse the underlying bytes buffer at its current
+        # position instead of copying via f.read().
+        if hasattr(f, 'remaining_buffer'):
+            buf_data, buf_offset = f.remaining_buffer()
+            reader = BytesIOReader(buf_data, buf_offset)
+        else:
+            reader = BytesIOReader(f.read())
         try:
             self.parsed_rows = colparser.parse_rows(reader, desc)
         except Exception as e:
             # Use explicitly the TupleRowParser to display better error messages for column decoding failures
             rowparser = TupleRowParser()
-            reader.buf_ptr = reader.buf
+            reader.buf_ptr = <char*>reader.buf + reader._initial_offset
             reader.pos = 0
             rowcount = read_int(reader)
             for i in range(rowcount):