Merge branch 'main' of https://github.com/lavrd/tantivy

Author: lavrd

ARCHITECTURE.md

@@ -254,7 +254,7 @@ The token positions of all of the terms are then stored in a separate file with
 The [TermInfo](src/postings/term_info.rs) gives an offset (expressed in position this time) in this file. As we iterate through the docset,
 we advance the position reader by the number of term frequencies of the current document.
 
-## [fieldnorms/](src/fieldnorms): Here is my doc, how many tokens in this field?
+## [fieldnorm/](src/fieldnorm): Here is my doc, how many tokens in this field?
 
 The [BM25](https://en.wikipedia.org/wiki/Okapi_BM25) formula also requires to know the number of tokens stored in a specific field for a given document. We store this information on one byte per document in the fieldnorm.
 The fieldnorm is therefore compressed. Values up to 40 are encoded unchanged.

Cargo.toml

@@ -22,7 +22,7 @@ once_cell = "1.10.0"
 regex = { version = "1.5.5", default-features = false, features = ["std", "unicode"] }
 aho-corasick = "1.0"
 tantivy-fst = "0.4.0"
-memmap2 = { version = "0.5.3", optional = true }
+memmap2 = { version = "0.6.0", optional = true }
 lz4_flex = { version = "0.10", default-features = false, features = ["checked-decode"], optional = true }
 brotli = { version = "3.3.4", optional = true }
 zstd = { version = "0.12", optional = true, default-features = false }

columnar/src/columnar/column_type.rs

@@ -34,7 +34,7 @@ impl fmt::Display for ColumnType {
             ColumnType::IpAddr => "ip",
             ColumnType::DateTime => "datetime",
         };
-        write!(f, "{}", short_str)
+        write!(f, "{short_str}")
     }
 }
 

columnar/src/dynamic_column.rs

@@ -26,14 +26,14 @@ impl fmt::Debug for DynamicColumn {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(f, "[{} {} |", self.get_cardinality(), self.column_type())?;
         match self {
-            DynamicColumn::Bool(col) => write!(f, " {:?}", col)?,
-            DynamicColumn::I64(col) => write!(f, " {:?}", col)?,
-            DynamicColumn::U64(col) => write!(f, " {:?}", col)?,
-            DynamicColumn::F64(col) => write!(f, "{:?}", col)?,
-            DynamicColumn::IpAddr(col) => write!(f, "{:?}", col)?,
-            DynamicColumn::DateTime(col) => write!(f, "{:?}", col)?,
-            DynamicColumn::Bytes(col) => write!(f, "{:?}", col)?,
-            DynamicColumn::Str(col) => write!(f, "{:?}", col)?,
+            DynamicColumn::Bool(col) => write!(f, " {col:?}")?,
+            DynamicColumn::I64(col) => write!(f, " {col:?}")?,
+            DynamicColumn::U64(col) => write!(f, " {col:?}")?,
+            DynamicColumn::F64(col) => write!(f, "{col:?}")?,
+            DynamicColumn::IpAddr(col) => write!(f, "{col:?}")?,
+            DynamicColumn::DateTime(col) => write!(f, "{col:?}")?,
+            DynamicColumn::Bytes(col) => write!(f, "{col:?}")?,
+            DynamicColumn::Str(col) => write!(f, "{col:?}")?,
         }
         write!(f, "]")
     }

common/src/byte_count.rs

@@ -37,7 +37,7 @@ impl ByteCount {
         for (suffix, threshold) in SUFFIX_AND_THRESHOLD.iter().rev() {
             if self.get_bytes() >= *threshold {
                 let unit_num = self.get_bytes() as f64 / *threshold as f64;
-                return format!("{:.2} {}", unit_num, suffix);
+                return format!("{unit_num:.2} {suffix}");
             }
         }
         format!("{:.2} B", self.get_bytes())

common/src/vint.rs

@@ -261,7 +261,7 @@ mod tests {
         let mut buffer2 = [0u8; 8];
         let len_vint = VInt(val as u64).serialize_into(&mut buffer);
         let res2 = serialize_vint_u32(val, &mut buffer2);
-        assert_eq!(&buffer[..len_vint], res2, "array wrong for {}", val);
+        assert_eq!(&buffer[..len_vint], res2, "array wrong for {val}");
     }
 
     #[test]

examples/index_from_multiple_threads.rs

@@ -96,7 +96,7 @@ fn main() -> tantivy::Result<()> {
             let mut index_writer_wlock = index_writer.write().unwrap();
             index_writer_wlock.commit()?
         };
-        println!("committed with opstamp {}", opstamp);
+        println!("committed with opstamp {opstamp}");
         thread::sleep(Duration::from_millis(500));
     }
 

