refactor: optimize packet writing by pre-allocating header space in payload buffers

Author: rusher

mariadb/impl/client/socket/write_stream.py

@@ -65,50 +65,64 @@ async def write_payload(self, payload: bytes, packet_type: str = "", reset_seque
         Write payload with MariaDB packet framing (async version)
         
         Args:
-            payload: Payload bytes to send
+            payload: Payload bytes with first 4 bytes reserved for header
             packet_type: Packet type for logging (e.g., "COM_QUERY")
             reset_sequence: Whether to reset sequence number before sending
         """
         if reset_sequence:
             self.sequence.set(-1)
 
-        payload_len = len(payload)
-        offset = 0
+        # Payload has 4 bytes reserved at start for header
+        payload_len = len(payload) - 4
+        data_offset = 4  # Data starts after reserved header space
 
         # Handle empty payload - still need to send header
         if payload_len == 0:
             seq = self.sequence.increment_and_get()
-            header = b'\x00\x00\x00' + bytes([seq])
+            # Write header into first 4 bytes
+            payload_buf = bytearray(payload)
+            payload_buf[0:4] = b'\x00\x00\x00' + bytes([seq])
 
             if logger.isEnabledFor(logging.DEBUG):
                 conn_id_str = f"[conn_id={self.connection_id}]" if self.connection_id >= 0 else ""
                 packet_type_str = f" {packet_type}" if packet_type else ""
-                logger.debug(hex_dump(header, f"SEND async: {conn_id_str}{packet_type_str}"))
+                logger.debug(hex_dump(bytes(payload_buf[0:4]), f"SEND async: {conn_id_str}{packet_type_str}"))
 
-            self.writer.write(header)
+            self.writer.write(payload_buf[0:4])
             await self.writer.drain()
             return
 
+        # Convert to bytearray for in-place header writing
+        payload_buf = bytearray(payload)
+        
         # Handle packet splitting for large payloads
-        while offset < payload_len:
-            chunk_size = min(MAX_PACKET_SIZE, payload_len - offset)
+        sent = 0
+        
+        while sent < payload_len:
+            chunk_size = min(MAX_PACKET_SIZE, payload_len - sent)
             seq = self.sequence.increment_and_get()
 
-            # Build header: 3-byte length + 1-byte sequence
-            header = chunk_size.to_bytes(3, 'little') + bytes([seq])
+            # Data for this chunk starts at data_offset + sent
+            chunk_start = data_offset + sent
+            chunk_end = chunk_start + chunk_size
+            
+            # Write header 4 bytes before the chunk data
+            header_pos = chunk_start - 4
+            payload_buf[header_pos] = chunk_size & 0xff
+            payload_buf[header_pos + 1] = (chunk_size >> 8) & 0xff
+            payload_buf[header_pos + 2] = (chunk_size >> 16) & 0xff
+            payload_buf[header_pos + 3] = seq
 
-            # Log if debug enabled (need to build full packet for logging)
+            # Log if debug enabled
             if logger.isEnabledFor(logging.DEBUG):
-                chunk = payload[offset:offset + chunk_size]
-                packet = header + chunk
+                packet = bytes(payload_buf[header_pos:chunk_end])
                 conn_id_str = f"[conn_id={self.connection_id}]" if self.connection_id >= 0 else ""
                 packet_type_str = f" {packet_type}" if packet_type else ""
                 logger.debug(hex_dump(packet, f"SEND async: {conn_id_str}{packet_type_str}"))
 
-            # Send header and chunk separately (more efficient - no concatenation)
-            self.writer.write(header)
-            self.writer.write(payload[offset:offset + chunk_size])
-            offset += chunk_size
+            # Send packet: header + chunk data
+            self.writer.write(payload_buf[header_pos:chunk_end])
+            sent += chunk_size
 
         # Flush all buffered data
         await self.writer.drain()
@@ -124,71 +138,79 @@ async def write_payload(self, payload: bytes, packet_type: str = "", reset_seque
 class SyncWriteStream(BaseWriteStream):
     """Sync write stream implementation using blocking socket operations"""
 
-    def __init__(self, sock: socket.socket, connection_id: int = -1):
-        """
-        Initialize sync write stream
-        
-        Args:
-            sock: Blocking socket
-            connection_id: Connection ID for logging
-        """
-        self.socket: socket.socket = sock
-        super().__init__(connection_id)
+    def __init__(self, socket: socket.socket, connection_id: int = -1):
+        """Initialize write stream with socket"""
+        self.socket = socket
+        self.sequence = MutableInt(-1)
+        self.connection_id = connection_id
+        # Check once if sendmsg is supported (Unix) or if we need sendall (Windows)
+        self.has_sendmsg = hasattr(socket, 'sendmsg')
 
     def write_payload(self, payload: bytes, packet_type: str = "", reset_sequence: bool = True) -> None:
         """
         Write payload with MariaDB packet framing (sync version)
         
         Args:
-            payload: Payload bytes to send
+            payload: Payload bytes with first 4 bytes reserved for header
             packet_type: Packet type for logging (e.g., "COM_QUERY")
             reset_sequence: Whether to reset sequence number before sending
         """
         if reset_sequence:
             self.sequence.set(-1)
 
-        payload_len = len(payload)
-        offset = 0
+        # Payload has 4 bytes reserved at start for header
+        payload_len = len(payload) - 4
+        data_offset = 4  # Data starts after reserved header space
 
         # Handle empty payload - still need to send header
         if payload_len == 0:
             seq = self.sequence.increment_and_get()
-            header = b'\x00\x00\x00' + bytes([seq])
+            # Write header into first 4 bytes
+            payload_buf = bytearray(payload)
+            payload_buf[0:4] = b'\x00\x00\x00' + bytes([seq])
 
             if logger.isEnabledFor(logging.DEBUG):
                 conn_id_str = f"[conn_id={self.connection_id}]" if self.connection_id >= 0 else ""
                 packet_type_str = f" {packet_type}" if packet_type else ""
-                logger.debug(hex_dump(header, f"SEND sync: {conn_id_str}{packet_type_str}"))
+                logger.debug(hex_dump(bytes(payload_buf[0:4]), f"SEND sync: {conn_id_str}{packet_type_str}"))
 
-            self.socket.sendall(header)
+            self.socket.sendall(payload_buf[0:4])
             return
 
+        # Convert to bytearray for in-place header writing
+        payload_buf = bytearray(payload)
+        
         # Handle packet splitting for large payloads
-        while offset < payload_len:
-            chunk_size = min(MAX_PACKET_SIZE, payload_len - offset)
+        sent = 0  # Track how much data we've sent
+        
+        
+        while sent < payload_len:
+            chunk_size = min(MAX_PACKET_SIZE, payload_len - sent)
             seq = self.sequence.increment_and_get()
 
-            # Build header: 3-byte length + 1-byte sequence
-            header = chunk_size.to_bytes(3, 'little') + bytes([seq])
-            chunk = payload[offset:offset + chunk_size]
+            # Data for this chunk starts at data_offset + sent
+            chunk_start = data_offset + sent
+            chunk_end = chunk_start + chunk_size
+            
+            # Write header 4 bytes before the chunk data
+            header_pos = chunk_start - 4
+            
+            payload_buf[header_pos] = chunk_size & 0xff
+            payload_buf[header_pos + 1] = (chunk_size >> 8) & 0xff
+            payload_buf[header_pos + 2] = (chunk_size >> 16) & 0xff
+            payload_buf[header_pos + 3] = seq
 
-            # Log if debug enabled (need full packet for logging)
+            # Log if debug enabled
             if logger.isEnabledFor(logging.DEBUG):
-                packet = header + chunk
+                packet = bytes(payload_buf[header_pos:chunk_end])
                 conn_id_str = f"[conn_id={self.connection_id}]" if self.connection_id >= 0 else ""
                 packet_type_str = f" {packet_type}" if packet_type else ""
                 logger.debug(hex_dump(packet, f"SEND sync: {conn_id_str}{packet_type_str}"))
 
-            # Send header and chunk in a single syscall using scatter-gather I/O
-            # sendmsg() is available on Unix and sends multiple buffers efficiently
-            try:
-                self.socket.sendmsg([header, chunk])
-            except Exception:
-                # Fallback for platforms without sendmsg (e.g., Windows)
-                self.socket.sendall(header)
-                self.socket.sendall(chunk)
+            # Send packet: header + chunk data
+            self.socket.sendall(payload_buf[header_pos:chunk_end])
 
-            offset += chunk_size
+            sent += chunk_size
 
         # If last packet was exactly MAX_PACKET_SIZE, send empty packet to signal end
         if payload_len % MAX_PACKET_SIZE == 0:

mariadb/impl/client/sync_client.py

@@ -64,7 +64,9 @@ def __init__(self, configuration: Configuration) -> None:
         self.sequence: MutableInt = MutableInt(-1)
 
         # Read buffer management
-        self._recv_buf: bytearray = bytearray(8192)
+        self._default_recv_buf: bytearray = bytearray(8192)
+        self._recv_buf: bytearray = self._default_recv_buf
+        
         self._recv_pos = 0
         self._recv_len = 0
 
@@ -79,7 +81,8 @@ def _ensure_space(self, needed):
         ALIGN = 16384 
         if (len(self._recv_buf) - self._recv_len >= needed):
             return
-        self._recv_buf.extend(bytearray((needed + ALIGN - 1) & ~(ALIGN - 1)))
+        self._recv_buf = self._recv_buf + bytearray((needed + ALIGN - 1) & ~(ALIGN - 1))
+                
 
     def _recv_into_buffer(self, size=0):
         """
@@ -97,20 +100,29 @@ def _recv_into_buffer(self, size=0):
 
         # Keep trying to read until we have enough data or there's nothing left
         try:
+            if self.logger.isEnabledFor(logging.DEBUG):
+                self.logger.debug(f"_recv_into_buffer: requesting size={size}, buffer_len={len(self._recv_buf)}, recv_len={self._recv_len}, recv_pos={self._recv_pos}")
+            
             if size == 0:
                 n = self.socket.recv_into(mv[self._recv_len + received:])
+                if self.logger.isEnabledFor(logging.DEBUG):
+                    self.logger.debug(f"_recv_into_buffer: received {n} bytes (no size limit)")
                 if n == 0:
                     raise ConnectionError("Connection reset by peer")
                 return n
             while received < size:
                 n = self.socket.recv_into(mv[self._recv_len + received:], size - received)
+                if self.logger.isEnabledFor(logging.DEBUG):
+                    self.logger.debug(f"_recv_into_buffer: received {n} bytes, total {received + n}/{size}")
                 if n == 0:
                     raise ConnectionError("Connection reset by peer")
                 received += n
             return received
 
         except socket.timeout:
-           raise TimeoutError("Socket recv timed out")
+            if self.logger.isEnabledFor(logging.DEBUG):
+                self.logger.debug(f"_recv_into_buffer: TIMEOUT after receiving {received} bytes (requested {size})")
+            raise TimeoutError("Socket recv timed out")
 
         except ConnectionResetError:
             raise ConnectionError("Connection reset by peer")
@@ -130,6 +142,7 @@ def read_payload(self):
         of the buffer
 
         """
+        from ..debug_utils import hex_dump
 
         # for faster local lookup
         PKT_HDR_SIZE=4
@@ -149,7 +162,7 @@ def read_payload(self):
 
         first_pos = self._recv_pos
         total_size = 0
-        payload_write_pos = None  # Track where to write compacted payload
+        packet_count = 0
 
         while True:
             bytes_in_buffer = self._recv_len - self._recv_pos
@@ -180,31 +193,57 @@ def read_payload(self):
                     continue
 
                 # We have complete packet (header + payload)
-                if payload_write_pos is None:
-                    # First packet - payload starts after first header
-                    payload_write_pos = first_pos + PKT_HDR_SIZE
-                elif self._recv_pos != payload_write_pos:
+                packet_count += 1
+                
+                # Log complete packet with data
+                if self.logger.isEnabledFor(logging.DEBUG):
+                    packet_data = bytes(self._recv_buf[self._recv_pos:self._recv_pos + PKT_HDR_SIZE + packet_length])
+                    conn_id_str = f"[conn_id={self.connection_id}]" if hasattr(self, 'connection_id') and self.connection_id >= 0 else ""
+                    self.logger.debug(hex_dump(packet_data, f"RECV sync: {conn_id_str} packet {packet_count} complete"))
+                
+                if packet_count > 1:
                     # Multi-packet: compact by removing intermediate header
-                    # Move this packet's payload to the write position
-                    payload_start = self._recv_pos + PKT_HDR_SIZE
-                    self._recv_buf[payload_write_pos:payload_write_pos + packet_length] = \
-                        self._recv_buf[payload_start:payload_start + packet_length]
+                    # Move this packet's payload immediately after previous payload
+                    payload_src = self._recv_pos + PKT_HDR_SIZE
+                    payload_dst = first_pos + PKT_HDR_SIZE + total_size
+                    if payload_src != payload_dst:
+                        # Calculate how much data is after this packet
+                        data_after_packet = self._recv_len - (self._recv_pos + PKT_HDR_SIZE + packet_length)
+                        # Move this packet's payload
+                        self._recv_buf[payload_dst:payload_dst + packet_length] = \
+                            self._recv_buf[payload_src:payload_src + packet_length]
+                        # Move any data after this packet
+                        if data_after_packet > 0:
+                            self._recv_buf[payload_dst + packet_length:payload_dst + packet_length + data_after_packet] = \
+                                self._recv_buf[self._recv_pos + PKT_HDR_SIZE + packet_length:self._recv_len]
+                        # After compaction, adjust buffer length to account for removed header
+                        self._recv_len -= PKT_HDR_SIZE
 
-                payload_write_pos += packet_length
                 total_size += packet_length
 
                 # Check if this is the last packet
                 if packet_length < MAX_PKT_SIZE:
                     # Last packet - return accumulated payload
+                    if self.logger.isEnabledFor(logging.DEBUG):
+                        conn_id_str = f"[conn_id={self.connection_id}]" if hasattr(self, 'connection_id') and self.connection_id >= 0 else ""
+                        self.logger.debug(f"RECV sync: {conn_id_str} complete multi-packet message: {packet_count} packets, {total_size} bytes total")
+                    
                     self._recv_pos = first_pos + PKT_HDR_SIZE + total_size
                     return memoryview(self._recv_buf[first_pos + PKT_HDR_SIZE:first_pos + PKT_HDR_SIZE + total_size])
 
                 # Multi-packet: advance to next packet header
-                self._recv_pos += PKT_HDR_SIZE + packet_length
+                # After compaction, the next header is immediately after current payload
+                if packet_count > 1:
+                    # After compaction, next header is at: first_pos + PKT_HDR_SIZE + total_size
+                    self._recv_pos = first_pos + PKT_HDR_SIZE + total_size
+                else:
+                    # First packet, no compaction yet
+                    self._recv_pos += PKT_HDR_SIZE + packet_length
             else:
                 self._recv_len += self._recv_into_buffer()
 
     def reset_buffer(self):
+        self._recv_buf = self._default_recv_buf
         self._recv_pos = 0
         self._recv_len = 0
 
@@ -449,16 +488,29 @@ def execute_many(self, messages: List[ClientMessage], config: 'Configuration' =
                 BATCH_SIZE = 1000
 
                 self.reset_buffer()
-                for i in range(0, len(messages), BATCH_SIZE):
-                    batch = messages[i:i + BATCH_SIZE]
-                    
-                    # Write batch
-                    for message in batch:
+                
+                # For large payloads (>1MB), process one at a time to avoid buffer issues
+                # For small payloads, batch for performance
+                has_large_payload = any(len(msg.payload(self.context)) > 1024 * 1024 for msg in messages[:min(10, len(messages))])
+                
+                if has_large_payload:
+                    # Process one command at a time for large payloads
+                    # This prevents TCP buffer issues and command mixing with multi-MB payloads
+                    for message in messages:
                         self.write_stream.write_payload(message.payload(self.context), message.type(), True)
-                    
-                    # Read responses for this batch
-                    for message in batch:
                         results.append(self._read_result(message.is_binary(), config, buffered, prepare_stmt_packet))
+                else:
+                    # Batch processing for small payloads
+                    for i in range(0, len(messages), BATCH_SIZE):
+                        batch = messages[i:i + BATCH_SIZE]
+                        
+                        # Write batch
+                        for message in batch:
+                            self.write_stream.write_payload(message.payload(self.context), message.type(), True)
+                        
+                        # Read responses for this batch
+                        for message in batch:
+                            results.append(self._read_result(message.is_binary(), config, buffered, prepare_stmt_packet))
 
             except DatabaseError as e:
                 raise e    

mariadb/impl/debug_utils.py

@@ -25,6 +25,15 @@ def hex_dump(data: Union[bytes, bytearray], descr: str = "") -> str:
     if not data:
         return ""
 
+    MAX_DUMP_SIZE = 1024
+    original_len = len(data)
+    truncated = False
+    
+    # Truncate if data is too large
+    if len(data) > MAX_DUMP_SIZE:
+        data = data[:MAX_DUMP_SIZE]
+        truncated = True
+    
     lines = [f"{descr}"]
 
     # Header
@@ -75,4 +84,8 @@ def hex_dump(data: Union[bytes, bytearray], descr: str = "") -> str:
     # Footer
     lines.append("+------+---------------------------------------------------+------------------+")
 
+    # Add truncation notice if data was truncated
+    if truncated:
+        lines.append(f"[DATA TRUNCATED: showing {MAX_DUMP_SIZE} of {original_len} bytes]")
+    
     return "\n".join(lines)