Implement quota-count

Cargo.toml

@@ -10,15 +10,15 @@ crate-type = ["cdylib", "lib"]
 
 [features]
 default = []
-experimental = []     # Experimental algorithms and APIs
-ahash = ["dep:ahash"] # Use ahash for HashMap and HashSet
+experimental = [] # Experimental algorithms and APIs
+fxhash = ["dep:rustc-hash"] # Use fxhash for HashMap and HashSet
 
 [dependencies]
 thiserror = "2.0"
 anyhow = "1.0"
 log = "0.4"
 serde = { version = "1.0", features = ["derive"] }
-ahash = { version = "0.8", features = ["serde"], optional = true }
+rustc-hash = { version = "2.1.1", optional = true }
 
 [dev-dependencies]
 log4rs = "1.4"

src/actions/hashmap_action_storage.rs

@@ -0,0 +1,104 @@
+use std::{fmt::Debug, hash::BuildHasher};
+
+use crate::{
+    actions::traits::{ActionId, ActionStorage},
+    events::traits::{EpochId, Uri},
+    util::hashmap::{HashMap, HashSet, RandomState},
+};
+
+#[derive(Debug, Clone)]
+pub struct UserActionState<E: EpochId, U: Uri, S: BuildHasher = RandomState> {
+    pub accessed_sites: HashMap<E, HashSet<U, S>, S>,
+}
+
+impl<E: EpochId, U: Uri, S: BuildHasher + Default> Default
+    for UserActionState<E, U, S>
+{
+    fn default() -> Self {
+        Self {
+            accessed_sites: HashMap::with_hasher(S::default()),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct HashMapActionStorage<
+    AID: ActionId,
+    E: EpochId,
+    U: Uri,
+    S: BuildHasher = RandomState,
+> {
+    pub actions: HashMap<AID, UserActionState<E, U, S>, S>,
+    pub quota_limit: Option<usize>,
+}
+
+impl<AID: ActionId, E: EpochId, U: Uri, S: BuildHasher + Default> Default
+    for HashMapActionStorage<AID, E, U, S>
+{
+    fn default() -> Self {
+        Self {
+            actions: HashMap::with_hasher(S::default()),
+            quota_limit: None,
+        }
+    }
+}
+
+impl<AID: ActionId, E: EpochId, U: Uri, S: BuildHasher + Default>
+    HashMapActionStorage<AID, E, U, S>
+{
+    pub fn new(quota_limit: Option<usize>) -> Self {
+        Self::with_hashmap_capacity(quota_limit, 0)
+    }
+
+    pub fn with_hashmap_capacity(
+        quota_limit: Option<usize>,
+        hashmap_capacity: usize,
+    ) -> Self {
+        Self {
+            actions: HashMap::with_capacity_and_hasher(
+                hashmap_capacity,
+                S::default(),
+            ),
+            quota_limit,
+        }
+    }
+}
+
+impl<AID, E, U, S> ActionStorage for HashMapActionStorage<AID, E, U, S>
+where
+    AID: ActionId,
+    E: EpochId,
+    U: Uri,
+    S: BuildHasher + Default,
+{
+    type ActionId = AID;
+    type EpochId = E;
+    type Uri = U;
+    type Error = anyhow::Error;
+
+    fn try_record_site(
+        &mut self,
+        action_id: Self::ActionId,
+        epoch: Self::EpochId,
+        site: &Self::Uri,
+    ) -> Result<bool, Self::Error> {
+        let state = self.actions.entry(action_id).or_default();
+        let epoch_sites = state
+            .accessed_sites
+            .entry(epoch)
+            .or_insert_with(|| HashSet::with_hasher(S::default()));
+
+        if epoch_sites.contains(site) {
+            return Ok(true);
+        }
+
+        if let Some(quota_limit) = self.quota_limit
+            && epoch_sites.len() >= quota_limit
+        {
+            return Ok(false);
+        }
+
+        epoch_sites.insert(site.clone());
+        Ok(true)
+    }
+}

src/actions/mod.rs

@@ -0,0 +1,2 @@
+pub mod hashmap_action_storage;
+pub mod traits;

src/actions/traits.rs

@@ -0,0 +1,24 @@
+use std::{fmt::Debug, hash::Hash};
+
+use crate::events::traits::{EpochId, Uri};
+
+/// Marker trait for user-action context identifiers
+pub trait ActionId: Copy + Eq + Hash + Debug {}
+impl<T: Copy + Eq + Hash + Debug> ActionId for T {}
+
+pub trait ActionStorage {
+    type ActionId: ActionId;
+    type EpochId: EpochId;
+    type Uri: Uri;
+    type Error: Debug;
+
+    /// Checks if a site can be recorded for the given action and epoch.
+    /// Returns Ok(true) if allowed (and records it), Ok(false) if
+    /// quota exceeded.
+    fn try_record_site(
+        &mut self,
+        action_id: Self::ActionId,
+        epoch: Self::EpochId,
+        site: &Self::Uri,
+    ) -> Result<bool, Self::Error>;
+}

src/budget/hashmap_filter_storage.rs

@@ -1,33 +1,54 @@
-use std::{fmt::Debug, hash::Hash};
+use std::{
+    collections::hash_map::Entry,
+    fmt::Debug,
+    hash::{BuildHasher, Hash},
+};
 
-use serde::{ser::SerializeStruct, Serialize};
+use serde::{Serialize, ser::SerializeStruct};
 
 use crate::{
     budget::traits::{Filter, FilterCapacities, FilterStorage},
-    util::hashmap::HashMap,
+    util::hashmap::{HashMap, RandomState},
 };
 
 /// Simple implementation of FilterStorage using a HashMap.
 /// Works for any Filter that implements the Filter trait.
-#[derive(Debug, Default)]
-pub struct HashMapFilterStorage<F, C>
+#[derive(Debug)]
+pub struct HashMapFilterStorage<F, C, S = RandomState>
 where
     C: FilterCapacities,
     F: Filter<C::Budget>,
+    S: BuildHasher + Default,
 {
     pub capacities: C,
-    pub filters: HashMap<C::FilterId, F>,
+    pub filters: HashMap<C::FilterId, F, S>,
+}
+
+impl<F, C, S> Default for HashMapFilterStorage<F, C, S>
+where
+    C: FilterCapacities + Default,
+    F: Filter<C::Budget>,
+    S: BuildHasher + Default,
+{
+    fn default() -> Self {
+        Self {
+            capacities: C::default(),
+            filters: HashMap::with_hasher(S::default()),
+        }
+    }
 }
 
-impl<F, C, FID> Serialize for HashMapFilterStorage<F, C>
+impl<F, C, FID, S> Serialize for HashMapFilterStorage<F, C, S>
 where
     C: FilterCapacities<FilterId = FID> + Serialize,
     F: Filter<C::Budget> + Serialize,
     FID: Serialize + Eq + Hash + Debug,
+    S: BuildHasher + Default,
+    HashMap<FID, F, S>: Serialize,
 {
-    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    fn serialize<Ser>(&self, serializer: Ser) -> Result<Ser::Ok, Ser::Error>
     where
-        S: serde::Serializer,
+        Ser: serde::Serializer,
     {
         let mut state =
             serializer.serialize_struct("HashMapFilterStorage", 2)?;
@@ -37,11 +58,32 @@ where
     }
 }
 
-impl<F, C> FilterStorage for HashMapFilterStorage<F, C>
+impl<F, C, S> HashMapFilterStorage<F, C, S>
+where
+    C: FilterCapacities,
+    F: Filter<C::Budget>,
+    S: BuildHasher + Default,
+{
+    pub fn with_hashmap_capacity(
+        capacities: C,
+        hashmap_capacity: usize,
+    ) -> Self {
+        Self {
+            capacities,
+            filters: HashMap::with_capacity_and_hasher(
+                hashmap_capacity,
+                S::default(),
+            ),
+        }
+    }
+}
+
+impl<F, C, S> FilterStorage for HashMapFilterStorage<F, C, S>
 where
     F: Filter<C::Budget, Error = anyhow::Error> + Clone,
     C: FilterCapacities<Error = anyhow::Error>,
     C::FilterId: Clone + Eq + Hash + Debug,
+    S: BuildHasher + Default,
 {
     type FilterId = C::FilterId;
     type Filter = F;
@@ -53,11 +95,7 @@ where
     where
         Self: Sized,
     {
-        let this = Self {
-            capacities,
-            filters: HashMap::new(),
-        };
-        Ok(this)
+        Ok(Self::with_hashmap_capacity(capacities, 0))
     }
 
     fn capacities(&self) -> &Self::Capacities {
@@ -80,6 +118,25 @@ where
         self.filters.insert(filter_id.clone(), filter);
         Ok(())
     }
+
+    /// Optimized version using Entry API
+    fn edit_filter_or_new<R>(
+        &mut self,
+        filter_id: &Self::FilterId,
+        f: impl FnOnce(&mut Self::Filter) -> Result<R, Self::Error>,
+    ) -> Result<R, Self::Error> {
+        let entry = self.filters.entry(filter_id.clone());
+
+        let filter = match entry {
+            Entry::Occupied(occupied) => occupied.into_mut(),
+            Entry::Vacant(vacant) => {
+                let capacity = self.capacities.capacity(filter_id)?;
+                vacant.insert(Self::Filter::new(capacity)?)
+            }
+        };
+
+        f(filter)
+    }
 }
 
 #[cfg(test)]

src/budget/pure_dp_filter.rs

@@ -67,7 +67,10 @@ impl Filter<PureDPBudget> for PureDPBudgetFilter {
         &mut self,
         budget: &PureDPBudget,
     ) -> Result<FilterStatus, Self::Error> {
-        debug!("The budget consumed in this epoch is {:?}, budget capacity for this epoch is  {:?}, and we need to consume this much budget {:?}", self.consumed, self.capacity, budget);
+        debug!(
+            "The budget consumed in this epoch is {:?}, budget capacity for this epoch is  {:?}, and we need to consume this much budget {:?}",
+            self.consumed, self.capacity, budget
+        );
 
         let status = self.can_consume(budget)?;
         if status == FilterStatus::Continue {

src/budget/traits.rs

@@ -65,10 +65,10 @@ pub trait FilterStorage {
     type Budget: Budget;
     type Filter: Filter<Self::Budget, Error = Self::Error>;
     type Capacities: FilterCapacities<
-        FilterId = Self::FilterId,
-        Budget = Self::Budget,
-        Error = Self::Error,
-    >;
+            FilterId = Self::FilterId,
+            Budget = Self::Budget,
+            Error = Self::Error,
+        >;
     type Error;
 
     /// Create a new filter storage with the given capacities for new filters.
@@ -144,10 +144,7 @@ pub trait FilterStorage {
         filter_id: &Self::FilterId,
         budget: &Self::Budget,
     ) -> Result<FilterStatus, Self::Error> {
-        let mut filter = self.get_filter_or_new(filter_id)?;
-        let status = filter.try_consume(budget)?;
-        self.set_filter(filter_id, filter)?;
-        Ok(status)
+        self.edit_filter_or_new(filter_id, |filter| filter.try_consume(budget))
     }
 
     /// Gets the remaining budget for a filter.

src/events/hashmap_event_storage.rs

@@ -1,29 +1,50 @@
+use std::hash::BuildHasher;
+
 use crate::{
     events::traits::{Event, EventStorage, RelevantEventSelector},
-    util::hashmap::HashMap,
+    util::hashmap::{HashMap, RandomState},
 };
 
 /// A simple in-memory event storage. Stores a mapping of epoch id to epoch
 /// events, where each epoch events is just a vec of events.
 /// Clones events when asked to retrieve events for an epoch.
-#[derive(Debug, Default)]
-pub struct HashMapEventStorage<E: Event> {
-    pub epochs: HashMap<E::EpochId, Vec<E>>,
+#[derive(Debug)]
+pub struct HashMapEventStorage<E: Event, S = RandomState>
+where
+    S: BuildHasher,
+{
+    pub epochs: HashMap<E::EpochId, Vec<E>, S>,
+}
+
+impl<E: Event, S: BuildHasher + Default> Default for HashMapEventStorage<E, S> {
+    fn default() -> Self {
+        Self {
+            epochs: HashMap::with_hasher(S::default()),
+        }
+    }
 }
 
 /// Simple in-memory event storage. Stores a mapping of epoch id to events
 /// in that epoch.
-impl<E: Event> HashMapEventStorage<E> {
+impl<E: Event, S: BuildHasher + Default> HashMapEventStorage<E, S> {
     pub fn new() -> Self {
+        Self::with_hashmap_capacity(0)
+    }
+
+    pub fn with_hashmap_capacity(hashmap_capacity: usize) -> Self {
         Self {
-            epochs: HashMap::new(),
+            epochs: HashMap::with_capacity_and_hasher(
+                hashmap_capacity,
+                S::default(),
+            ),
         }
     }
 }
 
-impl<E> EventStorage for HashMapEventStorage<E>
+impl<E, S> EventStorage for HashMapEventStorage<E, S>
 where
     E: Event + Clone,
+    S: BuildHasher + Default,
 {
     type Event = E;
     type Error = anyhow::Error;