pingora_cache/
variance.rs

1use std::{borrow::Cow, collections::BTreeMap};
2
3use blake2::Digest;
4
5use crate::key::{Blake2b128, HashBinary};
6
7/// A builder for variance keys, used for distinguishing multiple cached assets
8/// at the same URL. This is intended to be easily passed to helper functions,
9/// which can each populate a portion of the variance.
10pub struct VarianceBuilder<'a> {
11    values: BTreeMap<Cow<'a, str>, Cow<'a, [u8]>>,
12}
13
14impl<'a> VarianceBuilder<'a> {
15    /// Create an empty variance key. Has no variance by default - add some variance using
16    /// [`Self::add_value`].
17    pub fn new() -> Self {
18        VarianceBuilder {
19            values: BTreeMap::new(),
20        }
21    }
22
23    /// Add a byte string to the variance key. Not sensitive to insertion order.
24    /// `value` is intended to take either `&str` or `&[u8]`.
25    pub fn add_value(&mut self, name: &'a str, value: &'a (impl AsRef<[u8]> + ?Sized)) {
26        self.values
27            .insert(name.into(), Cow::Borrowed(value.as_ref()));
28    }
29
30    /// Move a byte string to the variance key. Not sensitive to insertion order. Useful when
31    /// writing helper functions which generate a value then add said value to the VarianceBuilder.
32    /// Without this, the helper function would have to move the value to the calling function
33    /// to extend its lifetime to at least match the VarianceBuilder.
34    pub fn add_owned_value(&mut self, name: &'a str, value: Vec<u8>) {
35        self.values.insert(name.into(), Cow::Owned(value));
36    }
37
38    /// Check whether this variance key actually has variance, or just refers to the root asset
39    pub fn has_variance(&self) -> bool {
40        !self.values.is_empty()
41    }
42
43    /// Hash this variance key. Returns [`None`] if [`Self::has_variance`] is false.
44    pub fn finalize(self) -> Option<HashBinary> {
45        const SALT: &[u8; 1] = &[0u8; 1];
46        if self.has_variance() {
47            let mut hash = Blake2b128::new();
48            for (name, value) in self.values.iter() {
49                hash.update(name.as_bytes());
50                hash.update(SALT);
51                hash.update(value);
52                hash.update(SALT);
53            }
54            Some(hash.finalize().into())
55        } else {
56            None
57        }
58    }
59}
60
61#[cfg(test)]
62mod test {
63    use super::*;
64
65    #[test]
66    fn test_basic() {
67        let key_empty = VarianceBuilder::new().finalize();
68        assert_eq!(None, key_empty);
69
70        let mut key_value = VarianceBuilder::new();
71        key_value.add_value("a", "a");
72        let key_value = key_value.finalize();
73
74        let mut key_owned_value = VarianceBuilder::new();
75        key_owned_value.add_owned_value("a", "a".as_bytes().to_vec());
76        let key_owned_value = key_owned_value.finalize();
77
78        assert_ne!(key_empty, key_value);
79        assert_ne!(key_empty, key_owned_value);
80        assert_eq!(key_value, key_owned_value);
81    }
82
83    #[test]
84    fn test_value_ordering() {
85        let mut key_abc = VarianceBuilder::new();
86        key_abc.add_value("a", "a");
87        key_abc.add_value("b", "b");
88        key_abc.add_value("c", "c");
89        let key_abc = key_abc.finalize().unwrap();
90
91        let mut key_bac = VarianceBuilder::new();
92        key_bac.add_value("b", "b");
93        key_bac.add_value("a", "a");
94        key_bac.add_value("c", "c");
95        let key_bac = key_bac.finalize().unwrap();
96
97        let mut key_cba = VarianceBuilder::new();
98        key_cba.add_value("c", "c");
99        key_cba.add_value("b", "b");
100        key_cba.add_value("a", "a");
101        let key_cba = key_cba.finalize().unwrap();
102
103        assert_eq!(key_abc, key_bac);
104        assert_eq!(key_abc, key_cba);
105    }
106
107    #[test]
108    fn test_value_overriding() {
109        let mut key_a = VarianceBuilder::new();
110        key_a.add_value("a", "a");
111        let key_a = key_a.finalize().unwrap();
112
113        let mut key_b = VarianceBuilder::new();
114        key_b.add_value("a", "b");
115        key_b.add_value("a", "a");
116        let key_b = key_b.finalize().unwrap();
117
118        assert_eq!(key_a, key_b);
119    }
120}