range_cache/
reader.rs

1//! Optional async read-through types.
2
3use std::{
4    collections::BTreeMap,
5    num::NonZeroUsize,
6    ops::Range,
7    sync::{Arc, Weak},
8};
9
10use async_trait::async_trait;
11use bytes::{Bytes, BytesMut};
12use futures_util::{StreamExt, TryStreamExt, stream};
13use parking_lot::Mutex;
14use tokio::sync::{Mutex as AsyncMutex, Semaphore};
15
16use crate::{RangeCache, RangeError, cache::ReadPlan};
17
18/// An immutable source capable of reading byte ranges.
19///
20/// Implementations must return exactly `range.len()` bytes on success. A
21/// [`CachedReader`] checks this contract before admitting a response.
22#[async_trait]
23pub trait RangeReader<K>: Send + Sync {
24    /// Source-specific error.
25    type Error: std::error::Error + Send + Sync + 'static;
26
27    /// Reads the requested byte range for `key`.
28    async fn read_range(&self, key: &K, range: Range<usize>) -> Result<Bytes, Self::Error>;
29}
30
31/// Configuration for [`CachedReader`].
32#[derive(Clone, Copy, Debug, Eq, PartialEq)]
33pub struct ReaderConfig {
34    max_fetch_concurrency: NonZeroUsize,
35}
36
37impl ReaderConfig {
38    /// Creates a configuration with an explicitly non-zero concurrency limit.
39    #[must_use]
40    pub const fn new(max_fetch_concurrency: NonZeroUsize) -> Self {
41        Self {
42            max_fetch_concurrency,
43        }
44    }
45
46    /// Returns the global maximum number of concurrent source fetches.
47    #[must_use]
48    pub const fn max_fetch_concurrency(self) -> NonZeroUsize {
49        self.max_fetch_concurrency
50    }
51}
52
53/// A read-through failure.
54#[derive(Debug, thiserror::Error)]
55pub enum ReadError<E> {
56    /// Invalid range or cache payload.
57    #[error(transparent)]
58    Range(#[from] RangeError),
59    /// Source-specific failure.
60    #[error("range source failed: {0}")]
61    Source(#[source] E),
62    /// The source returned fewer or more bytes than requested.
63    #[error("source returned {actual} bytes for {range:?}; expected {expected}")]
64    ShortRead {
65        /// Requested byte range.
66        range: Range<usize>,
67        /// Required byte length.
68        expected: usize,
69        /// Actual byte length.
70        actual: usize,
71    },
72}
73
74#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
75struct InFlightKey<K> {
76    key: K,
77    start: usize,
78    end: usize,
79}
80
81struct InFlightRegistry<K: Ord> {
82    entries: Mutex<BTreeMap<InFlightKey<K>, Weak<InFlightEntry>>>,
83}
84
85impl<K: Ord> Default for InFlightRegistry<K> {
86    fn default() -> Self {
87        Self {
88            entries: Mutex::new(BTreeMap::new()),
89        }
90    }
91}
92
93struct InFlightEntry {
94    gate: AsyncMutex<()>,
95    successful_response: Mutex<Option<Bytes>>,
96}
97
98impl<K: Ord + Clone> InFlightRegistry<K> {
99    fn register(self: &Arc<Self>, request: InFlightKey<K>) -> InFlightRegistration<K> {
100        let entry = {
101            let mut entries = self.entries.lock();
102            let current = entries.get(&request).and_then(Weak::upgrade);
103            let entry = current.unwrap_or_else(|| {
104                Arc::new(InFlightEntry {
105                    gate: AsyncMutex::new(()),
106                    successful_response: Mutex::new(None),
107                })
108            });
109            entries.insert(request.clone(), Arc::downgrade(&entry));
110            entry
111        };
112        InFlightRegistration {
113            registry: Arc::clone(self),
114            request,
115            entry,
116        }
117    }
118}
119
120struct InFlightRegistration<K: Ord> {
121    registry: Arc<InFlightRegistry<K>>,
122    request: InFlightKey<K>,
123    entry: Arc<InFlightEntry>,
124}
125
126impl<K: Ord> Drop for InFlightRegistration<K> {
127    fn drop(&mut self) {
128        let mut entries = self.registry.entries.lock();
129        let registered_is_self = entries
130            .get(&self.request)
131            .is_some_and(|registered| registered.ptr_eq(&Arc::downgrade(&self.entry)));
132        if registered_is_self && Arc::strong_count(&self.entry) == 1 {
133            entries.remove(&self.request);
134        }
135    }
136}
137
138/// An async read-through adapter over a [`RangeReader`].
139pub struct CachedReader<K: Ord, R: ?Sized> {
140    source: Arc<R>,
141    cache: RangeCache<K>,
142    config: ReaderConfig,
143    fetch_limit: Arc<Semaphore>,
144    in_flight: Arc<InFlightRegistry<K>>,
145}
146
147impl<K: Ord, R: ?Sized> Clone for CachedReader<K, R> {
148    fn clone(&self) -> Self {
149        Self {
150            source: Arc::clone(&self.source),
151            cache: self.cache.clone(),
152            config: self.config,
153            fetch_limit: Arc::clone(&self.fetch_limit),
154            in_flight: Arc::clone(&self.in_flight),
155        }
156    }
157}
158
159impl<K: Ord, R: ?Sized> CachedReader<K, R> {
160    /// Wraps `source` with the provided cache and concurrency policy.
161    #[must_use]
162    pub fn new(source: Arc<R>, cache: RangeCache<K>, config: ReaderConfig) -> Self {
163        Self {
164            source,
165            cache,
166            config,
167            fetch_limit: Arc::new(Semaphore::new(config.max_fetch_concurrency.get())),
168            in_flight: Arc::new(InFlightRegistry::default()),
169        }
170    }
171
172    /// Returns the shared cache.
173    #[must_use]
174    pub const fn cache(&self) -> &RangeCache<K> {
175        &self.cache
176    }
177
178    /// Returns the wrapped source.
179    #[must_use]
180    pub const fn source(&self) -> &Arc<R> {
181        &self.source
182    }
183
184    /// Returns the reader configuration.
185    #[must_use]
186    pub const fn config(&self) -> ReaderConfig {
187        self.config
188    }
189}
190
191impl<K, R> CachedReader<K, R>
192where
193    K: Ord + Clone,
194    R: RangeReader<K> + ?Sized,
195{
196    /// Reads a range, fetching and caching only its missing gaps.
197    ///
198    /// Identical in-flight key-and-gap requests share one source fetch. Other
199    /// overlapping requests remain independent. Responses rejected by the
200    /// cache capacity policy are still returned to the caller.
201    ///
202    /// # Errors
203    ///
204    /// Returns a validation error, source error, or [`ReadError::ShortRead`].
205    /// Source failures and invalid response lengths are never cached.
206    ///
207    /// # Panics
208    ///
209    /// Panics only if the privately owned fetch semaphore is unexpectedly
210    /// closed or an internal read-plan coverage invariant is violated.
211    pub async fn read(&self, key: &K, range: Range<usize>) -> Result<Bytes, ReadError<R::Error>> {
212        let (mut cached, missing) = match self.cache.read_plan(key, range.clone())? {
213            ReadPlan::Complete(bytes) => return Ok(bytes),
214            ReadPlan::Fetch { cached, missing } => (cached, missing),
215        };
216
217        let fetched = stream::iter(missing)
218            .map(|gap| self.fetch_gap(key, gap))
219            .buffer_unordered(self.config.max_fetch_concurrency.get())
220            .try_collect::<Vec<_>>()
221            .await?;
222        cached.extend(fetched);
223        cached.sort_unstable_by_key(|(chunk_range, _)| chunk_range.start);
224
225        if cached.len() == 1 && cached[0].0 == range {
226            return Ok(cached.pop().expect("single chunk exists").1);
227        }
228
229        let mut reconstructed = BytesMut::with_capacity(range.len());
230        let mut cursor = range.start;
231        for (chunk_range, bytes) in cached {
232            assert_eq!(chunk_range.start, cursor, "read plan has no gaps");
233            assert_eq!(chunk_range.len(), bytes.len(), "chunk length is exact");
234            reconstructed.extend_from_slice(&bytes);
235            cursor = chunk_range.end;
236        }
237        assert_eq!(cursor, range.end, "read plan covers the request");
238        Ok(reconstructed.freeze())
239    }
240
241    async fn fetch_gap(
242        &self,
243        key: &K,
244        range: Range<usize>,
245    ) -> Result<(Range<usize>, Bytes), ReadError<R::Error>> {
246        let registration = self.in_flight.register(InFlightKey {
247            key: key.clone(),
248            start: range.start,
249            end: range.end,
250        });
251        let _request_guard = registration.entry.gate.lock().await;
252
253        if let Some(bytes) = registration.entry.successful_response.lock().clone() {
254            return Ok((range, bytes));
255        }
256
257        if let Some(bytes) = self.cache.get(key, range.clone())? {
258            return Ok((range, bytes));
259        }
260
261        let _fetch_permit = self
262            .fetch_limit
263            .acquire()
264            .await
265            .expect("private fetch semaphore remains open");
266        let bytes = self
267            .source
268            .read_range(key, range.clone())
269            .await
270            .map_err(ReadError::Source)?;
271        let expected = range.len();
272        if bytes.len() != expected {
273            return Err(ReadError::ShortRead {
274                range,
275                expected,
276                actual: bytes.len(),
277            });
278        }
279
280        let _ = self
281            .cache
282            .insert(key.clone(), range.clone(), bytes.clone())?;
283        *registration.entry.successful_response.lock() = Some(bytes.clone());
284        Ok((range, bytes))
285    }
286}
287
288#[async_trait]
289impl<K, R> RangeReader<K> for CachedReader<K, R>
290where
291    K: Ord + Clone + Send + Sync,
292    R: RangeReader<K> + ?Sized,
293{
294    type Error = ReadError<R::Error>;
295
296    async fn read_range(&self, key: &K, range: Range<usize>) -> Result<Bytes, Self::Error> {
297        self.read(key, range).await
298    }
299}
300
301#[cfg(test)]
302mod tests {
303    use std::{convert::Infallible, num::NonZeroUsize, sync::Arc, time::Duration};
304
305    use async_trait::async_trait;
306    use bytes::Bytes;
307    use tokio::sync::Semaphore;
308
309    use super::{CachedReader, RangeReader, ReaderConfig};
310    use crate::{CacheCapacity, RangeCache};
311
312    struct BlockingSource {
313        started: Semaphore,
314    }
315
316    #[async_trait]
317    impl RangeReader<String> for BlockingSource {
318        type Error = Infallible;
319
320        async fn read_range(
321            &self,
322            _key: &String,
323            range: std::ops::Range<usize>,
324        ) -> Result<Bytes, Self::Error> {
325            self.started.add_permits(1);
326            tokio::time::sleep(Duration::from_secs(60)).await;
327            Ok(Bytes::from(vec![0; range.len()]))
328        }
329    }
330
331    #[tokio::test]
332    async fn cancelled_leader_removes_its_in_flight_registration() {
333        let source = Arc::new(BlockingSource {
334            started: Semaphore::new(0),
335        });
336        let reader = CachedReader::new(
337            Arc::clone(&source),
338            RangeCache::new(CacheCapacity::Unbounded),
339            ReaderConfig::new(NonZeroUsize::new(1).expect("non-zero")),
340        );
341        let task_reader = reader.clone();
342        let task = tokio::spawn(async move {
343            let key = String::from("key");
344            task_reader.read(&key, 0..4).await
345        });
346        source
347            .started
348            .acquire()
349            .await
350            .expect("source semaphore remains open")
351            .forget();
352        task.abort();
353        assert!(task.await.expect_err("task was cancelled").is_cancelled());
354
355        tokio::task::yield_now().await;
356        assert!(reader.in_flight.entries.lock().is_empty());
357    }
358}
range_cache/reader.rs

range_cache/
reader.rs
