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_castore::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::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        nar_refs: sui_cache::MemNarRefIndex,
118    }
119
120    impl CountingRemote {
121        fn with_entry(self, hash: &str, image_ref: &str) -> Self {
122            self.entries.lock().unwrap().insert(hash.to_string(), image_ref.to_string());
123            self
124        }
125    }
126
127    #[async_trait]
128    impl StorageBackend for CountingRemote {
129        async fn get_narinfo(&self, hash: &str) -> Result<Option<String>, CacheError> {
130            self.get_calls.fetch_add(1, Ordering::SeqCst);
131            Ok(self.entries.lock().unwrap().get(hash).cloned())
132        }
133        async fn put_narinfo_record(&self, hash: &str, content: &str) -> Result<(), CacheError> {
134            self.entries.lock().unwrap().insert(hash.to_string(), content.to_string());
135            Ok(())
136        }
137        async fn delete_narinfo_record(&self, hash: &str) -> Result<(), CacheError> {
138            self.entries.lock().unwrap().remove(hash);
139            Ok(())
140        }
141        async fn delete_nar_record(&self, _nar_path: &str) -> Result<(), CacheError> {
142            Ok(())
143        }
144        fn nar_ref_index(&self) -> &dyn sui_cache::NarRefIndex {
145            &self.nar_refs
146        }
147        async fn get_nar(&self, _path: &str) -> Result<Option<Vec<u8>>, CacheError> {
148            Ok(None)
149        }
150        async fn put_nar(&self, _path: &str, _data: &[u8]) -> Result<(), CacheError> {
151            Ok(())
152        }
153        /// An in-memory test double holds whole values by construction. The
154        /// declaration is required precisely so a *production* backend cannot
155        /// inherit this path by omission.
156        fn nar_residency(&self) -> sui_cache::NarResidency {
157            sui_cache::NarResidency::WholeValue
158        }
159
160        async fn list_narinfos(&self) -> Result<Vec<String>, CacheError> {
161            Ok(self.entries.lock().unwrap().keys().cloned().collect())
162        }
163    }
164
165    #[tokio::test]
166    async fn local_hit_never_calls_remote() {
167        let local = Arc::new(
168            MockLocalCacheStore::new().with_entry("h1", CachedArtifact { image_ref: "img:local".to_string() }),
169        );
170        let remote: Arc<CountingRemote> = Arc::new(CountingRemote::default());
171        let remote_dyn: Arc<dyn StorageBackend> = remote.clone();
172        let daemon = NodeCacheDaemon::new(local, remote_dyn);
173
174        let resp = daemon.handle_request(DaemonRequest::Get { content_hash: "h1".to_string() }).await;
175
176        match resp {
177            DaemonResponse::Get { artifact: Some(a) } => assert_eq!(a.image_ref, "img:local"),
178            other => panic!("expected local hit, got {other:?}"),
179        }
180        assert_eq!(remote.get_calls.load(Ordering::SeqCst), 0, "remote must not be touched on a local hit");
181    }
182
183    #[tokio::test]
184    async fn local_miss_calls_remote_exactly_once_and_persists_for_next_get() {
185        let local = Arc::new(MockLocalCacheStore::new());
186        let remote = Arc::new(CountingRemote::default().with_entry("h2", "img:remote"));
187        let remote_dyn: Arc<dyn StorageBackend> = remote.clone();
188        // Keep a handle to assert local persistence.
189        let local_check = Arc::clone(&local);
190        let daemon = NodeCacheDaemon::new(local, remote_dyn);
191
192        let resp = daemon.handle_request(DaemonRequest::Get { content_hash: "h2".to_string() }).await;
193        match resp {
194            DaemonResponse::Get { artifact: Some(a) } => assert_eq!(a.image_ref, "img:remote"),
195            other => panic!("expected remote-fetched hit, got {other:?}"),
196        }
197        assert_eq!(remote.get_calls.load(Ordering::SeqCst), 1);
198        assert!(local_check.contains("h2"), "remote hit must be persisted locally");
199
200        // A second Get for the same hash must now be a pure local hit —
201        // zero additional remote calls.
202        let resp2 = daemon.handle_request(DaemonRequest::Get { content_hash: "h2".to_string() }).await;
203        assert!(matches!(resp2, DaemonResponse::Get { artifact: Some(_) }));
204        assert_eq!(remote.get_calls.load(Ordering::SeqCst), 1, "second Get must not re-touch remote");
205    }
206
207    #[tokio::test]
208    async fn local_and_remote_miss_returns_none_without_error() {
209        let local = Arc::new(MockLocalCacheStore::new());
210        let remote: Arc<dyn StorageBackend> = Arc::new(CountingRemote::default());
211        let daemon = NodeCacheDaemon::new(local, remote);
212
213        let resp = daemon.handle_request(DaemonRequest::Get { content_hash: "nope".to_string() }).await;
214        assert!(matches!(resp, DaemonResponse::Get { artifact: None }));
215    }
216
217    #[tokio::test]
218    async fn put_persists_locally_only_never_touches_remote() {
219        let local = Arc::new(MockLocalCacheStore::new());
220        let remote = Arc::new(CountingRemote::default());
221        let remote_dyn: Arc<dyn StorageBackend> = remote.clone();
222        let local_check = Arc::clone(&local);
223        let daemon = NodeCacheDaemon::new(local, remote_dyn);
224
225        let resp = daemon
226            .handle_request(DaemonRequest::Put {
227                content_hash: "h3".to_string(),
228                artifact: CachedArtifact { image_ref: "img:new".to_string() },
229            })
230            .await;
231        assert!(matches!(resp, DaemonResponse::Put { ok: true }));
232        assert!(local_check.contains("h3"));
233        assert_eq!(remote.get_calls.load(Ordering::SeqCst), 0);
234    }
235
236    #[tokio::test]
237    async fn warm_on_already_local_hash_reports_already_local_and_skips_remote() {
238        let local = Arc::new(
239            MockLocalCacheStore::new().with_entry("h4", CachedArtifact { image_ref: "img:local".to_string() }),
240        );
241        let remote: Arc<dyn StorageBackend> = Arc::new(CountingRemote::default());
242        let daemon = NodeCacheDaemon::new(local, remote);
243
244        let resp = daemon.handle_request(DaemonRequest::Warm { content_hash: "h4".to_string() }).await;
245        assert!(matches!(resp, DaemonResponse::Warm { status: WarmStatus::AlreadyLocal }));
246    }
247
248    #[tokio::test]
249    async fn warm_on_missing_hash_schedules_background_fetch_that_persists() {
250        let local = Arc::new(MockLocalCacheStore::new());
251        let remote = Arc::new(CountingRemote::default().with_entry("h5", "img:warmed"));
252        let remote_dyn: Arc<dyn StorageBackend> = remote.clone();
253        let local_check = Arc::clone(&local);
254        let daemon = NodeCacheDaemon::new(local, remote_dyn);
255
256        let resp = daemon.handle_request(DaemonRequest::Warm { content_hash: "h5".to_string() }).await;
257        assert!(matches!(resp, DaemonResponse::Warm { status: WarmStatus::FetchScheduled }));
258
259        // Give the spawned background task a chance to run.
260        for _ in 0..50 {
261            if local_check.contains("h5") {
262                break;
263            }
264            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
265        }
266        assert!(local_check.contains("h5"), "background warm fetch must persist locally");
267        assert_eq!(remote.get_calls.load(Ordering::SeqCst), 1);
268    }
269}
sui_dockerfile_node_cache_daemon/server.rs

sui_dockerfile_node_cache_daemon/
server.rs
