Avoid using 2 io buffers when checksumming is not used

aboudreault · aboudreault · commit d992e818edf9 · 2020-10-16T10:08:14.000-04:00
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -5,6 +5,7 @@ Not released
 Features
 --------
 * Ensure the driver can connect when invalid peer hosts are in system.peers (PYTHON-1260)
+* Implement protocol v5 checksumming (PYTHON-1258)
 
 Bug Fixes
 ---------
diff --git a/cassandra/connection.py b/cassandra/connection.py
@@ -27,6 +27,8 @@
 from threading import Thread, Event, RLock, Condition
 import time
 import ssl
+import weakref
+
 
 if 'gevent.monkey' in sys.modules:
     from gevent.queue import Queue, Empty
@@ -610,6 +612,55 @@ def int_from_buf_item(i):
     int_from_buf_item = ord
 
 
+class _ConnectionIOBuffer(object):
+    """
+    Abstraction class to ease the use of the different connection io buffers. With
+    protocol V5 and checksumming, the data is read, validated and copied to another
+    cql frame buffer.
+    """
+    _io_buffer = None
+    _cql_frame_buffer = None
+    _connection = None
+
+    def __init__(self, connection):
+        self._io_buffer = io.BytesIO()
+        self._connection = weakref.proxy(connection)
+
+    @property
+    def io_buffer(self):
+        return self._io_buffer
+
+    @property
+    def cql_frame_buffer(self):
+        return self._cql_frame_buffer if self.is_checksumming_enabled else \
+            self._io_buffer
+
+    def set_checksumming_buffer(self):
+        self.reset_io_buffer()
+        self._cql_frame_buffer = io.BytesIO()
+
+    @property
+    def is_checksumming_enabled(self):
+        return self._connection._is_checksumming_enabled
+
+    def readable_io_bytes(self):
+        return self.io_buffer.tell()
+
+    def readable_cql_frame_bytes(self):
+        return self.cql_frame_buffer.tell()
+
+    def reset_io_buffer(self):
+        self._io_buffer = io.BytesIO(self._io_buffer.read())
+        self._io_buffer.seek(0, 2)  # 2 == SEEK_END
+
+    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
+        else:
+            self.reset_io_buffer()
+
+
 class Connection(object):
 
     CALLBACK_ERR_THREAD_THRESHOLD = 100
@@ -665,8 +716,6 @@ class Connection(object):
 
     allow_beta_protocol_version = False
 
-    _iobuf = None
-    _frame_iobuf = None
     _current_frame = None
 
     _socket = None
@@ -679,6 +728,11 @@ class Connection(object):
 
     _is_checksumming_enabled = False
 
