1use std::collections::{BTreeMap, HashMap, HashSet};
21use std::sync::{Arc, Mutex};
22
23use flint_store::{ObjectStore, StoreError};
24
25use super::chunk::{self, ChunkRef};
26use super::manifest::{
27 self, Entries, Hole, LeanEntry, LeanManifest, LoadedManifest, LoadedPointer, ManifestHandle, Pointer, Tombstone,
28};
29use super::{LeanConfig, LeanError, LeanResult};
30
31type EntryMap = BTreeMap<String, LeanEntry>;
32
33type Aliases = Arc<Vec<(String, String)>>;
36
37pub type AliasCache = Mutex<HashMap<String, Aliases>>;
40
41pub fn key_path<'a>(prefix: &str, key: &'a str) -> Option<&'a str> {
43 let rest = key.strip_prefix(prefix)?.strip_prefix("/files/")?;
44 Some(rest.rsplit_once('@')?.0)
45}
46
47fn aliases_in(prefix: &str, entries: &EntryMap) -> Vec<(String, String)> {
48 aliases_of(prefix, entries.iter())
49}
50
51fn aliases_of<'a>(prefix: &str, entries: impl Iterator<Item = (&'a String, &'a LeanEntry)>) -> Vec<(String, String)> {
52 entries
53 .filter(|(p, e)| key_path(prefix, &e.key) != Some(p.as_str()))
54 .map(|(p, e)| (p.clone(), e.key.clone()))
55 .collect()
56}
57
58pub fn chunk_entries<'a>(
61 doc: &'a EntryMap,
62 refs: &'a [ChunkRef],
63 i: usize,
64) -> impl Iterator<Item = (&'a String, &'a LeanEntry)> {
65 use std::ops::Bound;
66 let to = refs.get(i + 1).map_or(Bound::Unbounded, |n| Bound::Excluded(n.first.as_str()));
67 doc.range::<str, _>((Bound::Included(refs[i].first.as_str()), to))
68}
69
70pub fn retired_over(
75 cfg: &LeanConfig,
76 refs: &[ChunkRef],
77 doc: &EntryMap,
78 overlay: &BTreeMap<String, Option<LeanEntry>>,
79) -> Vec<String> {
80 let mut candidates: Vec<String> = vec![];
81 for (p, now) in overlay {
82 if let Some(was) = doc.get(p) {
83 if now.as_ref().is_none_or(|n| n.key != was.key) {
84 candidates.push(was.key.clone());
85 }
86 }
87 }
88 if candidates.is_empty() {
89 return vec![];
90 }
91 let aliases: Vec<Aliases> = {
92 let mut cache = cfg.alias_cache.lock().unwrap();
93 (0..refs.len())
94 .map(|i| {
95 cache
96 .entry(refs[i].addr.clone())
97 .or_insert_with(|| Arc::new(aliases_of(&cfg.prefix, chunk_entries(doc, refs, i))))
98 .clone()
99 })
100 .collect()
101 };
102 let mut by_key: HashMap<&str, Vec<&str>> = HashMap::new();
103 for (path, key) in aliases.iter().flat_map(|a| a.iter()) {
104 by_key.entry(key.as_str()).or_default().push(path.as_str());
105 }
106 let overlay_cites: HashSet<&str> = overlay.values().flatten().map(|e| e.key.as_str()).collect();
107 candidates
108 .into_iter()
109 .filter(|k| {
110 let cited = overlay_cites.contains(k.as_str())
111 || by_key.get(k.as_str()).is_some_and(|ps| ps.iter().any(|p| !overlay.contains_key(*p)))
112 || key_path(&cfg.prefix, k)
113 .filter(|q| !overlay.contains_key(*q))
114 .is_some_and(|q| doc.get(q).is_some_and(|e| &e.key == k));
115 !cited
116 })
117 .collect()
118}
119
120pub fn remember(cfg: &LeanConfig, refs: &[ChunkRef], entries: &EntryMap) {
125 let prefix = cfg.prefix.as_str();
126 let mut cache = cfg.alias_cache.lock().unwrap();
127 let mut it = entries.iter().peekable();
128 for (i, r) in refs.iter().enumerate() {
129 let next = refs.get(i + 1).map(|n| n.first.as_str());
130 let mut aliases = vec![];
131 while let Some((p, e)) = it.next_if(|(p, _)| next.is_none_or(|n| p.as_str() < n)) {
132 if key_path(prefix, &e.key) != Some(p.as_str()) {
133 aliases.push((p.clone(), e.key.clone()));
134 }
135 }
136 cache.entry(r.addr.clone()).or_insert_with(|| Arc::new(aliases));
137 }
138 let live: HashSet<&str> = refs.iter().map(|r| r.addr.as_str()).collect();
139 cache.retain(|a, _| live.contains(a.as_str()));
140}
141
142pub enum Opened {
144 View(DocView),
146 Full(Option<LoadedManifest>),
151}
152
153pub struct DocView {
154 store: Arc<dyn ObjectStore>,
155 cfg: LeanConfig,
156 pub pointer: Pointer,
157 pub etag: String,
158 pub last_modified_unix: Option<u64>,
159 pub refs: Vec<ChunkRef>,
160 pub tombstones: BTreeMap<String, Tombstone>,
161 parsed: HashMap<usize, Arc<EntryMap>>,
162 pub chunks_read: usize,
164}
165
166impl DocView {
167 pub async fn open(
172 store: Arc<dyn ObjectStore>,
173 cfg: &LeanConfig,
174 mut first: Option<Option<LoadedPointer>>,
175 ) -> LeanResult<Opened> {
176 for attempt in 0..manifest::LOAD_ATTEMPTS {
177 let read = match first.take() {
178 Some(already) => already,
179 None => manifest::load_pointer(store.as_ref(), cfg).await?,
180 };
181 let Some(lp) = read else {
182 return Ok(Opened::Full(manifest::load_from(store.as_ref(), cfg, Some(None)).await?));
183 };
184 let refs = match lp.pointer.entries()? {
185 Entries::Chunked(r) => r.to_vec(),
186 Entries::Single { .. } => {
187 return Ok(Opened::Full(manifest::load_from(store.as_ref(), cfg, Some(Some(lp))).await?));
188 }
189 };
190 chunk::check_list_order(&refs)?;
191 let listed = match store.list(&format!("{}/{}/chunks/", cfg.prefix, super::LEAN_DIR)).await {
192 Ok(l) => l,
193 Err(_) => return Ok(Opened::Full(manifest::load_from(store.as_ref(), cfg, Some(Some(lp))).await?)),
194 };
195 let present: HashSet<&str> = listed.iter().filter_map(|o| o.key.rsplit('/').next()).collect();
196 let hole = refs
197 .iter()
198 .find(|r| !present.contains(r.addr.as_str()))
199 .map(|r| Hole::Chunk { addr: r.addr.clone(), first: r.first.clone() })
200 .or_else(|| {
201 lp.pointer.tombstones.as_ref().filter(|a| !present.contains(a.as_str())).map(|a| Hole::Tombstones(a.clone()))
202 });
203 if let Some(hole) = hole {
204 let moved = match manifest::load_pointer(store.as_ref(), cfg).await? {
205 Some(now) => now.etag != lp.etag,
206 None => true,
207 };
208 if moved && attempt + 1 < manifest::LOAD_ATTEMPTS {
209 continue;
210 }
211 return Err(LeanError::State(hole.refusal(lp.pointer.seq)));
212 }
213 let tombstones = match lp.pointer.tombstones.clone() {
214 None => BTreeMap::new(),
215 Some(addr) => {
216 let body = fetch_body(store.as_ref(), cfg, &addr, lp.pointer.seq, || Hole::Tombstones(addr.clone()))
217 .await?;
218 chunk::decode_tombstones(&body, &addr)?
219 }
220 };
221 if let Some(c) = cfg.chunk_cache.as_deref() {
223 let mut live: HashSet<String> = refs.iter().map(|r| r.addr.clone()).collect();
224 live.extend(lp.pointer.tombstones.clone());
225 c.retain(&live);
226 }
227 return Ok(Opened::View(DocView {
228 store,
229 cfg: cfg.clone(),
230 pointer: lp.pointer,
231 etag: lp.etag,
232 last_modified_unix: lp.last_modified_unix,
233 refs,
234 tombstones,
235 parsed: HashMap::new(),
236 chunks_read: 0,
237 }));
238 }
239 unreachable!("the last attempt returns")
240 }
241
242 pub fn seq(&self) -> u64 {
243 self.pointer.seq
244 }
245
246 pub fn handle(&self) -> ManifestHandle {
248 ManifestHandle {
249 etag: self.etag.clone(),
250 legacy: false,
251 prev_chunks: self.refs.clone(),
252 prev_tombstones: self.pointer.tombstones.clone(),
253 }
254 }
255
256 pub fn home(&self, p: &str) -> usize {
258 self.refs.partition_point(|r| r.first.as_str() <= p).saturating_sub(1)
259 }
260
261 pub async fn chunk(&mut self, i: usize) -> LeanResult<Arc<EntryMap>> {
263 if let Some(c) = self.parsed.get(&i) {
264 return Ok(c.clone());
265 }
266 let r = self.refs[i].clone();
267 let seq = self.pointer.seq;
268 let body = fetch_body(self.store.as_ref(), &self.cfg, &r.addr, seq, || Hole::Chunk {
269 addr: r.addr.clone(),
270 first: r.first.clone(),
271 })
272 .await?;
273 let entries = Arc::new(chunk::parse_chunk(&self.refs, i, &body)?);
274 self.cfg
275 .alias_cache
276 .lock()
277 .unwrap()
278 .entry(r.addr.clone())
279 .or_insert_with(|| Arc::new(aliases_in(&self.cfg.prefix, &entries)));
280 self.chunks_read += 1;
281 self.parsed.insert(i, entries.clone());
282 Ok(entries)
283 }
284
285 pub async fn prefetch<'a>(&mut self, paths: impl IntoIterator<Item = &'a str>) -> LeanResult<()> {
288 if self.refs.is_empty() {
289 return Ok(());
290 }
291 let want: std::collections::BTreeSet<usize> =
292 paths.into_iter().map(|p| self.home(p)).filter(|i| !self.parsed.contains_key(i)).collect();
293 self.read_chunks(want.into_iter().collect()).await
294 }
295
296 async fn read_chunks(&mut self, want: Vec<usize>) -> LeanResult<()> {
297 use futures::stream::{self, StreamExt};
298 let seq = self.pointer.seq;
299 let (store, cfg) = (self.store.clone(), self.cfg.clone());
300 let got: Vec<(usize, LeanResult<bytes::Bytes>)> = stream::iter(want.into_iter().map(|i| {
301 let r = self.refs[i].clone();
302 let (store, cfg) = (store.clone(), cfg.clone());
303 async move {
304 let hole = || Hole::Chunk { addr: r.addr.clone(), first: r.first.clone() };
305 (i, fetch_body(store.as_ref(), &cfg, &r.addr, seq, hole).await)
306 }
307 }))
308 .buffer_unordered(self.cfg.chunk_fetch_parallelism.max(1))
309 .collect()
310 .await;
311 for (i, body) in got {
312 let entries = Arc::new(chunk::parse_chunk(&self.refs, i, &body?)?);
313 self.cfg
314 .alias_cache
315 .lock()
316 .unwrap()
317 .entry(self.refs[i].addr.clone())
318 .or_insert_with(|| Arc::new(aliases_in(&self.cfg.prefix, &entries)));
319 self.chunks_read += 1;
320 self.parsed.insert(i, entries);
321 }
322 Ok(())
323 }
324
325 pub async fn get(&mut self, p: &str) -> LeanResult<Option<LeanEntry>> {
326 if self.refs.is_empty() {
327 return Ok(None);
328 }
329 let i = self.home(p);
330 Ok(self.chunk(i).await?.get(p).cloned())
331 }
332
333 pub async fn first_under(&mut self, prefix: &str, skip: &dyn Fn(&str) -> bool) -> LeanResult<Option<String>> {
336 if self.refs.is_empty() {
337 return Ok(None);
338 }
339 let mut i = self.home(prefix);
340 while i < self.refs.len() {
341 if i > self.home(prefix) && !self.refs[i].first.starts_with(prefix) {
342 return Ok(None);
343 }
344 let c = self.chunk(i).await?;
345 for k in c.range::<str, _>((std::ops::Bound::Included(prefix), std::ops::Bound::Unbounded)).map(|(k, _)| k) {
346 if !k.starts_with(prefix) {
347 return Ok(None);
348 }
349 if !skip(k) {
350 return Ok(Some(k.clone()));
351 }
352 }
353 i += 1;
354 }
355 Ok(None)
356 }
357
358 pub async fn rechunk(
361 &mut self,
362 changes: &BTreeMap<String, Option<LeanEntry>>,
363 ) -> LeanResult<Vec<(ChunkRef, Option<Vec<u8>>)>> {
364 self.prefetch(changes.keys().map(|k| k.as_str())).await?;
365 loop {
366 let mut need: Option<usize> = None;
367 let parsed = &self.parsed;
368 let out = chunk::rechunk_with(
369 &self.refs,
370 |i| match parsed.get(&i) {
371 Some(c) => Ok(c.clone()),
372 None => {
373 need = Some(i);
374 Err(LeanError::State(format!("chunk {i} not read yet")))
375 }
376 },
377 changes,
378 self.cfg.chunk_target,
379 self.cfg.chunk_min,
380 self.cfg.chunk_max,
381 );
382 match (out, need) {
383 (Ok(r), _) => {
384 let mut cache = self.cfg.alias_cache.lock().unwrap();
387 for (c, body) in &r {
388 if let Some(b) = body {
389 let key = c.addr.clone();
390 if !cache.contains_key(&key) {
391 let body: chunk::ChunkBody = serde_json::from_slice(b)
392 .map_err(|e| LeanError::State(format!("chunk {} parse: {e}", c.addr)))?;
393 cache.insert(key, Arc::new(aliases_in(&self.cfg.prefix, &body.entries)));
394 }
395 }
396 }
397 return Ok(r);
398 }
399 (Err(_), Some(i)) => {
400 self.chunk(i).await?;
401 }
402 (Err(e), None) => return Err(e),
403 }
404 }
405 }
406
407 pub async fn retired(&mut self, overlay: &BTreeMap<String, Option<LeanEntry>>) -> LeanResult<Vec<String>> {
413 let mut candidates: Vec<String> = vec![];
414 for (p, now) in overlay {
415 if let Some(was) = self.get(p).await? {
416 if now.as_ref().is_none_or(|n| n.key != was.key) {
417 candidates.push(was.key);
418 }
419 }
420 }
421 if candidates.is_empty() {
422 return Ok(vec![]);
423 }
424 let aliases = self.aliases().await?;
425 let mut by_key: HashMap<&str, Vec<&str>> = HashMap::new();
426 for (path, key) in &aliases {
427 by_key.entry(key.as_str()).or_default().push(path.as_str());
428 }
429 let overlay_cites: HashSet<&str> = overlay.values().flatten().map(|e| e.key.as_str()).collect();
430 let mut out = vec![];
431 for k in candidates {
432 let cited = overlay_cites.contains(k.as_str())
433 || by_key.get(k.as_str()).is_some_and(|ps| ps.iter().any(|p| !overlay.contains_key(*p)))
434 || match key_path(&self.cfg.prefix, &k).filter(|q| !overlay.contains_key(*q)) {
435 Some(q) => self.get(q).await?.is_some_and(|e| e.key == k),
436 None => false,
437 };
438 if !cited {
439 out.push(k);
440 }
441 }
442 Ok(out)
443 }
444
445 pub async fn aliases(&mut self) -> LeanResult<Vec<(String, String)>> {
451 let prefix = self.cfg.prefix.clone();
452 let mut out: Vec<(String, String)> = vec![];
453 let mut missing: Vec<usize> = vec![];
454 {
455 let cache = self.cfg.alias_cache.lock().unwrap();
456 for (i, r) in self.refs.iter().enumerate() {
457 match cache.get(&r.addr) {
458 Some(a) => out.extend(a.iter().cloned()),
459 None => missing.push(i),
460 }
461 }
462 }
463 let unread: Vec<usize> = missing.iter().copied().filter(|i| !self.parsed.contains_key(i)).collect();
464 self.read_chunks(unread).await?;
465 for i in missing {
466 out.extend(aliases_in(&prefix, &self.parsed[&i]));
467 }
468 Ok(out)
469 }
470
471 pub fn retain_aliases(&self, live: &HashSet<&str>) {
474 self.cfg.alias_cache.lock().unwrap().retain(|a, _| live.contains(a.as_str()));
475 }
476
477 pub async fn assemble(&mut self) -> LeanResult<LoadedManifest> {
480 let missing: Vec<usize> = (0..self.refs.len()).filter(|i| !self.parsed.contains_key(i)).collect();
481 self.read_chunks(missing).await?;
482 let mut entries = BTreeMap::new();
483 for i in 0..self.refs.len() {
484 for (k, v) in self.parsed[&i].iter() {
485 entries.insert(k.clone(), v.clone());
486 }
487 }
488 Ok(LoadedManifest {
489 manifest: LeanManifest {
490 seq: self.pointer.seq,
491 entries,
492 sole_writer: self.pointer.sole_writer,
493 boundary_source: self.pointer.boundary_source.clone(),
494 tombstones: self.tombstones.clone(),
495 },
496 etag: self.etag.clone(),
497 last_modified_unix: self.last_modified_unix,
498 pointer: Some(self.pointer.clone()),
499 })
500 }
501}
502
503async fn fetch_body(
508 store: &dyn ObjectStore,
509 cfg: &LeanConfig,
510 addr: &str,
511 seq: u64,
512 hole: impl Fn() -> Hole,
513) -> LeanResult<bytes::Bytes> {
514 let cache = cfg.chunk_cache.as_deref();
515 if let Some(b) = cache.and_then(|c| c.get(addr)) {
516 if chunk::chunk_address(&b) == addr {
517 return Ok(b);
518 }
519 if let Some(c) = cache {
520 c.remove(addr);
521 }
522 }
523 let body = match store.get_whole(&cfg.chunk_key(addr), None).await {
524 Ok((_, b)) => b,
525 Err(StoreError::NotFound(_)) => return Err(LeanError::State(hole().refusal(seq))),
526 Err(e) => return Err(e.into()),
527 };
528 chunk::check_address_of(&body, addr)?;
529 if let Some(c) = cache {
530 c.put(addr, &body);
531 }
532 Ok(body)
533}
534
535pub enum CommitDocs {
538 Full { theirs: LeanManifest, installed: LeanManifest },
539 View {
540 theirs: DocView,
541 overlay: BTreeMap<String, Option<LeanEntry>>,
542 retired: Vec<String>,
544 seq: u64,
545 tombstones: BTreeMap<String, Tombstone>,
546 sole_writer: bool,
547 boundary_source: Option<String>,
548 },
549}
550
551impl CommitDocs {
552 pub fn installed_seq(&self) -> u64 {
553 match self {
554 CommitDocs::Full { installed, .. } => installed.seq,
555 CommitDocs::View { seq, .. } => *seq,
556 }
557 }
558
559 pub fn retired(&self) -> Vec<String> {
562 match self {
563 CommitDocs::Full { theirs, installed } => {
564 let still_cited: HashSet<&str> = installed.entries.values().map(|e| e.key.as_str()).collect();
565 let mut retired: Vec<String> = vec![];
566 for (path, was) in &theirs.entries {
567 let dropped = installed.entries.get(path).map(|now| now.key != was.key).unwrap_or(true);
568 if dropped && !still_cited.contains(was.key.as_str()) {
569 retired.push(was.key.clone());
570 }
571 }
572 retired
573 }
574 CommitDocs::View { retired, .. } => retired.clone(),
575 }
576 }
577
578 pub async fn installed_get(&mut self, p: &str) -> LeanResult<Option<LeanEntry>> {
579 match self {
580 CommitDocs::Full { installed, .. } => Ok(installed.entries.get(p).cloned()),
581 CommitDocs::View { theirs, overlay, .. } => match overlay.get(p) {
582 Some(v) => Ok(v.clone()),
583 None => theirs.get(p).await,
584 },
585 }
586 }
587
588 pub async fn theirs_get(&mut self, p: &str) -> LeanResult<Option<LeanEntry>> {
589 match self {
590 CommitDocs::Full { theirs, .. } => Ok(theirs.entries.get(p).cloned()),
591 CommitDocs::View { theirs, .. } => theirs.get(p).await,
592 }
593 }
594
595 pub async fn installed_at(&mut self, paths: &[String]) -> LeanResult<LeanManifest> {
598 let mut entries = BTreeMap::new();
599 if let CommitDocs::View { theirs, .. } = self {
600 theirs.prefetch(paths.iter().map(|p| p.as_str())).await?;
601 }
602 for p in paths {
603 if let Some(e) = self.installed_get(p).await? {
604 entries.insert(p.clone(), e);
605 }
606 }
607 Ok(LeanManifest { seq: self.installed_seq(), entries, ..Default::default() })
608 }
609
610 pub async fn installed_full(&mut self) -> LeanResult<std::borrow::Cow<'_, LeanManifest>> {
612 match self {
613 CommitDocs::Full { installed, .. } => Ok(std::borrow::Cow::Borrowed(installed)),
614 CommitDocs::View { theirs, overlay, seq, tombstones, sole_writer, boundary_source, .. } => {
615 let mut m = theirs.assemble().await?.manifest;
616 for (p, v) in overlay.iter() {
617 match v {
618 Some(e) => {
619 m.entries.insert(p.clone(), e.clone());
620 }
621 None => {
622 m.entries.remove(p);
623 }
624 }
625 }
626 m.seq = *seq;
627 m.tombstones = tombstones.clone();
628 m.sole_writer = *sole_writer;
629 m.boundary_source = boundary_source.clone();
630 Ok(std::borrow::Cow::Owned(m))
631 }
632 }
633 }
634}