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

micromegas_object_cache/
l1_store.rs
