add memory limit for aggregations

introduce AggregationLimits to set memory consumption limit and bucket limits
memory limit is checked during aggregation, bucket limit is checked before returning the aggregation request.

Author: PSeitz

examples/aggregation.rs

@@ -192,7 +192,7 @@ fn main() -> tantivy::Result<()> {
     //
 
     let agg_req: Aggregations = serde_json::from_str(agg_req_str)?;
-    let collector = AggregationCollector::from_aggs(agg_req, None);
+    let collector = AggregationCollector::from_aggs(agg_req, Default::default());
 
     let agg_res: AggregationResults = searcher.search(&AllQuery, &collector).unwrap();
     let res2: Value = serde_json::to_value(agg_res)?;
@@ -239,7 +239,7 @@ fn main() -> tantivy::Result<()> {
     .into_iter()
     .collect();
 
-    let collector = AggregationCollector::from_aggs(agg_req, None);
+    let collector = AggregationCollector::from_aggs(agg_req, Default::default());
     // We use the `AllQuery` which will pass all documents to the AggregationCollector.
     let agg_res: AggregationResults = searcher.search(&AllQuery, &collector).unwrap();
 
@@ -287,7 +287,7 @@ fn main() -> tantivy::Result<()> {
 
     let agg_req: Aggregations = serde_json::from_str(agg_req_str)?;
 
-    let collector = AggregationCollector::from_aggs(agg_req, None);
+    let collector = AggregationCollector::from_aggs(agg_req, Default::default());
 
     let agg_res: AggregationResults = searcher.search(&AllQuery, &collector).unwrap();
     let res: Value = serde_json::to_value(agg_res)?;

src/aggregation/agg_limits.rs

@@ -0,0 +1,91 @@
+use std::collections::HashMap;
+use std::sync::atomic::AtomicU64;
+use std::sync::Arc;
+
+use super::collector::DEFAULT_MEMORY_LIMIT;
+use super::{AggregationError, DEFAULT_BUCKET_LIMIT};
+use crate::TantivyError;
+
+/// An estimate for memory consumption
+pub trait MemoryConsumption {
+    fn memory_consumption(&self) -> usize;
+}
+
+impl<K, V, S> MemoryConsumption for HashMap<K, V, S> {
+    fn memory_consumption(&self) -> usize {
+        let num_items = self.len();
+        (std::mem::size_of::<K>() + std::mem::size_of::<V>()) * num_items
+    }
+}
+
+/// Aggregation memory limit after which the request fails. Defaults to DEFAULT_MEMORY_LIMIT
+/// (500MB). The limit is shared by all SegmentCollectors
+pub struct AggregationLimits {
+    /// The counter which is shared between the aggregations for one request.
+    memory_consumption: Arc<AtomicU64>,
+    /// The memory_limit in bytes
+    memory_limit: u64,
+    /// The maximum number of buckets _returned_
+    /// This is not counting intermediate buckets.
+    bucket_limit: u32,
+}
+impl Clone for AggregationLimits {
+    fn clone(&self) -> Self {
+        Self {
+            memory_consumption: Arc::clone(&self.memory_consumption),
+            memory_limit: self.memory_limit,
+            bucket_limit: self.bucket_limit,
+        }
+    }
+}
+
+impl Default for AggregationLimits {
+    fn default() -> Self {
+        Self {
+            memory_consumption: Default::default(),
+            memory_limit: DEFAULT_MEMORY_LIMIT,
+            bucket_limit: DEFAULT_BUCKET_LIMIT,
+        }
+    }
+}
+
+impl AggregationLimits {
+    /// *memory_limit*
+    /// memory_limit is defined in bytes.
+    /// Aggregation fails when the estimated memory consumption of the aggregation is higher than
+    /// memory_limit.     
+    /// memory_limit will default to `DEFAULT_MEMORY_LIMIT` (500MB)
+    ///
+    /// *bucket_limit*
+    /// Limits the maximum number of buckets returned from an aggregation request.
+    /// bucket_limit will default to `DEFAULT_BUCKET_LIMIT` (65000)
+    pub fn new(memory_limit: Option<u64>, bucket_limit: Option<u32>) -> Self {
+        Self {
+            memory_consumption: Default::default(),
+            memory_limit: memory_limit.unwrap_or(DEFAULT_MEMORY_LIMIT),
+            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(),
+                },
+            ));
+        }
+        Ok(())
+    }
+    pub(crate) fn add_memory_consumed(&self, num_bytes: u64) {
+        self.memory_consumption
+            .fetch_add(num_bytes, std::sync::atomic::Ordering::Relaxed);
+    }
+    pub fn get_memory_consumed(&self) -> u64 {
+        self.memory_consumption
+            .load(std::sync::atomic::Ordering::Relaxed)
+    }
+    pub fn get_bucket_limit(&self) -> u32 {
+        self.bucket_limit
+    }
+}

