perf: optimised marker operations removing heap allocations (#107)

Author: amunra

cpp_test/test_line_sender.cpp

@@ -178,29 +178,27 @@ TEST_CASE("line_sender c api basics")
         2.7,
         48121.5,
         4.3};
-    CHECK(
-        ::line_sender_buffer_column_f64_arr_byte_strides(
-            buffer,
-            arr_name,
-            rank,
-            shape,
-            strides,
-            reinterpret_cast<uint8_t*>(arr_data.data()),
-            sizeof(arr_data),
-            &err));
+    CHECK(::line_sender_buffer_column_f64_arr_byte_strides(
+        buffer,
+        arr_name,
+        rank,
+        shape,
+        strides,
+        reinterpret_cast<uint8_t*>(arr_data.data()),
+        sizeof(arr_data),
+        &err));
 
     line_sender_column_name arr_name2 = QDB_COLUMN_NAME_LITERAL("a2");
     intptr_t elem_strides[] = {6, 2, 1};
-    CHECK(
-        ::line_sender_buffer_column_f64_arr_elem_strides(
-            buffer,
-            arr_name2,
-            rank,
-            shape,
-            elem_strides,
-            reinterpret_cast<uint8_t*>(arr_data.data()),
-            sizeof(arr_data),
-            &err));
+    CHECK(::line_sender_buffer_column_f64_arr_elem_strides(
+        buffer,
+        arr_name2,
+        rank,
+        shape,
+        elem_strides,
+        reinterpret_cast<uint8_t*>(arr_data.data()),
+        sizeof(arr_data),
+        &err));
     CHECK(::line_sender_buffer_at_nanos(buffer, 10000000, &err));
     CHECK(server.recv() == 0);
     CHECK(::line_sender_buffer_size(buffer) == 266);

include/questdb/ingress/line_sender.h

@@ -515,8 +515,10 @@ bool line_sender_buffer_column_str(
  * @param[in] rank Number of dimensions of the array.
  * @param[in] shape Array of dimension sizes (length = `rank`).
  *                  Each element must be a positive integer.
- * @param[in] strides Array strides, in the unit of bytes. Strides can be negative.
- * @param[in] data_buffer Array data, laid out according to the provided shape and strides.
+ * @param[in] strides Array strides, in the unit of bytes. Strides can be
+ * negative.
+ * @param[in] data_buffer Array data, laid out according to the provided shape
+ * and strides.
  * @param[in] data_buffer_len Length of the array data block in bytes.
  * @param[out] err_out Set to an error object on failure (if non-NULL).
  * @return true on success, false on error.
@@ -543,8 +545,10 @@ bool line_sender_buffer_column_f64_arr_byte_strides(
  * @param[in] rank Number of dimensions of the array.
  * @param[in] shape Array of dimension sizes (length = `rank`).
  *                   Each element must be a positive integer.
- * @param[in] strides Array strides, in the unit of elements. Strides can be negative.
- * @param[in] data_buffer Array data, laid out according to the provided shape and strides.
+ * @param[in] strides Array strides, in the unit of elements. Strides can be
+ * negative.
+ * @param[in] data_buffer Array data, laid out according to the provided shape
+ * and strides.
  * @param[in] data_buffer_len Length of the array data block in bytes.
  * @param[out] err_out Set to an error object on failure (if non-NULL).
  * @return true on success, false on error.

include/questdb/ingress/line_sender.hpp

@@ -652,8 +652,8 @@ class line_sender_buffer
      * @param name    Column name.
      * @param shape   Array dimensions (e.g., [2,3] for a 2x3 matrix).
      * @param strides Strides for each dimension, in the unit specified by `B`.
-     * @param data    Array elements laid out in row-major order. Their number must
-     *                match the product of dimension sizes.
+     * @param data    Array elements laid out in row-major order. Their number
+     * must match the product of dimension sizes.
      */
     template <bool B, typename T, size_t N>
     line_sender_buffer& column(

questdb-rs/src/ingress/mod.rs

@@ -39,6 +39,7 @@ use std::collections::HashMap;
 use std::convert::Infallible;
 use std::fmt::{Debug, Display, Formatter, Write};
 use std::io::{self, BufRead, BufReader, ErrorKind, Write as IoWrite};
+use std::num::NonZeroUsize;
 use std::ops::Deref;
 use std::path::PathBuf;
 use std::slice::from_raw_parts_mut;
@@ -476,11 +477,13 @@ impl OpCase {
     }
 }
 
-#[derive(Debug, Clone)]
+// IMPORTANT: This struct MUST remain `Copy` to ensure that
+// there are no heap allocations when performing marker operations.
+#[derive(Debug, Clone, Copy)]
 struct BufferState {
     op_case: OpCase,
     row_count: usize,
-    first_table: Option<String>,
+    first_table_len: Option<NonZeroUsize>,
     transactional: bool,
 }
 
@@ -489,17 +492,10 @@ impl BufferState {
         Self {
             op_case: OpCase::Init,
             row_count: 0,
-            first_table: None,
+            first_table_len: None,
             transactional: true,
         }
     }
-
-    fn clear(&mut self) {
-        self.op_case = OpCase::Init;
-        self.row_count = 0;
-        self.first_table = None;
-        self.transactional = true;
-    }
 }
 
 /// A reusable buffer to prepare a batch of ILP messages.
@@ -688,7 +684,7 @@ impl Buffer {
                 )
             ));
         }
