common_cache/
paged.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Immutable byte pages backed by native cache loading and direct range reads.
16
17#[cfg(test)]
18mod tests;
19
20use std::collections::{BTreeMap, BTreeSet};
21use std::future::Future;
22use std::hash::Hash;
23use std::io;
24use std::num::NonZeroU64;
25use std::ops::Range;
26use std::sync::Arc;
27use std::task::Poll;
28
29use bytes::{Bytes, BytesMut};
30use common_base::range_read::RangeReader;
31use futures::channel::{mpsc, oneshot};
32use futures::{StreamExt, future};
33
34use crate::metrics::{CacheMetrics, CacheSizeMetric};
35use crate::page::PageLayout;
36use crate::{Cache, RemovalCause};
37
38/// Business-level events; callers retain their existing hit/miss metric units.
39pub enum PageEvent {
40    /// Unique logical page lookups in one request, before any missing loads execute.
41    Lookup { hits: usize, misses: usize },
42    /// One successful native page loader, excluding waiters and hits.
43    Loaded { bytes: u64 },
44    /// Backend weight, not process RSS; sampled after fetch/maintenance.
45    Residency { usage: u64, capacity: u64 },
46    /// Backend eviction, excluding unadmitted values released by their owners.
47    Evicted { cause: RemovalCause },
48}
49
50/// Caller-bound metrics without storage-engine dependencies or global registrations.
51#[derive(Clone)]
52pub struct PageCacheObserver {
53    pub bytes: CacheSizeMetric,
54    pub loads: Option<Arc<CacheMetrics>>,
55    pub events: Arc<dyn Fn(PageEvent) + Send + Sync>,
56}
57
58/// Byte accounting for one successful fetch. Hits are initial lookups, not waiters.
59#[derive(Debug, Default)]
60pub struct PageFetchStats {
61    pub cache_hit: usize,
62    pub cache_miss: usize,
63    /// Requested intersections already cached; overlapping requests count repeatedly.
64    pub cached_bytes: u64,
65    pub cached_fragments: usize,
66    /// Original request lengths, including duplicates and overlaps.
67    pub requested_bytes: u64,
68    /// Sum of unique logical page lengths, clipped at EOF, including cache hits.
69    pub aligned_bytes: u64,
70    /// Union of original request byte ranges, excluding overlaps.
71    pub unique_requested_bytes: u64,
72}
73
74/// Caches fixed-size byte pages of immutable content identified by `K`.
75/// File identity must include any version needed to distinguish different contents.
76/// Missing pages elected by this fetch are submitted through [`RangeReader::read_vec`].
77pub struct PagedCache<K> {
78    cache: Cache<(K, u64), Arc<CachedPage>>,
79    page_size: u64,
80    observer: PageCacheObserver,
81}
82
83impl<K: Hash + Eq + Clone + Send + Sync + 'static> PagedCache<K> {
84    /// Creates an independent cache budget. Zero page size is rejected by fetch.
85    pub fn new(
86        name: &str,
87        capacity: u64,
88        page_size: u64,
89        weigher: impl Fn(&(K, u64), &Bytes) -> u32 + Send + Sync + 'static,
90        observer: PageCacheObserver,
91    ) -> Self {
92        let events = observer.events.clone();
93        let mut cache = Cache::new(
94            name,
95            capacity,
96            move |key, value: &Arc<CachedPage>| weigher(key, &value.bytes),
97            move |_, _, cause| {
98                events(PageEvent::Evicted { cause });
99            },
100        )
101        .with_size_metric(observer.bytes.clone());
102        if let Some(metrics) = &observer.loads {
103            cache = cache.with_metrics(metrics.clone());
104        }
105        Self {
106            cache,
107            page_size,
108            observer,
109        }
110    }
111
112    /// Invalidates completed pages. Concurrent loads are not fenced.
113    pub fn invalidate_if(
114        &self,
115        predicate: impl Fn(&K) -> bool + Send + Sync + 'static,
116    ) -> io::Result<()> {
117        self.cache
118            .invalidate_if(move |(file, _), _| predicate(file))
119    }
120
121    /// Looks up a completed page without updating business lookup counters.
122    pub async fn contains(&self, file: &K, page_id: u64) -> bool {
123        self.cache.get(&(file.clone(), page_id)).await.is_some()
124    }
125
126    /// Recovers immutable file size from a requested page already in memory.
127    pub async fn known_file_size(&self, file: &K, ranges: &[Range<u64>]) -> Option<u64> {
128        let layout = PageLayout::new(NonZeroU64::new(self.page_size)?, u64::MAX);
129        for range in ranges {
130            for id in layout.page_ids(range.clone()).ok()? {
131                if let Some(page) = self.cache.get(&(file.clone(), id)).await {
132                    return Some(page.file_size);
133                }
134            }
135        }
136        None
137    }
138
139    /// Applies pending cache admission and eviction work.
140    pub async fn run_pending_tasks(&self) {
141        self.cache.run_pending_tasks().await;
142        self.observe_residency();
143    }
144
145    /// Fetches bytes in input order. The reader must address the same immutable
146    /// content as `file`. Each successful loader reads one complete missing page.
147    pub async fn fetch<R: RangeReader + ?Sized>(
148        &self,
149        file: K,
150        file_size: u64,
151        ranges: &[Range<u64>],
152        reader: &R,
153    ) -> Result<(Vec<Bytes>, PageFetchStats), Arc<io::Error>> {
154        if ranges.iter().any(|r| r.start > r.end || r.end > file_size) {
155            return Err(invalid_request());
156        }
157        let page_size = NonZeroU64::new(self.page_size).ok_or_else(invalid_request)?;
158        let layout = PageLayout::new(page_size, file_size);
159        let mut stats = PageFetchStats::default();
160        let mut ids = BTreeSet::new();
161        for range in ranges {
162            stats.requested_bytes = stats
163                .requested_bytes
164                .checked_add(range.end - range.start)
165                .ok_or_else(invalid_request)?;
166            ids.extend(layout.page_ids(range.clone()).map_err(Arc::new)?);
167        }
168        stats.unique_requested_bytes = unique_requested_bytes(ranges);
169        let mut pages = BTreeMap::new();
170        let mut missing = Vec::new();
171        for id in ids {
172            let range = layout.page_range(id).ok_or_else(invalid_request)?;
173            stats.aligned_bytes += range.end - range.start;
174            if let Some(page) = self.cache.get(&(file.clone(), id)).await {
175                check_file_size(&page, file_size)?;
176                pages.insert(id, page.bytes.clone());
177            } else {
178                missing.push(id);
179            }
180        }
181        for range in ranges {
182            for id in layout.page_ids(range.clone()).map_err(Arc::new)? {
183                if pages.contains_key(&id) {
184                    let slice = layout
185                        .request_slice(id, range.clone())
186                        .map_err(Arc::new)?
187                        .ok_or_else(invalid_request)?;
188                    stats.cached_bytes += slice.len() as u64;
189                    stats.cached_fragments += 1;
190                }
191            }
192        }
193        stats.cache_hit = pages.len();
194        stats.cache_miss = missing.len();
195        (self.observer.events)(PageEvent::Lookup {
196            hits: stats.cache_hit,
197            misses: stats.cache_miss,
198        });
199        if !missing.is_empty() {
200            let (requests, receiver) = mpsc::unbounded();
201            let batch_size = missing.len().max(1);
202            let mut loads = missing
203                .into_iter()
204                .map(|id| {
205                    let requests = requests.clone();
206                    let key = (file.clone(), id);
207                    Box::pin(future::maybe_done(async move {
208                        let range = layout.page_range(id).ok_or_else(invalid_request)?;
209                        let page = self
210                            .cache
211                            .try_get_with(&key, async {
212                                let (send, receive) =
213                                    oneshot::channel::<Result<Bytes, Arc<io::Error>>>();
214                                requests
215                                    .unbounded_send((range.clone(), send))
216                                    .map_err(|_| io::Error::other("Page read batch closed"))?;
217                                let bytes = receive
218                                    .await
219                                    .map_err(|_| io::Error::other("Page read batch cancelled"))?
220                                    .map_err(|error: Arc<io::Error>| {
221                                        io::Error::new(error.kind(), error)
222                                    })?;
223                                if bytes.len() as u64 != range.end - range.start {
224                                    return Err(io::Error::new(
225                                        io::ErrorKind::UnexpectedEof,
226                                        "Invalid page response length",
227                                    ));
228                                }
229                                // Detach readers' shared backing buffers from cache accounting.
230                                Ok(self
231                                    .create_cached_page(file_size, Bytes::copy_from_slice(&bytes)))
232                            })
233                            .await?;
234                        check_file_size(&page, file_size)?;
235                        Ok::<_, Arc<io::Error>>((id, page.bytes.clone()))
236                    }))
237                })
238                .collect::<Vec<_>>();
239            // Poll every page before driving I/O, so executor fairness limits do not
240            // split one fetch's immediately available misses into separate batches.
241            let loads = future::poll_fn(move |cx| {
242                let mut complete = true;
243                for load in &mut loads {
244                    complete &= load.as_mut().poll(cx).is_ready();
245                    if let Some(Err(error)) = load.as_mut().output_mut() {
246                        return Poll::Ready(Err(error.clone()));
247                    }
248                }
249                if !complete {
250                    return Poll::Pending;
251                }
252                Poll::Ready(
253                    loads
254                        .iter_mut()
255                        .map(|load| load.as_mut().take_output().ok_or_else(invalid_request)?)
256                        .collect::<Result<Vec<_>, _>>(),
257                )
258            });
259            // The channel contains only loaders elected by Moka for this fetch. Waiters
260            // enqueue nothing; a waiter promoted after cancellation can enqueue later.
261            drop(requests);
262            let read_batches = async {
263                let mut batches = receiver.ready_chunks(batch_size);
264                while let Some(batch) = batches.next().await {
265                    let (ranges, replies): (Vec<_>, Vec<_>) = batch.into_iter().unzip();
266                    let result = reader.read_vec(&ranges).await.and_then(|bytes| {
267                        if bytes.len() == ranges.len() {
268                            Ok(bytes)
269                        } else {
270                            Err(io::Error::new(
271                                io::ErrorKind::InvalidData,
272                                "Invalid page response count",
273                            ))
274                        }
275                    });
276                    match result {
277                        Ok(bytes) => {
278                            for (reply, bytes) in replies.into_iter().zip(bytes) {
279                                let _ = reply.send(Ok(bytes));
280                            }
281                        }
282                        Err(error) => {
283                            let error = Arc::new(error);
284                            for reply in replies {
285                                let _ = reply.send(Err(error.clone()));
286                            }
287                        }
288                    }
289                }
290                Ok::<(), Arc<io::Error>>(())
291            };
292            let (loaded, ()) = future::try_join(loads, read_batches).await?;
293            pages.extend(loaded);
294        }
295        let mut output = Vec::with_capacity(ranges.len());
296        for range in ranges {
297            if range.is_empty() {
298                output.push(Bytes::new());
299                continue;
300            }
301            let ids = layout.page_ids(range.clone()).map_err(Arc::new)?;
302            if ids.end - ids.start == 1 {
303                let slice = layout
304                    .request_slice(ids.start, range.clone())
305                    .map_err(Arc::new)?
306                    .ok_or_else(invalid_request)?;
307                let page = pages.get(&ids.start).ok_or_else(invalid_request)?;
308                if slice.end > page.len() {
309                    return Err(invalid_request());
310                }
311                output.push(page.slice(slice));
312                continue;
313            }
314            let len = usize::try_from(range.end - range.start).map_err(|_| invalid_request())?;
315            let mut assembled = BytesMut::with_capacity(len);
316            for id in ids {
317                let slice = layout
318                    .request_slice(id, range.clone())
319                    .map_err(Arc::new)?
320                    .ok_or_else(invalid_request)?;
321                let page = pages.get(&id).ok_or_else(invalid_request)?;
322                if slice.end > page.len() {
323                    return Err(invalid_request());
324                }
325                assembled.extend_from_slice(&page[slice]);
326            }
327            if assembled.len() != len {
328                return Err(invalid_request());
329            }
330            output.push(assembled.freeze());
331        }
332        self.observe_residency();
333        Ok((output, stats))
334    }
335
336    fn create_cached_page(&self, file_size: u64, bytes: Bytes) -> Arc<CachedPage> {
337        (self.observer.events)(PageEvent::Loaded {
338            bytes: bytes.len() as u64,
339        });
340        Arc::new(CachedPage { file_size, bytes })
341    }
342
343    fn observe_residency(&self) {
344        (self.observer.events)(PageEvent::Residency {
345            usage: self.cache.weighted_size(),
346            capacity: self.cache.capacity(),
347        });
348    }
349}
350
351struct CachedPage {
352    file_size: u64,
353    bytes: Bytes,
354}
355
356fn invalid_request() -> Arc<io::Error> {
357    Arc::new(io::Error::new(
358        io::ErrorKind::InvalidInput,
359        "Invalid cached page range",
360    ))
361}
362
363fn check_file_size(page: &CachedPage, expected: u64) -> Result<(), Arc<io::Error>> {
364    if page.file_size != expected {
365        return Err(Arc::new(io::Error::new(
366            io::ErrorKind::InvalidInput,
367            "Cached content identity used with inconsistent file sizes",
368        )));
369    }
370    Ok(())
371}
372
373/// Counts the union of request ranges already validated by fetch.
374fn unique_requested_bytes(ranges: &[Range<u64>]) -> u64 {
375    let mut ranges = ranges.to_vec();
376    ranges.sort_unstable_by_key(|range| range.start);
377    let mut end = 0;
378    let mut bytes = 0;
379    for range in ranges {
380        bytes += range.end.saturating_sub(range.start.max(end));
381        end = end.max(range.end);
382    }
383    bytes
384}
common_cache/paged.rs

common_cache/
paged.rs