+    @property
+    def _iobuf(self):
+        # backward compatibility, to avoid any change in the reactors
+        return self._io_buffer.io_buffer
+
     def __init__(self, host='127.0.0.1', port=9042, authenticator=None,
                  ssl_options=None, sockopts=None, compression=True,
                  cql_version=None, protocol_version=ProtocolVersion.MAX_SUPPORTED, is_control_connection=False,
@@ -702,8 +756,7 @@ def __init__(self, host='127.0.0.1', port=9042, authenticator=None,
         self.no_compact = no_compact
         self._push_watchers = defaultdict(set)
         self._requests = {}
-        self._iobuf = io.BytesIO()
-        self._frame_iobuf = io.BytesIO()
+        self._io_buffer = _ConnectionIOBuffer(self)
         self._continuous_paging_sessions = {}
         self._socket_writable = True
 
@@ -844,6 +897,12 @@ def _connect_socket(self):
             for args in self.sockopts:
                 self._socket.setsockopt(*args)
 
+    def _enable_checksumming(self):
+        self._io_buffer.set_checksumming_buffer()
+        self._is_checksumming_enabled = True
+        self._segment_codec = segment_codec_lz4 if self.compressor else segment_codec_no_compression
+        log.debug("Enabling protocol checksumming on connection (%s).", id(self))
+
     def close(self):
         raise NotImplementedError()
 
@@ -1032,7 +1091,7 @@ def control_conn_disposed(self):
 
     @defunct_on_error
     def _read_frame_header(self):
-        buf = self._frame_iobuf.getvalue()
+        buf = self._io_buffer.cql_frame_buffer.getvalue()
         pos = len(buf)
         if pos:
             version = int_from_buf_item(buf[0]) & PROTOCOL_VERSION_MASK
@@ -1048,28 +1107,19 @@ def _read_frame_header(self):
                 self._current_frame = _Frame(version, flags, stream, op, header_size, body_len + header_size)
         return pos
 
-    def _reset_frame(self):
-        self._frame_iobuf = io.BytesIO(self._frame_iobuf.read())
-        self._frame_iobuf.seek(0, 2)  # 2 == SEEK_END
-        self._current_frame = None
-
-    def _reset_io_buffer(self):
-        self._iobuf = io.BytesIO(self._iobuf.read())
-        self._iobuf.seek(0, 2)  # 2 == SEEK_END
-
     @defunct_on_error
     def _process_segment_buffer(self):
-        readable_bytes = self._iobuf.tell()
+        readable_bytes = self._io_buffer.readable_io_bytes()
         if readable_bytes >= self._segment_codec.header_length_with_crc:
             try:
-                self._iobuf.seek(0)
-                segment_header = self._segment_codec.decode_header(self._iobuf)
+                self._io_buffer.io_buffer.seek(0)
+                segment_header = self._segment_codec.decode_header(self._io_buffer.io_buffer)
                 if readable_bytes >= segment_header.segment_length:
                     segment = self._segment_codec.decode(self._iobuf, segment_header)
-                    self._frame_iobuf.write(segment.payload)
+                    self._io_buffer.cql_frame_buffer.write(segment.payload)
                 else:
                     # not enough data to read the segment
-                    self._iobuf.seek(0, 2)
+                    self._io_buffer.io_buffer.seek(0, 2)
             except CrcException as exc:
                 # re-raise an exception that inherits from ConnectionException
                 raise CrcMismatchException(str(exc), self.endpoint)
@@ -1078,21 +1128,15 @@ def process_io_buffer(self):
         while True:
             if self._is_checksumming_enabled:
                 self._process_segment_buffer()
-            else:
-                # We should probably refactor the IO buffering stuff out of the Connection
-                # class to handle this in a better way. That would make the segment and frame
-                # decoding code clearer.
-                self._frame_iobuf.write(self._iobuf.getvalue())
-
-            self._reset_io_buffer()
+                self._io_buffer.reset_io_buffer()
 
             if not self._current_frame:
                 pos = self._read_frame_header()
             else:
-                pos = self._frame_iobuf.tell()
+                pos = self._io_buffer.readable_cql_frame_bytes()
 
             if not self._current_frame or pos < self._current_frame.end_pos:
-                if self._is_checksumming_enabled and self._iobuf.tell():
+                if self._is_checksumming_enabled and self._io_buffer.readable_io_bytes():
                     # We have a multi-segments message and we need to read more
                     # data to complete the current cql frame
                     continue
@@ -1103,10 +1147,11 @@ def process_io_buffer(self):
                 return
             else:
                 frame = self._current_frame
-                self._frame_iobuf.seek(frame.body_offset)
-                msg = self._frame_iobuf.read(frame.end_pos - frame.body_offset)
+                self._io_buffer.cql_frame_buffer.seek(frame.body_offset)
+                msg = self._io_buffer.cql_frame_buffer.read(frame.end_pos - frame.body_offset)
                 self.process_msg(frame, msg)
-                self._reset_frame()
+                self._io_buffer.reset_cql_frame_buffer()
+                self._current_frame = None
 
     @defunct_on_error
     def process_msg(self, header, body):
@@ -1287,9 +1332,7 @@ def _handle_startup_response(self, startup_response, did_authenticate=False):
                 self.compressor = self._compressor
 
             if ProtocolVersion.has_checksumming_support(self.protocol_version):
-                self._is_checksumming_enabled = True
-                self._segment_codec = segment_codec_lz4 if self.compressor else segment_codec_no_compression
-                log.debug("Enabling protocol checksumming on connection (%s).", id(self))
+                self._enable_checksumming()
 
             self.connected_event.set()
         elif isinstance(startup_response, AuthenticateMessage):
diff --git a/tests/integration/standard/test_custom_protocol_handler.py b/tests/integration/standard/test_custom_protocol_handler.py
@@ -25,7 +25,7 @@
 
 from tests.integration import use_singledc, drop_keyspace_shutdown_cluster, \
     greaterthanorequalcass30, execute_with_long_wait_retry, greaterthanorequaldse51, greaterthanorequalcass3_10, \
-    greaterthanorequalcass31, TestCluster
+    TestCluster, greaterthanorequalcass40, requirecassandra
 from tests.integration.datatype_utils import update_datatypes, PRIMITIVE_DATATYPES
 from tests.integration.standard.utils import create_table_with_all_types, get_all_primitive_params
 from six import binary_type
@@ -124,7 +124,8 @@ def test_custom_raw_row_results_all_types(self):
         self.assertEqual(len(CustomResultMessageTracked.checked_rev_row_set), len(PRIMITIVE_DATATYPES)-1)
         cluster.shutdown()
 
-    @greaterthanorequalcass31
+    @requirecassandra
+    @greaterthanorequalcass40
     def test_protocol_divergence_v5_fail_by_continuous_paging(self):
         """
         Test to validate that V5 and DSE_V1 diverge. ContinuousPagingOptions is not supported by V5
@@ -170,7 +171,8 @@ def test_protocol_divergence_v4_fail_by_flag_uses_int(self):
         self._protocol_divergence_fail_by_flag_uses_int(ProtocolVersion.V4, uses_int_query_flag=False,
                                                         int_flag=True)
 
-    @greaterthanorequalcass3_10
+    @requirecassandra
+    @greaterthanorequalcass40
     def test_protocol_v5_uses_flag_int(self):
         """
         Test to validate that the _PAGE_SIZE_FLAG is treated correctly using write_uint for V5
@@ -196,7 +198,8 @@ def test_protocol_dsev1_uses_flag_int(self):
         self._protocol_divergence_fail_by_flag_uses_int(ProtocolVersion.DSE_V1, uses_int_query_flag=True,
                                                         int_flag=True)
 
-    @greaterthanorequalcass3_10
+    @requirecassandra
+    @greaterthanorequalcass40
     def test_protocol_divergence_v5_fail_by_flag_uses_int(self):
         """
         Test to validate that the _PAGE_SIZE_FLAG is treated correctly using write_uint for V5
diff --git a/tests/integration/standard/test_query.py b/tests/integration/standard/test_query.py
@@ -28,7 +28,7 @@
 from cassandra.policies import HostDistance, RoundRobinPolicy, WhiteListRoundRobinPolicy
 from tests.integration import use_singledc, PROTOCOL_VERSION, BasicSharedKeyspaceUnitTestCase, \
     greaterthanprotocolv3, MockLoggingHandler, get_supported_protocol_versions, local, get_cluster, setup_keyspace, \
-    USE_CASS_EXTERNAL, greaterthanorequalcass40, DSE_VERSION, TestCluster
+    USE_CASS_EXTERNAL, greaterthanorequalcass40, DSE_VERSION, TestCluster, requirecassandra
 from tests import notwindows
 from tests.integration import greaterthanorequalcass30, get_node
 
@@ -1408,6 +1408,8 @@ def test_setting_keyspace(self):
         """
         self._check_set_keyspace_in_statement(self.session)
 
+    @requirecassandra
+    @greaterthanorequalcass40
     def test_setting_keyspace_and_session(self):
         """
         Test we can still send the keyspace independently even the session
diff --git a/tests/unit/io/test_twistedreactor.py b/tests/unit/io/test_twistedreactor.py
@@ -148,12 +148,12 @@ def test_handle_read__incomplete(self):
         # incomplete header
         self.obj_ut._iobuf.write(b'\x84\x00\x00\x00\x00')
         self.obj_ut.handle_read()
-        self.assertEqual(self.obj_ut._frame_iobuf.getvalue(), b'\x84\x00\x00\x00\x00')
+        self.assertEqual(self.obj_ut._io_buffer.cql_frame_buffer.getvalue(), b'\x84\x00\x00\x00\x00')
 
         # full header, but incomplete body
         self.obj_ut._iobuf.write(b'\x00\x00\x00\x15')
         self.obj_ut.handle_read()
-        self.assertEqual(self.obj_ut._frame_iobuf.getvalue(),
+        self.assertEqual(self.obj_ut._io_buffer.cql_frame_buffer.getvalue(),
                          b'\x84\x00\x00\x00\x00\x00\x00\x00\x15')
         self.assertEqual(self.obj_ut._current_frame.end_pos, 30)
 
@@ -174,7 +174,7 @@ def test_handle_read__fullmessage(self):
         self.obj_ut._iobuf.write(
             b'\x84\x01\x00\x02\x03\x00\x00\x00\x15' + body + extra)
         self.obj_ut.handle_read()
-        self.assertEqual(self.obj_ut._frame_iobuf.getvalue(), extra)
+        self.assertEqual(self.obj_ut._io_buffer.cql_frame_buffer.getvalue(), extra)
         self.obj_ut.process_msg.assert_called_with(
             _Frame(version=4, flags=1, stream=2, opcode=3, body_offset=9, end_pos=9 + len(body)), body)
 
diff --git a/tests/unit/io/utils.py b/tests/unit/io/utils.py
@@ -309,14 +309,14 @@ def chunk(size):
 
         for message, expected_size in messages:
             message_chunks = message
-            c._iobuf = io.BytesIO()
+            c._io_buffer._io_buffer = io.BytesIO()
             c.process_io_buffer.reset_mock()
             c.handle_read(*self.null_handle_function_args)
-            c._iobuf.seek(0, os.SEEK_END)
+            c._io_buffer.io_buffer.seek(0, os.SEEK_END)
 
             # Ensure the message size is the good one and that the
             # message has been processed if it is non-empty
-            self.assertEqual(c._iobuf.tell(), expected_size)
+            self.assertEqual(c._io_buffer.io_buffer.tell(), expected_size)
             if expected_size == 0:
                 c.process_io_buffer.assert_not_called()
             else:
@@ -435,11 +435,11 @@ def test_partial_header_read(self):
 
         self.get_socket(c).recv.return_value = message[0:1]
         c.handle_read(*self.null_handle_function_args)
-        self.assertEqual(c._frame_iobuf.getvalue(), message[0:1])
+        self.assertEqual(c._io_buffer.cql_frame_buffer.getvalue(), message[0:1])
 
         self.get_socket(c).recv.return_value = message[1:]
         c.handle_read(*self.null_handle_function_args)
-        self.assertEqual(six.binary_type(), c._iobuf.getvalue())
+        self.assertEqual(six.binary_type(), c._io_buffer.io_buffer.getvalue())
 
         # let it write out a StartupMessage
         c.handle_write(*self.null_handle_function_args)
@@ -461,12 +461,12 @@ def test_partial_message_read(self):
         # read in the first nine bytes
         self.get_socket(c).recv.return_value = message[:9]
         c.handle_read(*self.null_handle_function_args)
-        self.assertEqual(c._frame_iobuf.getvalue(), message[:9])
+        self.assertEqual(c._io_buffer.cql_frame_buffer.getvalue(), message[:9])
 
         # ... then read in the rest
         self.get_socket(c).recv.return_value = message[9:]
         c.handle_read(*self.null_handle_function_args)
-        self.assertEqual(six.binary_type(), c._iobuf.getvalue())
+        self.assertEqual(six.binary_type(), c._io_buffer.io_buffer.getvalue())
 
         # let it write out a StartupMessage
         c.handle_write(*self.null_handle_function_args)
@@ -501,7 +501,7 @@ def test_mixed_message_and_buffer_sizes(self):
 
             for i in range(1, 15):
                 c.process_io_buffer.reset_mock()
-                c._iobuf = io.BytesIO()
+                c._io_buffer._io_buffer = io.BytesIO()
                 message = io.BytesIO(six.b('a') * (2**i))
 
                 def recv_side_effect(*args):
@@ -511,7 +511,7 @@ def recv_side_effect(*args):
 
                 self.get_socket(c).recv.side_effect = recv_side_effect
                 c.handle_read(*self.null_handle_function_args)
-                if c._iobuf.tell():
+                if c._io_buffer.io_buffer.tell():
                     c.process_io_buffer.assert_called_once()
                 else:
                     c.process_io_buffer.assert_not_called()
diff --git a/tests/unit/test_connection.py b/tests/unit/test_connection.py
@@ -100,7 +100,7 @@ def test_bad_protocol_version(self, *args):
         header = self.make_header_prefix(SupportedMessage, version=0x7f)
         options = self.make_options_body()
         message = self.make_msg(header, options)
-        c._iobuf = BytesIO()
+        c._iobuf._io_buffer = BytesIO()
         c._iobuf.write(message)
         c.process_io_buffer()
 
@@ -117,7 +117,7 @@ def test_negative_body_length(self, *args):
         # read in a SupportedMessage response
         header = self.make_header_prefix(SupportedMessage)
         message = header + int32_pack(-13)
-        c._iobuf = BytesIO()
+        c._iobuf._io_buffer = BytesIO()
         c._iobuf.write(message)
         c.process_io_buffer()
 
diff --git a/tests/unit/test_segment.py b/tests/unit/test_segment.py
diff --git a/tox.ini b/tox.ini