examples/iterating_docs_and_positions.rs

@@ -84,7 +84,7 @@ fn main() -> tantivy::Result<()> {
                 // Doc 0: TermFreq 2: [0, 4]
                 // Doc 2: TermFreq 1: [0]
                 // ```
-                println!("Doc {}: TermFreq {}: {:?}", doc_id, term_freq, positions);
+                println!("Doc {doc_id}: TermFreq {term_freq}: {positions:?}");
                 doc_id = segment_postings.advance();
             }
         }
@@ -125,7 +125,7 @@ fn main() -> tantivy::Result<()> {
                 // Once again these docs MAY contains deleted documents as well.
                 let docs = block_segment_postings.docs();
                 // Prints `Docs [0, 2].`
-                println!("Docs {:?}", docs);
+                println!("Docs {docs:?}");
                 block_segment_postings.advance();
             }
         }

examples/snippet.rs

@@ -56,7 +56,7 @@ fn main() -> tantivy::Result<()> {
     for (score, doc_address) in top_docs {
         let doc = searcher.doc(doc_address)?;
         let snippet = snippet_generator.snippet_from_doc(&doc);
-        println!("Document score {}:", score);
+        println!("Document score {score}:");
         println!(
             "title: {}",
             doc.get_first(title).unwrap().as_text().unwrap()

examples/stop_words.rs

@@ -106,7 +106,7 @@ fn main() -> tantivy::Result<()> {
 
     for (score, doc_address) in top_docs {
         let retrieved_doc = searcher.doc(doc_address)?;
-        println!("\n==\nDocument score {}:", score);
+        println!("\n==\nDocument score {score}:");
         println!("{}", schema.to_json(&retrieved_doc));
     }
 

ownedbytes/src/lib.rs

@@ -160,7 +160,7 @@ impl fmt::Debug for OwnedBytes {
         } else {
             self.as_slice()
         };
-        write!(f, "OwnedBytes({:?}, len={})", bytes_truncated, self.len())
+        write!(f, "OwnedBytes({bytes_truncated:?}, len={})", self.len())
     }
 }
 
@@ -259,12 +259,12 @@ mod tests {
     fn test_owned_bytes_debug() {
         let short_bytes = OwnedBytes::new(b"abcd".as_ref());
         assert_eq!(
-            format!("{:?}", short_bytes),
+            format!("{short_bytes:?}"),
             "OwnedBytes([97, 98, 99, 100], len=4)"
         );
         let long_bytes = OwnedBytes::new(b"abcdefghijklmnopq".as_ref());
         assert_eq!(
-            format!("{:?}", long_bytes),
+            format!("{long_bytes:?}"),
             "OwnedBytes([97, 98, 99, 100, 101, 102, 103, 104, 105, 106], len=17)"
         );
     }

query-grammar/src/query_grammar.rs

@@ -56,7 +56,7 @@ fn word<'a>() -> impl Parser<&'a str, Output = String> {
             !c.is_whitespace() && ![':', '^', '{', '}', '"', '[', ']', '(', ')'].contains(&c)
         })),
     )
-        .map(|(s1, s2): (char, String)| format!("{}{}", s1, s2))
+        .map(|(s1, s2): (char, String)| format!("{s1}{s2}"))
         .and_then(|s: String| match s.as_str() {
             "OR" | "AND " | "NOT" => Err(StringStreamError::UnexpectedParse),
             _ => Ok(s),
@@ -74,7 +74,7 @@ fn relaxed_word<'a>() -> impl Parser<&'a str, Output = String> {
             !c.is_whitespace() && !['{', '}', '"', '[', ']', '(', ')'].contains(&c)
         })),
     )
-        .map(|(s1, s2): (char, String)| format!("{}{}", s1, s2))
+        .map(|(s1, s2): (char, String)| format!("{s1}{s2}"))
 }
 
 /// Parses a date time according to rfc3339
@@ -178,9 +178,9 @@ fn negative_number<'a>() -> impl Parser<&'a str, Output = String> {
     )
         .map(|(s1, s2, s3): (char, String, Option<(char, String)>)| {
             if let Some(('.', s3)) = s3 {
-                format!("{}{}.{}", s1, s2, s3)
+                format!("{s1}{s2}.{s3}")
             } else {
-                format!("{}{}", s1, s2)
+                format!("{s1}{s2}")
             }
         })
 }
