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;
28use std::time::Duration;
29
30use bytes::{Bytes, BytesMut};
31use common_base::range_read::RangeReader;
32use futures::channel::{mpsc, oneshot};
33use futures::{StreamExt, future};
34
35use crate::metrics::{CacheMetrics, CacheSizeMetric, CacheSizeRegistration};
36use crate::page::PageLayout;
37use crate::{Cache, RemovalCause};
38
39/// Business-level events; callers retain their existing hit/miss metric units.
40pub enum PageEvent {
41    /// Unique logical page lookups in one request, before any missing loads execute.
42    Lookup { hits: usize, misses: usize },
43    /// One successful native page loader, excluding waiters and hits.
44    Loaded { bytes: u64 },
45    /// Backend eviction, excluding unadmitted values released by their owners.
46    Evicted { cause: RemovalCause },
47}
48
49/// Caller-bound metrics without storage-engine dependencies or global registrations.
50#[derive(Clone)]
51pub struct PageCacheObserver {
52    pub bytes: CacheSizeMetric,
53    pub capacity: 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    _capacity_registration: CacheSizeRegistration,
82}
83
84impl<K: Hash + Eq + Clone + Send + Sync + 'static> PagedCache<K> {
85    /// Creates an independent cache budget. Zero page size is rejected by fetch.
86    pub fn new(
87        name: &str,
88        capacity: u64,
89        page_size: u64,
90        weigher: impl Fn(&(K, u64), &Bytes) -> u32 + Send + Sync + 'static,
91        observer: PageCacheObserver,
92    ) -> Self {
93        Self::new_with_ttl(name, capacity, page_size, weigher, observer, None)
94    }
95
96    /// Creates a page cache with optional expiration measured from insertion.
97    pub fn new_with_ttl(
98        name: &str,
99        capacity: u64,
100        page_size: u64,
101        weigher: impl Fn(&(K, u64), &Bytes) -> u32 + Send + Sync + 'static,
102        observer: PageCacheObserver,
103        ttl: Option<Duration>,
104    ) -> Self {
105        let events = observer.events.clone();
106        let mut cache = Cache::new_with_ttl(
107            name,
108            capacity,
109            move |key, value: &Arc<CachedPage>| weigher(key, &value.bytes),
110            move |_, _, cause| {
111                events(PageEvent::Evicted { cause });
112            },
113            ttl,
114        )
115        .with_size_metric(observer.bytes.clone());
116        if let Some(metrics) = &observer.loads {
117            cache = cache.with_metrics(metrics.clone());
118        }
119        let capacity_registration = observer.capacity.register(move || capacity);
120        Self {
121            _capacity_registration: capacity_registration,
122            cache,
123            page_size,
124            observer,
125        }
126    }
127
128    /// Invalidates completed pages. Concurrent loads are not fenced.
129    pub fn invalidate_if(
130        &self,
131        predicate: impl Fn(&K) -> bool + Send + Sync + 'static,
132    ) -> io::Result<()> {
133        self.cache
134            .invalidate_if(move |(file, _), _| predicate(file))
135    }
136
137    /// Looks up a completed page without updating business lookup counters.
138    pub async fn contains(&self, file: &K, page_id: u64) -> bool {
139        self.cache.get(&(file.clone(), page_id)).await.is_some()
140    }
141
142    /// Recovers immutable file size from a requested page already in memory.
143    pub async fn known_file_size(&self, file: &K, ranges: &[Range<u64>]) -> Option<u64> {
144        let layout = PageLayout::new(NonZeroU64::new(self.page_size)?, u64::MAX);
145        for range in ranges {
146            for id in layout.page_ids(range.clone()).ok()? {
147                if let Some(page) = self.cache.get(&(file.clone(), id)).await {
148                    return Some(page.file_size);
149                }
150            }
151        }
152        None
153    }
154
155    /// Applies pending cache admission and eviction work.
156    pub async fn run_pending_tasks(&self) {
157        self.cache.run_pending_tasks().await;
158    }
159
160    /// Fetches bytes in input order. The reader must address the same immutable
161    /// content as `file`. Each successful loader reads one complete missing page.
162    pub async fn fetch<R: RangeReader + ?Sized>(
163        &self,
164        file: K,
165        file_size: u64,
166        ranges: &[Range<u64>],
167        reader: &R,
168    ) -> Result<(Vec<Bytes>, PageFetchStats), Arc<io::Error>> {
169        if ranges.iter().any(|r| r.start > r.end || r.end > file_size) {
170            return Err(invalid_request());
171        }
172        let page_size = NonZeroU64::new(self.page_size).ok_or_else(invalid_request)?;
173        let layout = PageLayout::new(page_size, file_size);
174        let mut stats = PageFetchStats::default();
175        let mut ids = BTreeSet::new();
176        for range in ranges {
177            stats.requested_bytes = stats
178                .requested_bytes
179                .checked_add(range.end - range.start)
180                .ok_or_else(invalid_request)?;
181            ids.extend(layout.page_ids(range.clone()).map_err(Arc::new)?);
182        }
183        stats.unique_requested_bytes = unique_requested_bytes(ranges);
184        let mut pages = BTreeMap::new();
185        let mut missing = Vec::new();
186        for id in ids {
187            let range = layout.page_range(id).ok_or_else(invalid_request)?;
188            stats.aligned_bytes += range.end - range.start;
189            if let Some(page) = self.cache.get(&(file.clone(), id)).await {
190                check_file_size(&page, file_size)?;
191                pages.insert(id, page.bytes.clone());
192            } else {
193                missing.push(id);
194            }
195        }
196        for range in ranges {
197            for id in layout.page_ids(range.clone()).map_err(Arc::new)? {
198                if pages.contains_key(&id) {
199                    let slice = layout
200                        .request_slice(id, range.clone())
201                        .map_err(Arc::new)?
202                        .ok_or_else(invalid_request)?;
203                    stats.cached_bytes += slice.len() as u64;
204                    stats.cached_fragments += 1;
205                }
206            }
207        }
208        stats.cache_hit = pages.len();
209        stats.cache_miss = missing.len();
210        (self.observer.events)(PageEvent::Lookup {
211            hits: stats.cache_hit,
212            misses: stats.cache_miss,
213        });
214        if !missing.is_empty() {
215            let (requests, receiver) = mpsc::unbounded();
216            let batch_size = missing.len().max(1);
217            let mut loads = missing
218                .into_iter()
219                .map(|id| {
220                    let requests = requests.clone();
221                    let key = (file.clone(), id);
222                    Box::pin(future::maybe_done(async move {
223                        let range = layout.page_range(id).ok_or_else(invalid_request)?;
224                        let page = self
225                            .cache
226                            .try_get_with(&key, async {
227                                let (send, receive) =
228                                    oneshot::channel::<Result<Bytes, Arc<io::Error>>>();
229                                requests
230                                    .unbounded_send((range.clone(), send))
231                                    .map_err(|_| io::Error::other("Page read batch closed"))?;
232                                let bytes = receive
233                                    .await
234                                    .map_err(|_| io::Error::other("Page read batch cancelled"))?
235                                    .map_err(|error: Arc<io::Error>| {
236                                        io::Error::new(error.kind(), error)
237                                    })?;
238                                if bytes.len() as u64 != range.end - range.start {
239                                    return Err(io::Error::new(
240                                        io::ErrorKind::UnexpectedEof,
241                                        "Invalid page response length",
242                                    ));
243                                }
244                                // Detach readers' shared backing buffers from cache accounting.
245                                Ok(self
246                                    .create_cached_page(file_size, Bytes::copy_from_slice(&bytes)))
247                            })
248                            .await?;
249                        check_file_size(&page, file_size)?;
250                        Ok::<_, Arc<io::Error>>((id, page.bytes.clone()))
251                    }))
252                })
253                .collect::<Vec<_>>();
254            // Poll every page before driving I/O, so executor fairness limits do not
255            // split one fetch's immediately available misses into separate batches.
256            let loads = future::poll_fn(move |cx| {
257                let mut complete = true;
258                for load in &mut loads {
259                    complete &= load.as_mut().poll(cx).is_ready();
260                    if let Some(Err(error)) = load.as_mut().output_mut() {
261                        return Poll::Ready(Err(error.clone()));
262                    }
263                }
264                if !complete {
265                    return Poll::Pending;
266                }
267                Poll::Ready(
268                    loads
269                        .iter_mut()
270                        .map(|load| load.as_mut().take_output().ok_or_else(invalid_request)?)
271                        .collect::<Result<Vec<_>, _>>(),
272                )
273            });
274            // The channel contains only loaders elected by Moka for this fetch. Waiters
275            // enqueue nothing; a waiter promoted after cancellation can enqueue later.
276            drop(requests);
277            let read_batches = async {
278                let mut batches = receiver.ready_chunks(batch_size);
279                while let Some(batch) = batches.next().await {
280                    let (ranges, replies): (Vec<_>, Vec<_>) = batch.into_iter().unzip();
281                    let result = reader.read_vec(&ranges).await.and_then(|bytes| {
282                        if bytes.len() == ranges.len() {
283                            Ok(bytes)
284                        } else {
285                            Err(io::Error::new(
286                                io::ErrorKind::InvalidData,
287                                "Invalid page response count",
288                            ))
289                        }
290                    });
291                    match result {
292                        Ok(bytes) => {
293                            for (reply, bytes) in replies.into_iter().zip(bytes) {
294                                let _ = reply.send(Ok(bytes));
295                            }
296                        }
297                        Err(error) => {
298                            let error = Arc::new(error);
299                            for reply in replies {
300                                let _ = reply.send(Err(error.clone()));
301                            }
302                        }
303                    }
304                }
305                Ok::<(), Arc<io::Error>>(())
306            };
307            let (loaded, ()) = future::try_join(loads, read_batches).await?;
308            pages.extend(loaded);
309        }
310        let mut output = Vec::with_capacity(ranges.len());
311        for range in ranges {
312            if range.is_empty() {
313                output.push(Bytes::new());
314                continue;
315            }
316            let ids = layout.page_ids(range.clone()).map_err(Arc::new)?;
317            if ids.end - ids.start == 1 {
318                let slice = layout
319                    .request_slice(ids.start, range.clone())
320                    .map_err(Arc::new)?
321                    .ok_or_else(invalid_request)?;
322                let page = pages.get(&ids.start).ok_or_else(invalid_request)?;
323                if slice.end > page.len() {
324                    return Err(invalid_request());
325                }
326                output.push(page.slice(slice));
327                continue;
328            }
329            let len = usize::try_from(range.end - range.start).map_err(|_| invalid_request())?;
330            let mut assembled = BytesMut::with_capacity(len);
331            for id in ids {
332                let slice = layout
333                    .request_slice(id, range.clone())
334                    .map_err(Arc::new)?
335                    .ok_or_else(invalid_request)?;
336                let page = pages.get(&id).ok_or_else(invalid_request)?;
337                if slice.end > page.len() {
338                    return Err(invalid_request());
339                }
340                assembled.extend_from_slice(&page[slice]);
341            }
342            if assembled.len() != len {
343                return Err(invalid_request());
344            }
345            output.push(assembled.freeze());
346        }
347        Ok((output, stats))
348    }
349
350    fn create_cached_page(&self, file_size: u64, bytes: Bytes) -> Arc<CachedPage> {
351        (self.observer.events)(PageEvent::Loaded {
352            bytes: bytes.len() as u64,
353        });
354        Arc::new(CachedPage { file_size, bytes })
355    }
356}
357
358struct CachedPage {
359    file_size: u64,
360    bytes: Bytes,
361}
362
363fn invalid_request() -> Arc<io::Error> {
364    Arc::new(io::Error::new(
365        io::ErrorKind::InvalidInput,
366        "Invalid cached page range",
367    ))
368}
369
370fn check_file_size(page: &CachedPage, expected: u64) -> Result<(), Arc<io::Error>> {
371    if page.file_size != expected {
372        return Err(Arc::new(io::Error::new(
373            io::ErrorKind::InvalidInput,
374            "Cached content identity used with inconsistent file sizes",
375        )));
376    }
377    Ok(())
378}
379
380/// Counts the union of request ranges already validated by fetch.
381fn unique_requested_bytes(ranges: &[Range<u64>]) -> u64 {
382    let mut ranges = ranges.to_vec();
383    ranges.sort_unstable_by_key(|range| range.start);
384    let mut end = 0;
385    let mut bytes = 0;
386    for range in ranges {
387        bytes += range.end.saturating_sub(range.start.max(end));
388        end = end.max(range.end);
389    }
390    bytes
391}
common_cache/paged.rs

common_cache/
paged.rs