src/aggregation/agg_req_with_accessor.rs

@@ -1,8 +1,5 @@
 //! This will enhance the request tree with access to the fastfield and metadata.
 
-use std::rc::Rc;
-use std::sync::atomic::AtomicU32;
-
 use columnar::{Column, ColumnType, StrColumn};
 
 use super::agg_req::{Aggregation, Aggregations, BucketAggregationType, MetricAggregation};
@@ -13,7 +10,7 @@ use super::metric::{
     AverageAggregation, CountAggregation, MaxAggregation, MinAggregation, StatsAggregation,
     SumAggregation,
 };
-use super::segment_agg_result::BucketCount;
+use super::segment_agg_result::AggregationLimits;
 use super::VecWithNames;
 use crate::{SegmentReader, TantivyError};
 
@@ -45,16 +42,15 @@ pub struct BucketAggregationWithAccessor {
     pub(crate) field_type: ColumnType,
     pub(crate) bucket_agg: BucketAggregationType,
     pub(crate) sub_aggregation: AggregationsWithAccessor,
-    pub(crate) bucket_count: BucketCount,
+    pub(crate) limits: AggregationLimits,
 }
 
 impl BucketAggregationWithAccessor {
     fn try_from_bucket(
         bucket: &BucketAggregationType,
         sub_aggregation: &Aggregations,
         reader: &SegmentReader,
-        bucket_count: Rc<AtomicU32>,
-        max_bucket_count: u32,
+        limits: AggregationLimits,
     ) -> crate::Result<BucketAggregationWithAccessor> {
         let mut str_dict_column = None;
         let (accessor, field_type) = match &bucket {
@@ -82,15 +78,11 @@ impl BucketAggregationWithAccessor {
             sub_aggregation: get_aggs_with_accessor_and_validate(
                 &sub_aggregation,
                 reader,
-                bucket_count.clone(),
-                max_bucket_count,
+                &limits.clone(),
             )?,
             bucket_agg: bucket.clone(),
             str_dict_column,
-            bucket_count: BucketCount {
-                bucket_count,
-                max_bucket_count,
-            },
+            limits,
         })
     }
 }
