micromegas_object_cache/
l1_store.rs

1use std::ops::Range;
2use std::sync::{Arc, OnceLock};
3
4use async_trait::async_trait;
5use bytes::Bytes;
6use futures::stream::{self, BoxStream};
7use micromegas_tracing::prelude::*;
8use object_store::{
9    Attributes, CopyOptions, GetOptions, GetRange, GetResult, GetResultPayload, ListResult,
10    MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload,
11    PutResult, path::Path,
12};
13
14use crate::backend::RangeCacheBackend;
15use crate::bounded_memory_backend::BoundedMemoryBackend;
16use crate::range_cache::{
17    DEFAULT_BLOCK_SIZE, DEFAULT_MAX_COALESCED_GET_BYTES, DEFAULT_PROMOTE_WHOLE_BATCH, RangeCache,
18};
19
20/// Environment variable sizing the shared in-process L1 RAM budget, in MB.
21/// `0` disables L1 (`l1_wrap` then returns the origin store unchanged).
22/// Part of the `MICROMEGAS_OBJECT_CACHE_*` family: `_L1_MB` sizes this
23/// in-process tier, `_RAM_MB` sizes the `object-cache-srv`'s own RAM tier.
24const ENV_OBJECT_CACHE_L1_MB: &str = "MICROMEGAS_OBJECT_CACHE_L1_MB";
25/// Matches the old whole-file `FileCache`'s default budget.
26const DEFAULT_OBJECT_CACHE_L1_MB: u64 = 200;
27
28/// Total number of origin GETs the in-process L1 cache allows concurrently.
29/// L1 issues only demand reads (`RangeCache::get_range`/`get_ranges`), so
30/// this is the only fetch-permit knob that matters here (see
31/// `L1_DEMAND_RESERVED_FETCH_PERMITS`). Sized smaller than the dedicated
32/// object-cache-srv's default (32): an in-process L1 shares the process with
33/// DataFusion's own working set, so a smaller concurrent-fetch /
34/// transient-buffer footprint (see `L1_FETCH_MEMORY_BUDGET_BYTES`) is
35/// preferable.
36const L1_TOTAL_FETCH_PERMITS: usize = 16;
37/// Placeholder satisfying `RangeCache::new`'s `demand_reserved < total`
38/// assertion. It sizes the prefetch-only pools of both fetch budgets -- the
39/// count pool, and the byte pool via a derived reservation of
40/// `L1_DEMAND_RESERVED_FETCH_PERMITS * max_run_bytes` (32 MiB here) -- so it
41/// must also leave `L1_FETCH_MEMORY_BUDGET_BYTES` at least one max run above
42/// that reservation. L1 never issues prefetch reads, so demand always draws
43/// on the full pools and the value has no runtime effect.
44const L1_DEMAND_RESERVED_FETCH_PERMITS: usize = 4;
45/// Transient origin-GET buffer memory bound for one `L1CacheStore`'s
46/// `RangeCache` (128 MiB at these defaults): `l1_wrap` builds a separate
47/// `L1CacheStore`/`RangeCache` per call (see `l1_wrap` below), so this bounds
48/// one `L1CacheStore` instance's fetch buffers, not the whole process.
49const L1_FETCH_MEMORY_BUDGET_BYTES: u64 =
50    L1_TOTAL_FETCH_PERMITS as u64 * DEFAULT_MAX_COALESCED_GET_BYTES;
51
52/// The shared in-process L1 RAM budget, lazily sized once (on first use) from
53/// `MICROMEGAS_OBJECT_CACHE_L1_MB`. `None` when L1 is disabled (budget `0`).
54static SHARED_L1_BACKEND: OnceLock<Option<Arc<BoundedMemoryBackend>>> = OnceLock::new();
55
56fn shared_l1_backend() -> Option<Arc<BoundedMemoryBackend>> {
57    SHARED_L1_BACKEND
58        .get_or_init(|| {
59            let mb = match std::env::var(ENV_OBJECT_CACHE_L1_MB) {
60                Ok(s) => s.parse::<u64>().unwrap_or_else(|_| {
61                    warn!(
62                        "Invalid {ENV_OBJECT_CACHE_L1_MB} value '{s}', using default {DEFAULT_OBJECT_CACHE_L1_MB} MB"
63                    );
64                    DEFAULT_OBJECT_CACHE_L1_MB
65                }),
66                Err(_) => DEFAULT_OBJECT_CACHE_L1_MB,
67            };
68            if mb == 0 {
69                info!("{ENV_OBJECT_CACHE_L1_MB}=0, in-process L1 cache disabled");
70                None
71            } else {
72                info!("in-process L1 cache enabled, budget={mb}MB");
73                Some(Arc::new(BoundedMemoryBackend::new(
74                    (mb * 1024 * 1024) as usize,
75                )))
76            }
77        })
78        .clone()
79}
80
81/// Wrap `origin` with the in-process L1 cache: a `RangeCache` over the shared
82/// bounded RAM backend, namespaced by `ns` (e.g. `"lakehouse"`, `"static"`) so
83/// distinct wrap sites share one RAM budget without their keys colliding.
84///
85/// Returns `origin` unchanged when L1 is disabled (`MICROMEGAS_OBJECT_CACHE_L1_MB=0`).
86pub fn l1_wrap(origin: Arc<dyn ObjectStore>, ns: &str) -> Arc<dyn ObjectStore> {
87    match shared_l1_backend() {
88        Some(backend) => Arc::new(L1CacheStore::new(origin, backend, ns.to_string())),
89        None => origin,
90    }
91}
92
93/// An `ObjectStore` adapter fronting `origin` with a RAM-backed `RangeCache`.
94///
95/// `get_opts`/`get_ranges` are served from the cache (falling back to
96/// `origin` on any cache error, mirroring `CacheClientStore`'s
97/// graceful-degradation contract); everything else -- `put`, `list`,
98/// `delete`, `copy`, and preconditioned or `head` gets -- passes straight
99/// through to `origin`, which is never itself cached by this store.
100pub struct L1CacheStore {
101    cache: RangeCache,
102    origin: Arc<dyn ObjectStore>,
103}
104
105impl L1CacheStore {
106    pub fn new(
107        origin: Arc<dyn ObjectStore>,
108        backend: Arc<dyn RangeCacheBackend>,
109        ns: String,
110    ) -> Self {
111        let cache = RangeCache::new(
112            origin.clone(),
113            backend,
114            DEFAULT_BLOCK_SIZE,
115            ns,
116            L1_TOTAL_FETCH_PERMITS,
117            L1_DEMAND_RESERVED_FETCH_PERMITS,
118            L1_FETCH_MEMORY_BUDGET_BYTES,
119            DEFAULT_MAX_COALESCED_GET_BYTES,
120            DEFAULT_PROMOTE_WHOLE_BATCH,
121        );
122        Self { cache, origin }
123    }
124
125    fn key(location: &Path) -> String {
126        location.as_ref().to_string()
127    }
128
129    async fn fallback_get_opts(
130        &self,
131        location: &Path,
132        options: GetOptions,
133        error: anyhow::Error,
134    ) -> object_store::Result<GetResult> {
135        imetric!("l1_cache_fallback", "count", 1_u64);
136        debug!("L1 cache miss for {location}, falling back to origin: {error}");
137        self.origin.get_opts(location, options).await
138    }
139}
140
141impl std::fmt::Debug for L1CacheStore {
142    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
143        write!(f, "L1CacheStore({})", self.origin)
144    }
145}
146
147impl std::fmt::Display for L1CacheStore {
148    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
149        write!(f, "L1CacheStore({})", self.origin)
150    }
151}
152
153/// Build a single-chunk streaming `GetResult` from already-fetched bytes.
154fn single_chunk_get_result(
155    data: Bytes,
156    range: Range<u64>,
157    object_size: u64,
158    location: &Path,
159) -> GetResult {
160    let meta = ObjectMeta {
161        location: location.clone(),
162        last_modified: chrono::Utc::now(),
163        size: object_size,
164        e_tag: None,
165        version: None,
166    };
167    let payload = GetResultPayload::Stream(Box::pin(stream::once(async move { Ok(data) })));
168    GetResult {
169        payload,
170        meta,
171        range,
172        attributes: Attributes::default(),
173    }
174}
175
176#[async_trait]
177impl ObjectStore for L1CacheStore {
178    async fn put_opts(
179        &self,
180        location: &Path,
181        payload: PutPayload,
182        opts: PutOptions,
183    ) -> object_store::Result<PutResult> {
184        self.origin.put_opts(location, payload, opts).await
185    }
186
187    async fn put_multipart_opts(
188        &self,
189        location: &Path,
190        opts: PutMultipartOptions,
191    ) -> object_store::Result<Box<dyn MultipartUpload>> {
192        self.origin.put_multipart_opts(location, opts).await
193    }
194
195    async fn get_opts(
196        &self,
197        location: &Path,
198        options: GetOptions,
199    ) -> object_store::Result<GetResult> {
200        // The L1 protocol carries no conditional/version preconditions and a
201        // head-only request needs no bytes: both go straight to origin.
202        if options.if_match.is_some()
203            || options.if_none_match.is_some()
204            || options.if_modified_since.is_some()
205            || options.if_unmodified_since.is_some()
206            || options.version.is_some()
207            || options.head
208        {
209            return self.origin.get_opts(location, options).await;
210        }
211
212        let key = Self::key(location);
213        let range = options.range.clone();
214
215        let resolved: anyhow::Result<(Bytes, Range<u64>, u64)> = async {
216            match &range {
217                Some(GetRange::Bounded(r)) => {
218                    let r = r.clone();
219                    let data = self.cache.get_range(&key, r.clone()).await?;
220                    // Resolve the true object size so `meta.size` matches
221                    // what real object stores report even for ranged gets
222                    // (our only in-tree caller, `ObjectStore::get_range`'s
223                    // default impl, reads via `.bytes()` and only consults
224                    // `range`, but callers of `get_opts` directly are
225                    // entitled to a correct `meta.size`).
226                    let size = self.cache.size(&key).await?;
227                    Ok((data, r, size))
228                }
229                other => {
230                    let size = self.cache.size(&key).await?;
231                    let resolved_range = match other {
232                        None => 0..size,
233                        Some(GetRange::Offset(offset)) => *offset..size,
234                        Some(GetRange::Suffix(suffix)) => size.saturating_sub(*suffix)..size,
235                        Some(GetRange::Bounded(_)) => unreachable!("handled above"),
236                    };
237                    let data = self.cache.get_range(&key, resolved_range.clone()).await?;
238                    Ok((data, resolved_range, size))
239                }
240            }
241        }
242        .await;
243
244        match resolved {
245            Ok((data, range, size)) => Ok(single_chunk_get_result(data, range, size, location)),
246            Err(e) => self.fallback_get_opts(location, options, e).await,
247        }
248    }
249
250    async fn get_ranges(
251        &self,
252        location: &Path,
253        ranges: &[Range<u64>],
254    ) -> object_store::Result<Vec<Bytes>> {
255        if ranges.is_empty() {
256            return Ok(vec![]);
257        }
258        let key = Self::key(location);
259        match self.cache.get_ranges(&key, ranges).await {
260            Ok(results) => Ok(results),
261            Err(e) => {
262                imetric!("l1_cache_fallback", "count", 1_u64);
263                debug!("L1 cache miss for {location} (ranges), falling back to origin: {e}");
264                self.origin.get_ranges(location, ranges).await
265            }
266        }
267    }
268
269    fn delete_stream(
270        &self,
271        locations: BoxStream<'static, object_store::Result<Path>>,
272    ) -> BoxStream<'static, object_store::Result<Path>> {
273        self.origin.delete_stream(locations)
274    }
275
276    fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, object_store::Result<ObjectMeta>> {
277        self.origin.list(prefix)
278    }
279
280    async fn list_with_delimiter(&self, prefix: Option<&Path>) -> object_store::Result<ListResult> {
281        self.origin.list_with_delimiter(prefix).await
282    }
283
284    async fn copy_opts(
285        &self,
286        from: &Path,
287        to: &Path,
288        options: CopyOptions,
289    ) -> object_store::Result<()> {
290        self.origin.copy_opts(from, to, options).await
291    }
292}
micromegas_object_cache/l1_store.rs

micromegas_object_cache/
l1_store.rs