-        self.marker = Some((self.output.len(), self.state.clone()));
+        self.marker = Some((self.output.len(), self.state));
         Ok(())
     }
 
@@ -721,7 +717,7 @@ impl Buffer {
     /// [`capacity`](Buffer::capacity).
     pub fn clear(&mut self) {
         self.output.clear();
-        self.state.clear();
+        self.state = BufferState::new();
         self.marker = None;
     }
 
@@ -798,16 +794,31 @@ impl Buffer {
         let name: TableName<'a> = name.try_into()?;
         self.validate_max_name_len(name.name)?;
         self.check_op(Op::Table)?;
+        let table_begin = self.output.len();
         write_escaped_unquoted(&mut self.output, name.name);
+        let table_end = self.output.len();
         self.state.op_case = OpCase::TableWritten;
 
         // A buffer stops being transactional if it targets multiple tables.
-        if let Some(first_table) = &self.state.first_table {
-            if first_table != name.name {
+        if let Some(first_table_len) = &self.state.first_table_len {
+            let first_table = &self.output[0..(first_table_len.get())];
+            let this_table = &self.output[table_begin..table_end];
+            if first_table != this_table {
                 self.state.transactional = false;
             }
         } else {
-            self.state.first_table = Some(name.name.to_owned());
+            debug_assert!(table_begin == 0);
+
+            // This is a bit confusing, so worth explaining:
+            // `NonZeroUsize::new(table_end)` will return `None` if `table_end` is 0,
+            // but we know that `table_end` is never 0 here, we just need an option type
+            // anyway, so we don't bother unwrapping it to then wrap it again.
+            let first_table_len = NonZeroUsize::new(table_end);
+
+            // Instead we just assert that it's `Some`.
+            debug_assert!(first_table_len.is_some());
+
+            self.state.first_table_len = first_table_len;
         }
         Ok(self)
     }

questdb-rs/src/tests/sender.rs

@@ -307,6 +307,78 @@ fn test_row_count() -> TestResult {
     Ok(())
 }
 
+#[test]
+fn test_transactional() -> TestResult {
+    let mut buffer = Buffer::new(ProtocolVersion::V2);
+
+    // transactional since there are no recorded tables yet
+    assert_eq!(buffer.row_count(), 0);
+    assert!(buffer.transactional());
+
+    buffer.set_marker()?;
+    buffer.table("table 1.test")?.symbol("a", "b")?.at_now()?;
+    assert_eq!(buffer.row_count(), 1); // tables={'table 1.test'}
+
+    // still transactional since there is only one single table.
+    assert!(buffer.transactional());
+
+    buffer.table("table 2.test")?.symbol("c", "d")?.at_now()?;
+
+    // not transactional since we have both tables "x" and "y".
+    assert_eq!(buffer.row_count(), 2); // tables={'table 1.test', 'table 2.test'}
+    assert!(!buffer.transactional());
+
+    buffer.rewind_to_marker()?;
+    // no tables, so we're transactional again
+    assert_eq!(buffer.row_count(), 0); // tables=[]
+    assert!(buffer.transactional());
+    assert!(buffer.is_empty());
+
+    // We add another new and different table, so we are still transactional.
+    buffer.table("test=table=3")?.symbol("e", "f")?.at_now()?;
+    assert_eq!(buffer.row_count(), 1); // tables={'test=table=3'}
+    assert!(buffer.transactional());
+
+    // Same table again, so we are still transactional.
+    buffer.table("test=table=3")?.symbol("g", "h")?.at_now()?;
+    assert_eq!(buffer.row_count(), 2); // tables={'test=table=3'}
+    assert!(buffer.transactional());
+
+    buffer.set_marker()?;
+    // We add a new different table: Name differs in length.
+    buffer.table("test=table=3 ")?.symbol("i", "j")?.at_now()?;
+    assert_eq!(buffer.row_count(), 3); // tables={'test=table=3', 'test=table=3 '}
+    assert!(!buffer.transactional());
+
+    buffer.rewind_to_marker()?;
+    assert_eq!(buffer.row_count(), 2); // tables={'test=table=3'}
+    assert!(buffer.transactional());
+
+    buffer.set_marker()?;
+    // We add a new different table: Name differs in content, but not in length.
+    buffer.table("test=table=4")?.symbol("k", "l")?.at_now()?;
+    assert_eq!(buffer.row_count(), 3); // tables={'test=table=3', 'test=table=4'}
+    assert!(!buffer.transactional());
+
+    buffer.rewind_to_marker()?;
+    assert_eq!(buffer.row_count(), 2); // tables={'test=table=3'}
+    assert!(buffer.transactional());
+
+    buffer.clear();
+    assert_eq!(buffer.row_count(), 0); // tables=[]
+    assert!(buffer.transactional());
+    assert!(buffer.is_empty());
+
+    // We add three rows of the same new table, so we are still transactional.
+    buffer.table("test=table=5")?.symbol("m", "n")?.at_now()?;
+    buffer.table("test=table=5")?.symbol("o", "p")?.at_now()?;
+    buffer.table("test=table=5")?.symbol("q", "r")?.at_now()?;
+    assert_eq!(buffer.row_count(), 3); // tables={'test=table=5'}
+    assert!(buffer.transactional());
+
+    Ok(())
+}
+
 #[test]
 fn test_auth_inconsistent_keys() -> TestResult {
     test_bad_key("fLKYEaoEb9lrn3nkwLDA-M_xnuFOdSt9y0Z7_vWSHLU", // d