pingora_limits/
inflight.rs

1// Copyright 2024 Cloudflare, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! The inflight module defines the [Inflight] type which estimates the count of events occurring
16//! at any point in time.
17
18use crate::estimator::Estimator;
19use crate::{hash, RandomState};
20use std::hash::Hash;
21use std::sync::Arc;
22
23/// An `Inflight` type tracks the frequency of actions that are actively occurring. When the value
24/// is dropped from scope, the count will automatically decrease.
25pub struct Inflight {
26    estimator: Arc<Estimator>,
27    hasher: RandomState,
28}
29
30// fixed parameters for simplicity: hashes: h, slots: n
31// Time complexity for a lookup operation is O(h). Space complexity is O(h*n)
32// False positive ratio is 1/(n^h)
33// We choose a small h and a large n to keep lookup cheap and FP ratio low
34const HASHES: usize = 4;
35const SLOTS: usize = 8192;
36
37impl Inflight {
38    /// Create a new `Inflight`.
39    pub fn new() -> Self {
40        Inflight {
41            estimator: Arc::new(Estimator::new(HASHES, SLOTS)),
42            hasher: RandomState::new(),
43        }
44    }
45
46    /// Increment `key` by the value given. The return value is a tuple of a [Guard] and the
47    /// estimated count.
48    pub fn incr<T: Hash>(&self, key: T, value: isize) -> (Guard, isize) {
49        let guard = Guard {
50            estimator: self.estimator.clone(),
51            id: hash(key, &self.hasher),
52            value,
53        };
54        let estimation = guard.incr();
55        (guard, estimation)
56    }
57}
58
59/// A `Guard` is returned when an `Inflight` key is incremented via [Inflight::incr].
60pub struct Guard {
61    estimator: Arc<Estimator>,
62    // store the hash instead of the actual key to save space
63    id: u64,
64    value: isize,
65}
66
67impl Guard {
68    /// Increment the key's value that the `Guard` was created from.
69    pub fn incr(&self) -> isize {
70        self.estimator.incr(self.id, self.value)
71    }
72
73    /// Get the estimated count of the key that the `Guard` was created from.
74    pub fn get(&self) -> isize {
75        self.estimator.get(self.id)
76    }
77}
78
79impl Drop for Guard {
80    fn drop(&mut self) {
81        self.estimator.decr(self.id, self.value)
82    }
83}
84
85impl std::fmt::Debug for Guard {
86    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
87        f.debug_struct("Guard")
88            .field("id", &self.id)
89            .field("value", &self.value)
90            // no need to dump shared estimator
91            .finish()
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn inflight_count() {
101        let inflight = Inflight::new();
102        let (g1, v) = inflight.incr("a", 1);
103        assert_eq!(v, 1);
104        let (g2, v) = inflight.incr("a", 2);
105        assert_eq!(v, 3);
106
107        drop(g1);
108
109        assert_eq!(g2.get(), 2);
110
111        drop(g2);
112
113        let (_, v) = inflight.incr("a", 1);
114        assert_eq!(v, 1);
115    }
116}
pingora_limits/inflight.rs

pingora_limits/
inflight.rs