@@ -130,8 +122,7 @@ impl MetricAggregationWithAccessor {
 pub(crate) fn get_aggs_with_accessor_and_validate(
     aggs: &Aggregations,
     reader: &SegmentReader,
-    bucket_count: Rc<AtomicU32>,
-    max_bucket_count: u32,
+    limits: &AggregationLimits,
 ) -> crate::Result<AggregationsWithAccessor> {
     let mut metrics = vec![];
     let mut buckets = vec![];
@@ -143,8 +134,7 @@ pub(crate) fn get_aggs_with_accessor_and_validate(
                     &bucket.bucket_agg,
                     &bucket.sub_aggregation,
                     reader,
-                    Rc::clone(&bucket_count),
-                    max_bucket_count,
+                    limits.clone(),
                 )?,
             )),
             Aggregation::Metric(metric) => metrics.push((

src/aggregation/agg_result.rs

@@ -11,6 +11,7 @@ use super::agg_req::BucketAggregationInternal;
 use super::bucket::GetDocCount;
 use super::intermediate_agg_result::{IntermediateBucketResult, IntermediateMetricResult};
 use super::metric::{SingleMetricResult, Stats};
+use super::segment_agg_result::AggregationLimits;
 use super::Key;
 use crate::TantivyError;
 
@@ -19,6 +20,13 @@ use crate::TantivyError;
 pub struct AggregationResults(pub FxHashMap<String, AggregationResult>);
 
 impl AggregationResults {
+    pub(crate) fn get_bucket_count(&self) -> u64 {
+        self.0
+            .values()
+            .map(|agg| agg.get_bucket_count())
+            .sum::<u64>()
+    }
+
     pub(crate) fn get_value_from_aggregation(
         &self,
         name: &str,
@@ -47,6 +55,13 @@ pub enum AggregationResult {
 }
 
 impl AggregationResult {
+    pub(crate) fn get_bucket_count(&self) -> u64 {
+        match self {
+            AggregationResult::BucketResult(bucket) => bucket.get_bucket_count(),
+            AggregationResult::MetricResult(_) => 0,
+        }
+    }
+
     pub(crate) fn get_value_from_aggregation(
         &self,
         _name: &str,
@@ -153,9 +168,28 @@ pub enum BucketResult {
 }
 
 impl BucketResult {
-    pub(crate) fn empty_from_req(req: &BucketAggregationInternal) -> crate::Result<Self> {
+    pub(crate) fn get_bucket_count(&self) -> u64 {
+        match self {
+            BucketResult::Range { buckets } => {
+                buckets.iter().map(|bucket| bucket.get_bucket_count()).sum()
+            }
+            BucketResult::Histogram { buckets } => {
+                buckets.iter().map(|bucket| bucket.get_bucket_count()).sum()
+            }
+            BucketResult::Terms {
+                buckets,
+                sum_other_doc_count: _,
+                doc_count_error_upper_bound: _,
+            } => buckets.iter().map(|bucket| bucket.get_bucket_count()).sum(),
+        }
+    }
+
+    pub(crate) fn empty_from_req(
+        req: &BucketAggregationInternal,
+        limits: &AggregationLimits,
+    ) -> crate::Result<Self> {
         let empty_bucket = IntermediateBucketResult::empty_from_req(&req.bucket_agg);
-        empty_bucket.into_final_bucket_result(req)
+        empty_bucket.into_final_bucket_result(req, limits)
     }
 }
 
@@ -170,6 +204,15 @@ pub enum BucketEntries<T> {
     HashMap(FxHashMap<String, T>),
 }
 
+impl<T> BucketEntries<T> {
+    fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = &T> + 'a> {
+        match self {
+            BucketEntries::Vec(vec) => Box::new(vec.iter()),
+            BucketEntries::HashMap(map) => Box::new(map.values()),
+        }
+    }
+}
+
 /// This is the default entry for a bucket, which contains a key, count, and optionally
 /// sub-aggregations.
 ///
@@ -209,6 +252,11 @@ pub struct BucketEntry {
     /// Sub-aggregations in this bucket.
     pub sub_aggregation: AggregationResults,
 }
+impl BucketEntry {
+    pub(crate) fn get_bucket_count(&self) -> u64 {
+        1 + self.sub_aggregation.get_bucket_count()
+    }
+}
 impl GetDocCount for &BucketEntry {
     fn doc_count(&self) -> u64 {
         self.doc_count
@@ -272,3 +320,8 @@ pub struct RangeBucketEntry {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub to_as_string: Option<String>,
 }
+impl RangeBucketEntry {
+    pub(crate) fn get_bucket_count(&self) -> u64 {
+        1 + self.sub_aggregation.get_bucket_count()
+    }
+}