@@ -419,9 +419,7 @@ mod test {
     fn assert_nearly_equals(expected: f64, val: f64) {
         assert!(
             nearly_equals(val, expected),
-            "Got {}, expected {}.",
-            val,
-            expected
+            "Got {val}, expected {expected}."
         );
     }
 
@@ -468,7 +466,7 @@ mod test {
 
     fn test_parse_query_to_ast_helper(query: &str, expected: &str) {
         let query = parse_to_ast().parse(query).unwrap().0;
-        let query_str = format!("{:?}", query);
+        let query_str = format!("{query:?}");
         assert_eq!(query_str, expected);
     }
 
@@ -554,7 +552,7 @@ mod test {
     fn test_occur_leaf() {
         let ((occur, ast), _) = super::occur_leaf().parse("+abc").unwrap();
         assert_eq!(occur, Some(Occur::Must));
-        assert_eq!(format!("{:?}", ast), "\"abc\"");
+        assert_eq!(format!("{ast:?}"), "\"abc\"");
     }
 
     #[test]
@@ -613,7 +611,7 @@ mod test {
         let escaped_special_chars_re = Regex::new(ESCAPED_SPECIAL_CHARS_PATTERN).unwrap();
         for special_char in SPECIAL_CHARS.iter() {
             assert_eq!(
-                escaped_special_chars_re.replace_all(&format!("\\{}", special_char), "$1"),
+                escaped_special_chars_re.replace_all(&format!("\\{special_char}"), "$1"),
                 special_char.to_string()
             );
         }

query-grammar/src/user_input_ast.rs

@@ -28,7 +28,7 @@ impl Debug for UserInputLeaf {
                 ref upper,
             } => {
                 if let Some(ref field) = field {
-                    write!(formatter, "\"{}\":", field)?;
+                    write!(formatter, "\"{field}\":")?;
                 }
                 lower.display_lower(formatter)?;
                 write!(formatter, " TO ")?;
@@ -37,14 +37,14 @@ impl Debug for UserInputLeaf {
             }
             UserInputLeaf::Set { field, elements } => {
                 if let Some(ref field) = field {
-                    write!(formatter, "\"{}\": ", field)?;
+                    write!(formatter, "\"{field}\": ")?;
                 }
                 write!(formatter, "IN [")?;
                 for (i, element) in elements.iter().enumerate() {
                     if i != 0 {
                         write!(formatter, " ")?;
                     }
-                    write!(formatter, "\"{}\"", element)?;
+                    write!(formatter, "\"{element}\"")?;
                 }
                 write!(formatter, "]")
             }
@@ -63,7 +63,7 @@ pub struct UserInputLiteral {
 impl fmt::Debug for UserInputLiteral {
     fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
         if let Some(ref field) = self.field_name {
-            write!(formatter, "\"{}\":", field)?;
+            write!(formatter, "\"{field}\":")?;
         }
         write!(formatter, "\"{}\"", self.phrase)?;
         if self.slop > 0 {
@@ -83,16 +83,16 @@ pub enum UserInputBound {
 impl UserInputBound {
     fn display_lower(&self, formatter: &mut fmt::Formatter) -> Result<(), fmt::Error> {
         match *self {
-            UserInputBound::Inclusive(ref word) => write!(formatter, "[\"{}\"", word),
-            UserInputBound::Exclusive(ref word) => write!(formatter, "{{\"{}\"", word),
+            UserInputBound::Inclusive(ref word) => write!(formatter, "[\"{word}\""),
+            UserInputBound::Exclusive(ref word) => write!(formatter, "{{\"{word}\""),
             UserInputBound::Unbounded => write!(formatter, "{{\"*\""),
         }
     }
 
     fn display_upper(&self, formatter: &mut fmt::Formatter) -> Result<(), fmt::Error> {
         match *self {
-            UserInputBound::Inclusive(ref word) => write!(formatter, "\"{}\"]", word),
-            UserInputBound::Exclusive(ref word) => write!(formatter, "\"{}\"}}", word),
+            UserInputBound::Inclusive(ref word) => write!(formatter, "\"{word}\"]"),
+            UserInputBound::Exclusive(ref word) => write!(formatter, "\"{word}\"}}"),
             UserInputBound::Unbounded => write!(formatter, "\"*\"}}"),
         }
     }
@@ -163,9 +163,9 @@ fn print_occur_ast(
     formatter: &mut fmt::Formatter,
 ) -> fmt::Result {
     if let Some(occur) = occur_opt {
-        write!(formatter, "{}{:?}", occur, ast)?;
+        write!(formatter, "{occur}{ast:?}")?;
     } else {
-        write!(formatter, "*{:?}", ast)?;
+        write!(formatter, "*{ast:?}")?;
     }
     Ok(())
 }
@@ -187,8 +187,8 @@ impl fmt::Debug for UserInputAst {
                 }
                 Ok(())
             }
-            UserInputAst::Leaf(ref subquery) => write!(formatter, "{:?}", subquery),
-            UserInputAst::Boost(ref leaf, boost) => write!(formatter, "({:?})^{}", leaf, boost),
+            UserInputAst::Leaf(ref subquery) => write!(formatter, "{subquery:?}"),
+            UserInputAst::Boost(ref leaf, boost) => write!(formatter, "({leaf:?})^{boost}"),
         }
     }
 }

src/aggregation/agg_limits.rs

@@ -1,12 +1,11 @@
 use std::collections::HashMap;
-use std::sync::atomic::AtomicU64;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 
 use common::ByteCount;
 
 use super::collector::DEFAULT_MEMORY_LIMIT;
 use super::{AggregationError, DEFAULT_BUCKET_LIMIT};
-use crate::TantivyError;
 
 /// An estimate for memory consumption. Non recursive
 pub trait MemoryConsumption {
@@ -68,28 +67,68 @@ impl AggregationLimits {
             bucket_limit: bucket_limit.unwrap_or(DEFAULT_BUCKET_LIMIT),
         }
     }
-    pub(crate) fn validate_memory_consumption(&self) -> crate::Result<()> {
-        if self.get_memory_consumed() > self.memory_limit {
-            return Err(TantivyError::AggregationError(
-                AggregationError::MemoryExceeded {
-                    limit: self.memory_limit,
-                    current: self.get_memory_consumed(),
-                },
-            ));
+
+    /// Create a new ResourceLimitGuard, that will release the memory when dropped.
+    pub fn new_guard(&self) -> ResourceLimitGuard {
+        ResourceLimitGuard {
+            /// The counter which is shared between the aggregations for one request.
+            memory_consumption: Arc::clone(&self.memory_consumption),
+            /// The memory_limit in bytes
+            memory_limit: self.memory_limit,
+            allocated_with_the_guard: 0,
         }
-        Ok(())
     }
-    pub(crate) fn add_memory_consumed(&self, num_bytes: u64) {
-        self.memory_consumption
-            .fetch_add(num_bytes, std::sync::atomic::Ordering::Relaxed);
-    }
-    /// Returns the estimated memory consumed by the aggregations
-    pub fn get_memory_consumed(&self) -> ByteCount {
+
+    pub(crate) fn add_memory_consumed(&self, num_bytes: u64) -> crate::Result<()> {
         self.memory_consumption
-            .load(std::sync::atomic::Ordering::Relaxed)
-            .into()
+            .fetch_add(num_bytes, Ordering::Relaxed);
+        validate_memory_consumption(&self.memory_consumption, self.memory_limit)?;
+        Ok(())
     }
+
     pub(crate) fn get_bucket_limit(&self) -> u32 {
         self.bucket_limit
     }
 }
+
+fn validate_memory_consumption(
+    memory_consumption: &AtomicU64,
+    memory_limit: ByteCount,
+) -> Result<(), AggregationError> {
+    // Load the estimated memory consumed by the aggregations
+    let memory_consumed: ByteCount = memory_consumption.load(Ordering::Relaxed).into();
+    if memory_consumed > memory_limit {
+        return Err(AggregationError::MemoryExceeded {
+            limit: memory_limit,
+            current: memory_consumed,
+        });
+    }
+    Ok(())
+}
+
+pub struct ResourceLimitGuard {
+    /// The counter which is shared between the aggregations for one request.
+    memory_consumption: Arc<AtomicU64>,
+    /// The memory_limit in bytes
+    memory_limit: ByteCount,
+    /// Allocated memory with this guard.
+    allocated_with_the_guard: u64,
+}
+
+impl ResourceLimitGuard {
+    pub(crate) fn add_memory_consumed(&self, num_bytes: u64) -> crate::Result<()> {
+        self.memory_consumption
+            .fetch_add(num_bytes, Ordering::Relaxed);
+        validate_memory_consumption(&self.memory_consumption, self.memory_limit)?;
+        Ok(())
+    }
+}
+
+impl Drop for ResourceLimitGuard {
+    /// Removes the memory consumed tracked by this _instance_ of AggregationLimits.
+    /// This is used to clear the segment specific memory consumption all at once.
+    fn drop(&mut self) {
+        self.memory_consumption
+            .fetch_sub(self.allocated_with_the_guard, Ordering::Relaxed);
+    }
+}