snix_store/pathinfoservice/
lru.rs

1use async_stream::try_stream;
2use futures::stream::BoxStream;
3use lru::LruCache;
4use nix_compat::nixbase32;
5use std::num::NonZeroUsize;
6use std::sync::Arc;
7use tokio::sync::RwLock;
8use tonic::async_trait;
9use tracing::instrument;
10
11use snix_castore::Error;
12use snix_castore::composition::{CompositionContext, ServiceBuilder};
13
14use super::{PathInfo, PathInfoService};
15
16pub struct LruPathInfoService {
17    instance_name: String,
18    lru: Arc<RwLock<LruCache<[u8; 20], PathInfo>>>,
19}
20
21impl LruPathInfoService {
22    pub fn with_capacity(instance_name: String, capacity: NonZeroUsize) -> Self {
23        Self {
24            instance_name,
25            lru: Arc::new(RwLock::new(LruCache::new(capacity))),
26        }
27    }
28}
29
30#[async_trait]
31impl PathInfoService for LruPathInfoService {
32    #[instrument(level = "trace", skip_all, fields(path_info.digest = nixbase32::encode(&digest), instance_name = %self.instance_name))]
33    async fn get(&self, digest: [u8; 20]) -> Result<Option<PathInfo>, Error> {
34        Ok(self.lru.write().await.get(&digest).cloned())
35    }
36
37    #[instrument(level = "trace", skip_all, fields(path_info.root_node = ?path_info.node, instance_name = %self.instance_name))]
38    async fn put(&self, path_info: PathInfo) -> Result<PathInfo, Error> {
39        self.lru
40            .write()
41            .await
42            .put(*path_info.store_path.digest(), path_info.clone());
43
44        Ok(path_info)
45    }
46
47    fn list(&self) -> BoxStream<'static, Result<PathInfo, Error>> {
48        let lru = self.lru.clone();
49        Box::pin(try_stream! {
50            let lru = lru.read().await;
51            let it = lru.iter();
52
53            for (_k,v) in it {
54                yield v.clone()
55            }
56        })
57    }
58}
59
60#[derive(serde::Deserialize, Debug)]
61#[serde(deny_unknown_fields)]
62pub struct LruPathInfoServiceConfig {
63    capacity: NonZeroUsize,
64}
65
66impl TryFrom<url::Url> for LruPathInfoServiceConfig {
67    type Error = Box<dyn std::error::Error + Send + Sync>;
68    fn try_from(_url: url::Url) -> Result<Self, Self::Error> {
69        Err(Error::StorageError(
70            "Instantiating a LruPathInfoService from a url is not supported".into(),
71        )
72        .into())
73    }
74}
75
76#[async_trait]
77impl ServiceBuilder for LruPathInfoServiceConfig {
78    type Output = dyn PathInfoService;
79    async fn build<'a>(
80        &'a self,
81        instance_name: &str,
82        _context: &CompositionContext,
83    ) -> Result<Arc<dyn PathInfoService>, Box<dyn std::error::Error + Send + Sync + 'static>> {
84        Ok(Arc::new(LruPathInfoService::with_capacity(
85            instance_name.to_string(),
86            self.capacity,
87        )))
88    }
89}
90
91#[cfg(test)]
92mod test {
93    use nix_compat::store_path::StorePath;
94    use std::{num::NonZeroUsize, sync::LazyLock};
95
96    use crate::{
97        fixtures::PATH_INFO,
98        pathinfoservice::{LruPathInfoService, PathInfo, PathInfoService},
99    };
100    static PATHINFO_2: LazyLock<PathInfo> = LazyLock::new(|| {
101        let mut p = PATH_INFO.clone();
102        p.store_path = StorePath::from_name_and_digest_fixed("dummy", [1; 20]).unwrap();
103        p
104    });
105
106    static PATHINFO_2_DIGEST: LazyLock<[u8; 20]> =
107        LazyLock::new(|| *PATHINFO_2.store_path.digest());
108
109    #[tokio::test]
110    async fn evict() {
111        let svc = LruPathInfoService::with_capacity("test".into(), NonZeroUsize::new(1).unwrap());
112
113        // pathinfo_1 should not be there
114        assert!(
115            svc.get(*PATH_INFO.store_path.digest())
116                .await
117                .expect("no error")
118                .is_none()
119        );
120
121        // insert it
122        svc.put(PATH_INFO.clone()).await.expect("no error");
123
124        // now it should be there.
125        assert_eq!(
126            Some(PATH_INFO.clone()),
127            svc.get(*PATH_INFO.store_path.digest())
128                .await
129                .expect("no error")
130        );
131
132        // insert pathinfo_2. This will evict pathinfo 1
133        svc.put(PATHINFO_2.clone()).await.expect("no error");
134
135        // now pathinfo 2 should be there.
136        assert_eq!(
137            Some(PATHINFO_2.clone()),
138            svc.get(*PATHINFO_2_DIGEST).await.expect("no error")
139        );
140
141        // … but pathinfo 1 not anymore.
142        assert!(
143            svc.get(*PATH_INFO.store_path.digest())
144                .await
145                .expect("no error")
146                .is_none()
147        );
148    }
149}