zero_mysql/
buffer.rs

1/// A set of reusable buffers for MySQL protocol communication
2///
3/// `Conn` uses a single `BufferSet` for all its operations.
4#[derive(Debug)]
5pub struct BufferSet {
6    /// Bytes are valid during Conn.
7    pub initial_handshake: Vec<u8>,
8
9    /// General-purpose read buffer
10    /// Bytes are valid during an operation.
11    pub read_buffer: Vec<u8>,
12
13    /// General-purpose write buffer
14    /// It always has at least 4 bytes which is reserved for the first packet header.
15    /// It is followed by payload bytes without considering 16MB split.
16    /// Layout: [4-byte header space][payload that is possibly larger than 16MB]
17    /// Bytes are valid during an operation.
18    pub write_buffer: Vec<u8>,
19
20    /// ColumnDefinition packets in one buffer
21    /// Bytes are valid during an operation.
22    pub column_definition_buffer: Vec<u8>,
23}
24
25impl BufferSet {
26    /// Create a new empty buffer set
27    pub fn new() -> Self {
28        Self {
29            initial_handshake: Vec::new(),
30            read_buffer: Vec::new(),
31            write_buffer: vec![0; 4],
32            column_definition_buffer: Vec::new(),
33        }
34    }
35
36    /// Create a new buffer set with the initial handshake packet
37    pub fn with_initial_handshake(initial_handshake: Vec<u8>) -> Self {
38        Self {
39            initial_handshake,
40            read_buffer: Vec::new(),
41            write_buffer: vec![0; 4],
42            column_definition_buffer: Vec::new(),
43        }
44    }
45
46    /// Clear the write buffer, reserve 4 bytes for the header, and return mutable access.
47    #[inline]
48    pub fn new_write_buffer(&mut self) -> &mut Vec<u8> {
49        self.write_buffer.clear();
50        self.write_buffer.extend_from_slice(&[0_u8; 4]);
51        &mut self.write_buffer
52    }
53
54    /// Get mutable access to the write buffer.
55    #[inline]
56    pub fn write_buffer_mut(&mut self) -> &mut Vec<u8> {
57        &mut self.write_buffer
58    }
59
60    /// Get the write buffer for reading.
61    #[inline]
62    pub fn write_buffer(&self) -> &[u8] {
63        &self.write_buffer
64    }
65
66    /// Get the payload length (total buffer length minus 4-byte header).
67    #[inline]
68    pub fn payload_len(&self) -> usize {
69        self.write_buffer.len().saturating_sub(4)
70    }
71}
72
73impl Default for BufferSet {
74    fn default() -> Self {
75        Self::new()
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn test_buffer_set_new() {
85        let buffers = BufferSet::new();
86        assert!(buffers.initial_handshake.is_empty());
87        assert!(buffers.read_buffer.is_empty());
88        assert_eq!(buffers.write_buffer().len(), 4); // pre-allocated header space
89    }
90
91    #[test]
92    fn test_buffer_set_default() {
93        let buffers = BufferSet::default();
94        assert!(buffers.initial_handshake.is_empty());
95    }
96
97    #[test]
98    fn test_buffer_set_with_initial_handshake() {
99        let handshake = vec![0x0a, 0x35, 0x2e, 0x37];
100        let buffers = BufferSet::with_initial_handshake(handshake.clone());
101        assert_eq!(buffers.initial_handshake, handshake);
102        assert!(buffers.read_buffer.is_empty());
103    }
104
105    #[test]
106    fn test_new_write_buffer() {
107        let mut buffers = BufferSet::new();
108        let buf = buffers.new_write_buffer();
109        // Should have 4 bytes reserved for header
110        assert_eq!(buf.len(), 4);
111        assert_eq!(buffers.payload_len(), 0);
112    }
113
114    #[test]
115    fn test_write_buffer_mut() {
116        let mut buffers = BufferSet::new();
117        buffers.new_write_buffer().extend_from_slice(b"SELECT 1");
118        // 4 header bytes + 8 payload bytes
119        assert_eq!(buffers.write_buffer().len(), 12);
120        assert_eq!(buffers.payload_len(), 8);
121    }
122
123    #[test]
124    fn test_write_buffer() {
125        let mut buffers = BufferSet::new();
126        buffers.new_write_buffer().extend_from_slice(b"test");
127        let packet = buffers.write_buffer();
128        assert_eq!(packet.len(), 8); // 4 header + 4 payload
129        assert_eq!(&packet[4..], b"test");
130    }
131
132    #[test]
133    fn test_buffer_reuse() {
134        let mut buffers = BufferSet::new();
135
136        // Write to buffers
137        buffers.read_buffer.extend_from_slice(b"test data");
138        buffers.new_write_buffer().extend_from_slice(b"query");
139
140        assert_eq!(buffers.read_buffer.len(), 9);
141        assert_eq!(buffers.payload_len(), 5);
142
143        // Clear and reuse
144        buffers.read_buffer.clear();
145        buffers.new_write_buffer();
146
147        assert_eq!(buffers.read_buffer.len(), 0);
148        assert_eq!(buffers.payload_len(), 0);
149
150        // Capacity should be preserved
151        assert!(buffers.read_buffer.capacity() >= 9);
152    }
153}