Skip to main content

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