perf: optimised marker operations removing heap allocations [backport] (#108)

Author: amunra

questdb-rs/src/ingress/mod.rs

@@ -34,6 +34,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::str::FromStr;
@@ -430,11 +431,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,
 }
 
@@ -443,17 +446,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.
@@ -494,11 +490,11 @@ impl BufferState {
 ///   * A row always starts with [`table`](Buffer::table).
 ///   * A row must contain at least one [`symbol`](Buffer::symbol) or
 ///     column (
-///       [`column_bool`](Buffer::column_bool),
-///       [`column_i64`](Buffer::column_i64),
-///       [`column_f64`](Buffer::column_f64),
-///       [`column_str`](Buffer::column_str),
-///       [`column_ts`](Buffer::column_ts)).
+///     [`column_bool`](Buffer::column_bool),
+///     [`column_i64`](Buffer::column_i64),
+///     [`column_f64`](Buffer::column_f64),
+///     [`column_str`](Buffer::column_str),
+///     [`column_ts`](Buffer::column_ts)).
 ///   * Symbols must appear before columns.
 ///   * A row must be terminated with either [`at`](Buffer::at) or
 ///     [`at_now`](Buffer::at_now).
@@ -630,7 +626,7 @@ impl Buffer {
                 )
             ));
         }
-        self.marker = Some((self.output.len(), self.state.clone()));
+        self.marker = Some((self.output.len(), self.state));
         Ok(())
     }
 
@@ -663,7 +659,7 @@ impl Buffer {
     /// [`capacity`](Buffer::capacity).
     pub fn clear(&mut self) {
         self.output.clear();
-        self.state.clear();
+        self.state = BufferState::new();
         self.marker = None;
     }
 
@@ -729,16 +725,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

@@ -151,6 +151,78 @@ fn test_row_count() -> TestResult {
     Ok(())
 }
 
+#[test]
+fn test_transactional() -> TestResult {
+    let mut buffer = Buffer::new();
+
+    // 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