sui_dockerfile_node_cache_daemon/
server.rs

1//! Request-handling core: local-first, remote-on-miss, persist-and-return.
2//!
3//! This module is transport-agnostic — it answers a
4//! [`DaemonRequest`](crate::protocol::DaemonRequest) and produces a
5//! [`DaemonResponse`](crate::protocol::DaemonResponse) given a
6//! [`LocalCacheStore`] and a remote
7//! [`StorageBackend`](sui_cache::storage::StorageBackend). The UDS
8//! listen loop (`listen`, in this crate's `lib.rs`) is a thin shell
9//! around [`NodeCacheDaemon::handle_request`].
10
11use std::sync::Arc;
12
13use sui_cache::storage::StorageBackend;
14
15use crate::protocol::{CachedArtifact, DaemonRequest, DaemonResponse, WarmStatus};
16use crate::store::LocalCacheStore;
17use crate::DaemonError;
18
19/// The daemon's request-handling core. Generic over the local store so
20/// tests can inject [`crate::store::MockLocalCacheStore`] and
21/// production wires [`crate::store::RealLocalCacheStore`]; the remote
22/// tier is always the existing `dyn StorageBackend` trait object,
23/// reused unmodified from `sui-cache`.
24pub struct NodeCacheDaemon<L: LocalCacheStore> {
25    local: Arc<L>,
26    remote: Arc<dyn StorageBackend>,
27}
28
29impl<L: LocalCacheStore + 'static> NodeCacheDaemon<L> {
30    #[must_use]
31    pub fn new(local: Arc<L>, remote: Arc<dyn StorageBackend>) -> Self {
32        Self { local, remote }
33    }
34
35    /// Answer one request. Never panics: every fallible path is
36    /// surfaced as [`DaemonResponse::Error`], not an unwind.
37    pub async fn handle_request(&self, request: DaemonRequest) -> DaemonResponse {
38        match request {
39            DaemonRequest::Get { content_hash } => self.handle_get(&content_hash).await,
40            DaemonRequest::Put { content_hash, artifact } => self.handle_put(&content_hash, &artifact).await,
41            DaemonRequest::Warm { content_hash } => self.handle_warm(content_hash).await,
42        }
43    }
44
45    async fn handle_get(&self, content_hash: &str) -> DaemonResponse {
46        match self.local.get(content_hash).await {
47            Ok(Some(artifact)) => return DaemonResponse::Get { artifact: Some(artifact) },
48            Ok(None) => {}
49            Err(e) => return DaemonResponse::Error { message: e.to_string() },
50        }
51
52        // Local miss: reach into the remote tier exactly once, persist
53        // locally on a hit, and return — the caller never needs to know
54        // which tier actually served the answer.
55        match self.fetch_from_remote_and_persist(content_hash).await {
56            Ok(artifact) => DaemonResponse::Get { artifact },
57            Err(e) => DaemonResponse::Error { message: e.to_string() },
58        }
59    }
60
61    async fn handle_put(&self, content_hash: &str, artifact: &CachedArtifact) -> DaemonResponse {
62        match self.local.put(content_hash, artifact).await {
63            Ok(()) => DaemonResponse::Put { ok: true },
64            Err(e) => DaemonResponse::Error { message: e.to_string() },
65        }
66    }
67
68    async fn handle_warm(&self, content_hash: String) -> DaemonResponse {
69        match self.local.get(&content_hash).await {
70            Ok(Some(_)) => return DaemonResponse::Warm { status: WarmStatus::AlreadyLocal },
71            Ok(None) => {}
72            Err(e) => return DaemonResponse::Error { message: e.to_string() },
73        }
74
75        // Fire-and-forget: schedule the fetch-and-persist on the
76        // runtime and respond immediately. A follow-up `Get` for the
77        // same hash observes the result once the task completes.
78        let local = Arc::clone(&self.local);
79        let remote = Arc::clone(&self.remote);
80        tokio::spawn(async move {
81            let daemon = NodeCacheDaemon { local, remote };
82            if let Err(e) = daemon.fetch_from_remote_and_persist(&content_hash).await {
83                tracing::warn!(target: "sui-dockerfile-node-cache-daemon", content_hash, error = %e, "background warm fetch failed");
84            }
85        });
86        DaemonResponse::Warm { status: WarmStatus::FetchScheduled }
87    }
88
89    /// Fetch a hash from the remote tier and, on a hit, persist it
90    /// locally before returning. Called at most once per `Get`/`Warm` —
91    /// never re-queries the remote tier for a hash already resolved
92    /// locally in this call.
93    async fn fetch_from_remote_and_persist(&self, content_hash: &str) -> Result<Option<CachedArtifact>, DaemonError> {
94        let remote_hit = self.remote.get_narinfo(content_hash).await?;
95        let Some(image_ref) = remote_hit else {
96            return Ok(None);
97        };
98        let artifact = CachedArtifact { image_ref };
99        self.local.put(content_hash, &artifact).await?;
100        Ok(Some(artifact))
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107    use crate::store::MockLocalCacheStore;
108    use async_trait::async_trait;
109    use std::sync::atomic::{AtomicUsize, Ordering};
110    use sui_cache::CacheError;
111
112    /// Counting remote mock — lets tests assert "called exactly once".
113    #[derive(Default)]
114    struct CountingRemote {
115        entries: std::sync::Mutex<std::collections::BTreeMap<String, String>>,
116        get_calls: AtomicUsize,
117    }
118
119    impl CountingRemote {
120        fn with_entry(self, hash: &str, image_ref: &str) -> Self {
121            self.entries.lock().unwrap().insert(hash.to_string(), image_ref.to_string());
122            self
123        }
124    }
125
126    #[async_trait]
127    impl StorageBackend for CountingRemote {
128        async fn get_narinfo(&self, hash: &str) -> Result<Option<String>, CacheError> {
129            self.get_calls.fetch_add(1, Ordering::SeqCst);
130            Ok(self.entries.lock().unwrap().get(hash).cloned())
131        }
132        async fn put_narinfo(&self, hash: &str, content: &str) -> Result<(), CacheError> {
133            self.entries.lock().unwrap().insert(hash.to_string(), content.to_string());
134            Ok(())
135        }
136        async fn get_nar(&self, _path: &str) -> Result<Option<Vec<u8>>, CacheError> {
137            Ok(None)
138        }
139        async fn put_nar(&self, _path: &str, _data: &[u8]) -> Result<(), CacheError> {
140            Ok(())
141        }
142        async fn delete(&self, hash: &str) -> Result<(), CacheError> {
143            self.entries.lock().unwrap().remove(hash);
144            Ok(())
145        }
146        async fn list_narinfos(&self) -> Result<Vec<String>, CacheError> {
147            Ok(self.entries.lock().unwrap().keys().cloned().collect())
148        }
149    }
150
151    #[tokio::test]
152    async fn local_hit_never_calls_remote() {
153        let local = Arc::new(
154            MockLocalCacheStore::new().with_entry("h1", CachedArtifact { image_ref: "img:local".to_string() }),
155        );
156        let remote: Arc<CountingRemote> = Arc::new(CountingRemote::default());
157        let remote_dyn: Arc<dyn StorageBackend> = remote.clone();
158        let daemon = NodeCacheDaemon::new(local, remote_dyn);
159
160        let resp = daemon.handle_request(DaemonRequest::Get { content_hash: "h1".to_string() }).await;
161
162        match resp {
163            DaemonResponse::Get { artifact: Some(a) } => assert_eq!(a.image_ref, "img:local"),
164            other => panic!("expected local hit, got {other:?}"),
165        }
166        assert_eq!(remote.get_calls.load(Ordering::SeqCst), 0, "remote must not be touched on a local hit");
167    }
168
169    #[tokio::test]
170    async fn local_miss_calls_remote_exactly_once_and_persists_for_next_get() {
171        let local = Arc::new(MockLocalCacheStore::new());
172        let remote = Arc::new(CountingRemote::default().with_entry("h2", "img:remote"));
173        let remote_dyn: Arc<dyn StorageBackend> = remote.clone();
174        // Keep a handle to assert local persistence.
175        let local_check = Arc::clone(&local);
176        let daemon = NodeCacheDaemon::new(local, remote_dyn);
177
178        let resp = daemon.handle_request(DaemonRequest::Get { content_hash: "h2".to_string() }).await;
179        match resp {
180            DaemonResponse::Get { artifact: Some(a) } => assert_eq!(a.image_ref, "img:remote"),
181            other => panic!("expected remote-fetched hit, got {other:?}"),
182        }
183        assert_eq!(remote.get_calls.load(Ordering::SeqCst), 1);
184        assert!(local_check.contains("h2"), "remote hit must be persisted locally");
185
186        // A second Get for the same hash must now be a pure local hit —
187        // zero additional remote calls.
188        let resp2 = daemon.handle_request(DaemonRequest::Get { content_hash: "h2".to_string() }).await;
189        assert!(matches!(resp2, DaemonResponse::Get { artifact: Some(_) }));
190        assert_eq!(remote.get_calls.load(Ordering::SeqCst), 1, "second Get must not re-touch remote");
191    }
192
193    #[tokio::test]
194    async fn local_and_remote_miss_returns_none_without_error() {
195        let local = Arc::new(MockLocalCacheStore::new());
196        let remote: Arc<dyn StorageBackend> = Arc::new(CountingRemote::default());
197        let daemon = NodeCacheDaemon::new(local, remote);
198
199        let resp = daemon.handle_request(DaemonRequest::Get { content_hash: "nope".to_string() }).await;
200        assert!(matches!(resp, DaemonResponse::Get { artifact: None }));
201    }
202
203    #[tokio::test]
204    async fn put_persists_locally_only_never_touches_remote() {
205        let local = Arc::new(MockLocalCacheStore::new());
206        let remote = Arc::new(CountingRemote::default());
207        let remote_dyn: Arc<dyn StorageBackend> = remote.clone();
208        let local_check = Arc::clone(&local);
209        let daemon = NodeCacheDaemon::new(local, remote_dyn);
210
211        let resp = daemon
212            .handle_request(DaemonRequest::Put {
213                content_hash: "h3".to_string(),
214                artifact: CachedArtifact { image_ref: "img:new".to_string() },
215            })
216            .await;
217        assert!(matches!(resp, DaemonResponse::Put { ok: true }));
218        assert!(local_check.contains("h3"));
219        assert_eq!(remote.get_calls.load(Ordering::SeqCst), 0);
220    }
221
222    #[tokio::test]
223    async fn warm_on_already_local_hash_reports_already_local_and_skips_remote() {
224        let local = Arc::new(
225            MockLocalCacheStore::new().with_entry("h4", CachedArtifact { image_ref: "img:local".to_string() }),
226        );
227        let remote: Arc<dyn StorageBackend> = Arc::new(CountingRemote::default());
228        let daemon = NodeCacheDaemon::new(local, remote);
229
230        let resp = daemon.handle_request(DaemonRequest::Warm { content_hash: "h4".to_string() }).await;
231        assert!(matches!(resp, DaemonResponse::Warm { status: WarmStatus::AlreadyLocal }));
232    }
233
234    #[tokio::test]
235    async fn warm_on_missing_hash_schedules_background_fetch_that_persists() {
236        let local = Arc::new(MockLocalCacheStore::new());
237        let remote = Arc::new(CountingRemote::default().with_entry("h5", "img:warmed"));
238        let remote_dyn: Arc<dyn StorageBackend> = remote.clone();
239        let local_check = Arc::clone(&local);
240        let daemon = NodeCacheDaemon::new(local, remote_dyn);
241
242        let resp = daemon.handle_request(DaemonRequest::Warm { content_hash: "h5".to_string() }).await;
243        assert!(matches!(resp, DaemonResponse::Warm { status: WarmStatus::FetchScheduled }));
244
245        // Give the spawned background task a chance to run.
246        for _ in 0..50 {
247            if local_check.contains("h5") {
248                break;
249            }
250            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
251        }
252        assert!(local_check.contains("h5"), "background warm fetch must persist locally");
253        assert_eq!(remote.get_calls.load(Ordering::SeqCst), 1);
254    }
255}
sui_dockerfile_node_cache_daemon/server.rs

sui_dockerfile_node_cache_daemon/
server.rs
