micromegas_object_cache/
prefetch.rs

1use async_trait::async_trait;
2use object_store::path::Path;
3use serde::{Deserialize, Serialize};
4use std::sync::Arc;
5
6/// Wire format for `POST /prefetch`: `Content-Type: application/x-ndjson`, one
7/// `PrefetchItem` JSON object per `\n`-terminated line. There is no wrapper
8/// type — the body is a stream of lines, parsed and enqueued incrementally so
9/// the server never buffers the whole batch.
10///
11/// One key to warm, whole-object or ranged.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct PrefetchItem {
14    pub key: String,
15    /// The object's file size, supplied by the caller. Both triggers already know
16    /// it: `Partition.file_size` (persisted in PostgreSQL) for query/write warming,
17    /// and `PartitionWriteResult.file_size` for the write path. Supplying it lets
18    /// the server drive fills through `prefetch_blocks(key, file_size, indices)`
19    /// with no origin HEAD (prefetch targets cold objects, so a server-side
20    /// `size()` would force an avoidable HEAD).
21    ///
22    /// Contract: this must be the object's exact current size. The server
23    /// trusts it without verification; an undersized value stores a truncated
24    /// final block under the same block key demand reads use. `RangeCache`
25    /// mitigates this with a hit-path length guard that detects and refetches a
26    /// wrong-length block on the next correctly-sized read. An oversized value
27    /// is safe — the origin GET past EOF fails and nothing is stored.
28    pub size: u64,
29    /// None or empty = warm the whole object `[0, size)`. Present = warm only these
30    /// ranges (validated against `size`).
31    #[serde(default, skip_serializing_if = "Option::is_none")]
32    pub ranges: Option<Vec<[u64; 2]>>,
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct PrefetchResponse {
37    /// Enqueued onto the prefetch queue.
38    pub accepted: usize,
39    /// Failed key/prefix/range validation, skipped.
40    pub rejected: usize,
41    /// Queue full, load-shed.
42    pub dropped: usize,
43}
44
45/// Dyn-compatible capability seam so downstream consumers that hold
46/// `Arc<dyn ObjectStore>` (not `CacheClientStore` directly) can still drive
47/// prefetch without downcasting.
48///
49/// `Debug` is a supertrait so `Arc<dyn ObjectPrefetch>` can be embedded in a
50/// `#[derive(Debug)]` struct (e.g. `DataLakeConnection`).
51#[async_trait]
52pub trait ObjectPrefetch: Send + Sync + std::fmt::Debug {
53    async fn prefetch(&self, items: Vec<PrefetchItem>) -> anyhow::Result<PrefetchResponse>;
54}
55
56/// Prepends a root prefix to each `PrefetchItem`'s key before delegating, so a
57/// warm keyed by a lake-root-relative path (`views/…`) targets the same cache
58/// key a demand read produces through `object_store::PrefixStore` (`root/views/…`).
59/// This mirrors, for the prefetch path, what `PrefixStore` does for reads.
60#[derive(Debug)]
61pub struct PrefixPrefetch {
62    inner: Arc<dyn ObjectPrefetch>,
63    prefix: Path,
64}
65
66impl PrefixPrefetch {
67    pub fn new(inner: Arc<dyn ObjectPrefetch>, prefix: Path) -> Self {
68        Self { inner, prefix }
69    }
70}
71
72#[async_trait]
73impl ObjectPrefetch for PrefixPrefetch {
74    async fn prefetch(&self, mut items: Vec<PrefetchItem>) -> anyhow::Result<PrefetchResponse> {
75        for item in &mut items {
76            // Compose exactly as PrefixStore does: chain the prefix's path parts
77            // with the key's parts, so the resulting string equals the key
78            // CacheClientStore.get() sees for a demand read of the same object.
79            let full = Path::from_iter(
80                self.prefix
81                    .parts()
82                    .chain(Path::from(item.key.as_str()).parts()),
83            );
84            item.key = full.as_ref().to_string();
85        }
86        self.inner.prefetch(items).await
87    }
88}
micromegas_object_cache/prefetch.rs

micromegas_object_cache/
prefetch.rs
