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}