1use std::path::{Path, PathBuf};
12
13#[cfg(test)]
14use std::cell::Cell;
15
16use readcon_core::types::ConFrame;
17
18use crate::corpus::ConCorpus;
19use crate::error::{Error, Result};
20use crate::keys::{FrameKey, TrajId};
21use crate::select::Select;
22
23pub const DEFAULT_N_SHARDS: u32 = 64;
25
26const MANIFEST: &str = "shards.json";
28
29#[cfg(test)]
30thread_local! {
31 pub(crate) static FAIL_DEST_MAN_COPY: Cell<bool> = const { Cell::new(false) };
32}
33
34#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
35pub struct ShardManifest {
36 pub n_shards: u32,
37 pub version: u32,
38}
39
40pub struct ShardedConCorpus {
42 root: PathBuf,
43 n_shards: u32,
44 shards: Vec<Option<ConCorpus>>,
46}
47
48impl ShardedConCorpus {
49 pub fn open_existing(root: impl AsRef<Path>) -> Result<Self> {
51 let root = root.as_ref();
52 if !root.join(MANIFEST).is_file() {
53 return Err(Error::Message(format!(
54 "missing shards.json: {}",
55 root.display()
56 )));
57 }
58 Self::open(root, 1)
59 }
60
61 pub fn open(root: impl AsRef<Path>, n_shards: u32) -> Result<Self> {
63 let root = root.as_ref().to_path_buf();
64 std::fs::create_dir_all(&root)?;
65 let manifest_path = root.join(MANIFEST);
66 let n_shards = if manifest_path.is_file() {
67 let s = std::fs::read_to_string(&manifest_path)?;
68 let m: ShardManifest = serde_json::from_str(&s)?;
69 m.n_shards
70 } else {
71 if n_shards == 0 {
72 return Err(Error::Message("n_shards must be >= 1".into()));
73 }
74 let m = ShardManifest {
75 n_shards,
76 version: 1,
77 };
78 std::fs::write(&manifest_path, serde_json::to_string_pretty(&m)?)?;
79 n_shards
80 };
81 let mut shards = Vec::with_capacity(n_shards as usize);
82 shards.resize_with(n_shards as usize, || None);
83 Ok(Self {
84 root,
85 n_shards,
86 shards,
87 })
88 }
89
90 pub fn canonical_dump(&mut self) -> Result<Vec<u8>> {
93 let mut out = Vec::new();
94 for sid in 0..self.n_shards {
95 if !self.shard_has_data(sid) {
96 continue;
97 }
98 let d = self.shard_mut(sid)?.canonical_dump()?;
99 out.extend_from_slice(&sid.to_le_bytes());
100 out.extend_from_slice(&(d.len() as u64).to_le_bytes());
101 out.extend_from_slice(&d);
102 }
103 Ok(out)
104 }
105
106 pub fn n_shards(&self) -> u32 {
107 self.n_shards
108 }
109
110 pub fn root(&self) -> &Path {
111 &self.root
112 }
113
114 #[inline]
115 pub fn shard_for_traj(traj_id: TrajId, n_shards: u32) -> u32 {
116 (traj_id % u64::from(n_shards)) as u32
117 }
118
119 fn shard_path(&self, shard_id: u32) -> PathBuf {
120 self.root.join(format!("shard_{shard_id:04}"))
121 }
122
123 fn shard_has_data(&self, shard_id: u32) -> bool {
124 self.shard_path(shard_id).join("data.mdb").is_file()
125 }
126
127 fn release_shards(&mut self) {
128 for slot in &mut self.shards {
129 if let Some(c) = slot.take() {
130 c.close();
131 }
132 }
133 }
134
135 pub fn shard_mut(&mut self, shard_id: u32) -> Result<&ConCorpus> {
137 if shard_id >= self.n_shards {
138 return Err(Error::Message(format!(
139 "shard_id {shard_id} >= n_shards {}",
140 self.n_shards
141 )));
142 }
143 let i = shard_id as usize;
144 if self.shards[i].is_none() {
145 let p = self.shard_path(shard_id);
146 self.shards[i] = Some(ConCorpus::open(p)?);
147 }
148 Ok(self.shards[i].as_ref().unwrap())
149 }
150
151 pub fn open_shard_for_traj(
153 root: impl AsRef<Path>,
154 traj_id: TrajId,
155 ) -> Result<(u32, ConCorpus)> {
156 let root = root.as_ref();
157 let manifest_path = root.join(MANIFEST);
158 let n_shards = if manifest_path.is_file() {
159 let m: ShardManifest = serde_json::from_str(&std::fs::read_to_string(&manifest_path)?)?;
160 m.n_shards
161 } else {
162 let _ = Self::open(root, DEFAULT_N_SHARDS)?;
163 DEFAULT_N_SHARDS
164 };
165 let sid = Self::shard_for_traj(traj_id, n_shards);
166 let corpus = ConCorpus::open(root.join(format!("shard_{sid:04}")))?;
167 Ok((sid, corpus))
168 }
169
170 pub fn open_shard(root: impl AsRef<Path>, shard_id: u32) -> Result<ConCorpus> {
173 let root = root.as_ref();
174 let manifest_path = root.join(MANIFEST);
175 if !manifest_path.is_file() {
176 return Err(Error::Message(format!(
177 "missing shards.json: {}",
178 root.display()
179 )));
180 }
181 let m: ShardManifest = serde_json::from_str(&std::fs::read_to_string(&manifest_path)?)?;
182 let n_shards = m.n_shards;
183 if shard_id >= n_shards {
184 return Err(Error::Message(format!(
185 "shard_id {shard_id} >= n_shards {n_shards}"
186 )));
187 }
188 ConCorpus::open(root.join(format!("shard_{shard_id:04}")))
189 }
190
191 pub fn append_trajectory_path(
192 &mut self,
193 traj_id: TrajId,
194 file: impl AsRef<Path>,
195 ) -> Result<u32> {
196 self.append_trajectory_path_units(traj_id, file, None)
197 }
198
199 pub fn append_trajectory_path_units(
200 &mut self,
201 traj_id: TrajId,
202 file: impl AsRef<Path>,
203 units: Option<serde_json::Value>,
204 ) -> Result<u32> {
205 let sid = Self::shard_for_traj(traj_id, self.n_shards);
206 let c = self.shard_mut(sid)?;
207 c.append_trajectory_path_units(traj_id, file, units)
208 }
209
210 pub fn append_trajectory_str(
211 &mut self,
212 traj_id: TrajId,
213 contents: &str,
214 source: impl Into<String>,
215 ) -> Result<u32> {
216 let sid = Self::shard_for_traj(traj_id, self.n_shards);
217 let c = self.shard_mut(sid)?;
218 c.append_trajectory_str(traj_id, contents, source)
219 }
220
221 pub fn append_trajectory_frames(
222 &mut self,
223 traj_id: TrajId,
224 frames: &[ConFrame],
225 source: impl Into<String>,
226 ) -> Result<u32> {
227 let sid = Self::shard_for_traj(traj_id, self.n_shards);
228 let c = self.shard_mut(sid)?;
229 c.append_trajectory_frames(traj_id, frames, source)
230 }
231
232 pub fn select(&mut self, sel: &Select) -> Result<Vec<FrameKey>> {
234 self.release_shards();
235 let mut out = Vec::new();
236 for sid in 0..self.n_shards {
237 if !self.shard_has_data(sid) {
238 continue;
239 }
240 let c = ConCorpus::open_readonly(self.shard_path(sid))?;
241 out.extend(c.select(sel)?);
242 c.close();
243 }
244 out.sort();
245 if let Some(lim) = sel.limit {
246 out.truncate(lim);
247 }
248 Ok(out)
249 }
250
251 pub fn get_frame_text(&mut self, key: FrameKey) -> Result<String> {
252 self.release_shards();
253 let sid = Self::shard_for_traj(key.traj_id, self.n_shards);
254 if !self.shard_has_data(sid) {
255 return Err(Error::Message(format!(
256 "shard_{sid:04} is not a corpus directory"
257 )));
258 }
259 let c = ConCorpus::open_readonly(self.shard_path(sid))?;
260 let text = c.get_frame_text(key)?;
261 c.close();
262 Ok(text)
263 }
264
265 pub fn reindex_all(&mut self) -> Result<u32> {
266 let mut n = 0u32;
267 for sid in 0..self.n_shards {
268 if self.shard_has_data(sid) {
269 n += self.shard_mut(sid)?.reindex()?;
270 }
271 }
272 Ok(n)
273 }
274
275 pub fn drain_to(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> Result<u32> {
279 let src = src.as_ref();
280 let dst = dst.as_ref();
281 let man = src.join(MANIFEST);
282 if !man.is_file() {
283 return Err(Error::Message("drain: missing shards.json".into()));
284 }
285 let dest_was_new = !dst.exists();
286 let dest_man = dst.join(MANIFEST);
287 let mut copied_manifest = false;
288 let mut created = Vec::new();
289 let written = (|| -> Result<u32> {
290 std::fs::create_dir_all(dst)?;
291 if dest_man.is_file() {
292 let existing: ShardManifest =
293 serde_json::from_str(&std::fs::read_to_string(&dest_man)?)?;
294 let incoming: ShardManifest =
295 serde_json::from_str(&std::fs::read_to_string(&man)?)?;
296 if existing.n_shards != incoming.n_shards {
297 return Err(Error::Message(
298 "drain: dest shards.json n_shards does not match src".into(),
299 ));
300 }
301 }
302 let m: ShardManifest = serde_json::from_str(&std::fs::read_to_string(&man)?)?;
303 for i in 0..m.n_shards {
304 let name = format!("shard_{i:04}");
305 if src.join(&name).join("data.mdb").is_file()
306 && dst.join(&name).join("data.mdb").is_file()
307 {
308 return Err(Error::Message(format!(
309 "drain: dest {name} exists; refuse overwrite. Drain each node to a unique dest, then join-drained."
310 )));
311 }
312 }
313 copied_manifest = !dest_man.is_file();
314 if copied_manifest {
315 #[cfg(test)]
316 if FAIL_DEST_MAN_COPY.with(|f| f.replace(false)) {
317 return Err(Error::Message("drain: dest_man copy".into()));
318 }
319 std::fs::copy(&man, &dest_man)?;
320 }
321 let mut n = 0u32;
322 for i in 0..m.n_shards {
323 let name = format!("shard_{i:04}");
324 let from = src.join(&name);
325 if !from.join("data.mdb").is_file() {
326 continue;
327 }
328 let to = dst.join(&name);
329 created.push(to.clone());
330 let ro = ConCorpus::open_readonly(&from)?;
331 ro.snapshot_to(&to)?;
332 ro.close();
333 n += 1;
334 }
335 Ok(n)
336 })();
337 if written.is_err() {
338 if dest_was_new {
339 let _ = std::fs::remove_dir_all(dst);
340 } else {
341 for p in &created {
342 let _ = std::fs::remove_dir_all(p);
343 }
344 if copied_manifest {
345 let _ = std::fs::remove_file(&dest_man);
346 }
347 }
348 }
349 written
350 }
351}
352
353#[cfg(test)]
354mod tests {
355 use super::*;
356 use std::sync::Arc;
357 use std::thread;
358
359 fn fixture(name: &str) -> PathBuf {
360 PathBuf::from(env!("CARGO_MANIFEST_DIR"))
361 .join("resources/test")
362 .join(name)
363 }
364
365 #[test]
366 fn parallel_writers_different_shards() {
367 let dir = tempfile::tempdir().unwrap();
368 let root = dir.path().join("hpc");
369 let n_shards = 8u32;
371 ShardedConCorpus::open(&root, n_shards).unwrap();
372 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
373 let root = Arc::new(root);
374 let mut joins = Vec::new();
375 for sid in 0..n_shards {
376 let root = Arc::clone(&root);
377 let text = text.clone();
378 joins.push(thread::spawn(move || {
379 let db = ShardedConCorpus::open_shard(root.as_path(), sid).unwrap();
381 let traj = u64::from(sid); db.append_trajectory_str(traj, &text, format!("shard{sid}"))
383 .unwrap()
384 }));
385 }
386 let mut ns = Vec::new();
387 for j in joins {
388 ns.push(j.join().unwrap());
389 }
390 assert!(ns.iter().all(|&n| n >= 1));
391 let mut fan = ShardedConCorpus::open(root.as_path(), n_shards).unwrap();
392 let keys = fan.select(&Select::new().require_symbol("Cu")).unwrap();
393 drop(fan);
394 let drained = dir.path().join("pfs");
395 let ncopy = ShardedConCorpus::drain_to(root.as_path(), &drained).unwrap();
396 assert_eq!(ncopy, n_shards);
397 assert!(drained.join("shards.json").is_file());
398 assert!(drained.join("shard_0000").join("data.mdb").is_file());
399 assert!(!drained.join("shard_0000").join("lock.mdb").is_file());
400 let dest_sz = std::fs::metadata(drained.join("shard_0000").join("data.mdb"))
401 .unwrap()
402 .len();
403 assert!(
404 dest_sz < 64 * 1024 * 1024,
405 "compact snapshot must not materialize the 2 GiB map, got {dest_sz}"
406 );
407 assert!(ShardedConCorpus::drain_to(root.as_path(), &drained).is_err());
408 assert!(drained.join("shards.json").is_file());
409 assert_eq!(keys.len(), 8);
410 let joined = dir.path().join("joined");
411 let mut drained_root = ShardedConCorpus::open(&drained, n_shards).unwrap();
412 let njoin = drained_root.join_to_single_env(&joined).unwrap();
413 assert!(njoin >= n_shards);
414 let single = ConCorpus::open(&joined).unwrap();
415 let jk = single.select(&Select::new().require_symbol("Cu")).unwrap();
416 assert_eq!(jk.len(), keys.len());
417 }
418
419 #[test]
420 fn open_shard_for_traj_writes_manifest() {
421 let dir = tempfile::tempdir().unwrap();
422 let root = dir.path().join("fresh");
423 let (sid, db) = ShardedConCorpus::open_shard_for_traj(&root, 0).unwrap();
424 assert_eq!(sid, 0);
425 drop(db);
426 assert!(root.join("shards.json").is_file());
427 }
428
429 #[test]
430 fn join_drained_duplicate_traj_does_not_create_dest() {
431 let dir = tempfile::tempdir().unwrap();
432 let a = dir.path().join("a");
433 let b = dir.path().join("b");
434 ShardedConCorpus::open(&a, 1).unwrap();
435 ShardedConCorpus::open(&b, 1).unwrap();
436 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
437 ShardedConCorpus::open_shard(&a, 0)
438 .unwrap()
439 .append_trajectory_str(7, &text, "a")
440 .unwrap();
441 ShardedConCorpus::open_shard(&b, 0)
442 .unwrap()
443 .append_trajectory_str(7, &text, "b")
444 .unwrap();
445 let dest = dir.path().join("out");
446 let err = join_drained_roots(&[a, b], &dest).unwrap_err();
447 assert!(err.to_string().contains("traj_id"), "{err}");
448 assert!(!dest.exists());
449 }
450
451 #[test]
452 fn join_drained_roots_missing_source_errors() {
453 let dir = tempfile::tempdir().unwrap();
454 let missing = dir.path().join("nope");
455 let dest = dir.path().join("out");
456 let err = join_drained_roots(&[missing], &dest).unwrap_err();
457 assert!(err.to_string().contains("missing shards.json"), "{err}");
458 assert!(!dest.exists());
459 }
460
461 #[test]
462 fn drain_refuse_does_not_write_manifest() {
463 let dir = tempfile::tempdir().unwrap();
464 let root = dir.path().join("hpc");
465 ShardedConCorpus::open(&root, 2).unwrap();
466 let db = ShardedConCorpus::open_shard(&root, 0).unwrap();
467 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
468 db.append_trajectory_str(0, &text, "s0").unwrap();
469 drop(db);
470 let dest = dir.path().join("pfs");
471 std::fs::create_dir_all(dest.join("shard_0000")).unwrap();
472 std::fs::write(dest.join("shard_0000").join("data.mdb"), b"x").unwrap();
473 let err = ShardedConCorpus::drain_to(&root, &dest).unwrap_err();
474 assert!(err.to_string().contains("refuse overwrite"), "{err}");
475 assert!(
476 !dest.join("shards.json").is_file(),
477 "refuse must not leave dest shards.json"
478 );
479 }
480
481 #[test]
482 fn traj_routing_stable() {
483 assert_eq!(ShardedConCorpus::shard_for_traj(0, 64), 0);
484 assert_eq!(ShardedConCorpus::shard_for_traj(65, 64), 1);
485 }
486
487 #[test]
491 fn multi_shard_writers_select_matches_single_env_baseline() {
492 let dir = tempfile::tempdir().unwrap();
493 let root = dir.path().join("hpc_scale");
494 let baseline = dir.path().join("single_env");
495 let n_shards = 4u32;
496 ShardedConCorpus::open(&root, n_shards).unwrap();
497 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
498 let root_a = Arc::new(root.clone());
499 let mut joins = Vec::new();
500 for sid in 0..n_shards {
501 let root = Arc::clone(&root_a);
502 let text = text.clone();
503 joins.push(thread::spawn(move || {
504 let db = ShardedConCorpus::open_shard(root.as_path(), sid).unwrap();
505 let traj = u64::from(sid);
506 db.append_trajectory_str(traj, &text, format!("s{sid}"))
507 .unwrap()
508 }));
509 }
510 let mut frames_per_traj = Vec::new();
511 for j in joins {
512 frames_per_traj.push(j.join().unwrap());
513 }
514 assert!(frames_per_traj.iter().all(|&n| n >= 1));
515
516 let single = ConCorpus::open(&baseline).unwrap();
518 for sid in 0..n_shards {
519 let traj = u64::from(sid);
520 let n = single
521 .append_trajectory_str(traj, &text, format!("s{sid}"))
522 .unwrap();
523 assert_eq!(n, frames_per_traj[sid as usize]);
524 }
525
526 let mut fan = ShardedConCorpus::open(&root, n_shards).unwrap();
527 let sharded_keys = fan.select(&Select::new().require_symbol("Cu")).unwrap();
528 let base_keys = single.select(&Select::new().require_symbol("Cu")).unwrap();
529 assert_eq!(sharded_keys.len(), base_keys.len());
530 let mut sk: Vec<_> = sharded_keys
531 .iter()
532 .map(|k| (k.traj_id, k.frame_idx))
533 .collect();
534 let mut bk: Vec<_> = base_keys.iter().map(|k| (k.traj_id, k.frame_idx)).collect();
535 sk.sort_unstable();
536 bk.sort_unstable();
537 assert_eq!(sk, bk, "fan-out select must match single-env membership");
538
539 for (tid, fidx) in &sk {
541 let key = crate::keys::FrameKey {
542 traj_id: *tid,
543 frame_idx: *fidx,
544 };
545 let a = fan.get_frame_text(key).unwrap();
546 let b = single.get_frame_text(key).unwrap();
547 assert_eq!(a, b);
548 }
549 }
550}
551
552#[derive(Clone, Copy, Debug, PartialEq, Eq)]
554pub enum CorpusExportKind {
555 ShardedLmdb,
557 SingleEnvLmdb,
559 ExtXyz,
561}
562
563impl CorpusExportKind {
564 pub fn as_str(self) -> &'static str {
565 match self {
566 Self::ShardedLmdb => "sharded-lmdb",
567 Self::SingleEnvLmdb => "single-env-lmdb",
568 Self::ExtXyz => "extxyz",
569 }
570 }
571}
572
573impl ShardedConCorpus {
574 pub fn join_to_single_env(&mut self, dst: impl AsRef<Path>) -> Result<u32> {
579 let dst = dst.as_ref();
580 if dst.exists() {
581 return Err(Error::Message(format!(
582 "join dest exists: {}",
583 dst.display()
584 )));
585 }
586 self.release_shards();
587 let mut preview = std::collections::BTreeMap::new();
588 for sid in 0..self.n_shards {
589 if !self.shard_has_data(sid) {
590 continue;
591 }
592 let c = ConCorpus::open_readonly(self.shard_path(sid))?;
593 for fk in c.list_frame_keys()? {
594 preview_traj(&mut preview, fk.traj_id, u64::from(sid))?;
595 }
596 c.close();
597 }
598 let n = (|| -> Result<u32> {
599 let out = ConCorpus::open(dst)?;
600 let mut seen_traj = std::collections::BTreeSet::new();
601 let n = append_sharded_into(self, &out, &mut seen_traj)?;
602 out.close();
603 Ok(n)
604 })();
605 rollback_new_dest(dst, n)
606 }
607
608 pub fn export_extxyz(
610 &mut self,
611 sel: &Select,
612 out: impl AsRef<Path>,
613 energy_key: &str,
614 ) -> Result<u32> {
615 let joined = std::env::temp_dir().join(format!(
616 "readcon_db_join_{}_{}",
617 std::process::id(),
618 std::time::SystemTime::now()
619 .duration_since(std::time::UNIX_EPOCH)
620 .map(|d| d.as_nanos())
621 .unwrap_or(0)
622 ));
623 if joined.exists() {
624 return Err(Error::Message(
625 "export_extxyz: temp join dest exists".into(),
626 ));
627 }
628 struct RemoveOnDrop(PathBuf);
629 impl Drop for RemoveOnDrop {
630 fn drop(&mut self) {
631 let _ = std::fs::remove_dir_all(&self.0);
632 }
633 }
634 let _guard = RemoveOnDrop(joined.clone());
635 self.join_to_single_env(&joined)?;
636 let db = ConCorpus::open_readonly(&joined)?;
637 let keys = db.select(sel)?;
638 let n = db.export_extxyz(&keys, out, energy_key)?;
639 db.close();
640 Ok(n as u32)
641 }
642
643 pub fn split_single_to_sharded(
646 single: &ConCorpus,
647 dst_root: impl AsRef<Path>,
648 n_shards: u32,
649 ) -> Result<u32> {
650 if n_shards == 0 {
651 return Err(Error::Message("n_shards must be >= 1".into()));
652 }
653 let dst_root = dst_root.as_ref();
654 if dst_root.exists() {
655 return Err(Error::Message(format!(
656 "split dest exists: {}",
657 dst_root.display()
658 )));
659 }
660 let n = (|| -> Result<u32> {
661 let mut sharded = ShardedConCorpus::open(dst_root, n_shards)?;
662 let keys = single.list_frame_keys()?;
663 let mut by_traj: std::collections::BTreeMap<u64, Vec<FrameKey>> =
664 std::collections::BTreeMap::new();
665 for fk in keys {
666 by_traj.entry(fk.traj_id).or_default().push(fk);
667 }
668 let mut n = 0u32;
669 for (tid, mut fks) in by_traj {
670 fks.sort();
671 let mut concat = String::new();
672 for fk in &fks {
673 concat.push_str(&single.get_frame_text(*fk)?);
674 }
675 let nf = sharded.append_trajectory_str(tid, &concat, "split-from-single")?;
676 n += nf;
677 }
678 Ok(n)
679 })();
680 rollback_new_dest(dst_root, n)
681 }
682}
683
684fn append_sharded_into(
685 sh: &mut ShardedConCorpus,
686 out: &ConCorpus,
687 seen_traj: &mut std::collections::BTreeSet<u64>,
688) -> Result<u32> {
689 let mut n = 0u32;
690 for sid in 0..sh.n_shards {
691 if !sh.shard_has_data(sid) {
692 continue;
693 }
694 let shard = ConCorpus::open_readonly(sh.shard_path(sid))?;
695 let keys = shard.list_frame_keys()?;
696 let mut by_traj: std::collections::BTreeMap<u64, Vec<FrameKey>> =
697 std::collections::BTreeMap::new();
698 for fk in keys {
699 by_traj.entry(fk.traj_id).or_default().push(fk);
700 }
701 for (tid, mut fks) in by_traj {
702 if !seen_traj.insert(tid) {
703 return Err(Error::Message(format!(
704 "traj_id {tid} appears in multiple shards or join sources"
705 )));
706 }
707 fks.sort();
708 let mut concat = String::new();
709 for fk in &fks {
710 concat.push_str(&shard.get_frame_text(*fk)?);
711 }
712 n += out.append_trajectory_str(tid, &concat, format!("join-from-shard-{sid}"))?;
713 }
714 shard.close();
715 }
716 Ok(n)
717}
718
719pub fn join_drained_roots(sources: &[PathBuf], dst: impl AsRef<Path>) -> Result<u32> {
722 if sources.is_empty() {
723 return Err(Error::Message("join-drained: no sources".into()));
724 }
725 let dst = dst.as_ref();
726 if dst.exists() {
727 return Err(Error::Message(format!(
728 "join-drained dest exists: {}",
729 dst.display()
730 )));
731 }
732 for src in sources {
733 if !src.join(MANIFEST).is_file() {
734 return Err(Error::Message(format!(
735 "join-drained: missing shards.json: {}",
736 src.display()
737 )));
738 }
739 }
740 {
741 let mut preview = std::collections::BTreeMap::new();
742 for (si, src) in sources.iter().enumerate() {
743 let sh = ShardedConCorpus::open_existing(src)?;
744 for sid in 0..sh.n_shards {
745 if !sh.shard_has_data(sid) {
746 continue;
747 }
748 let shard = ConCorpus::open_readonly(sh.shard_path(sid))?;
749 let owner = ((si as u64) << 32) | u64::from(sid);
750 for fk in shard.list_frame_keys()? {
751 preview_traj(&mut preview, fk.traj_id, owner)?;
752 }
753 shard.close();
754 }
755 }
756 }
757 let n = (|| -> Result<u32> {
758 let out = ConCorpus::open(dst)?;
759 let mut n = 0u32;
760 let mut seen = std::collections::BTreeSet::new();
761 for src in sources {
762 let mut sh = ShardedConCorpus::open_existing(src)?;
763 n += append_sharded_into(&mut sh, &out, &mut seen)?;
764 }
765 out.close();
766 Ok(n)
767 })();
768 rollback_new_dest(dst, n)
769}
770
771pub fn join_corpus_dirs(sources: &[PathBuf], dst: impl AsRef<Path>) -> Result<u32> {
773 let dst = dst.as_ref();
774 if dst.exists() {
775 return Err(Error::Message(format!(
776 "join dest exists: {}",
777 dst.display()
778 )));
779 }
780 for src in sources {
781 if !src.join("data.mdb").is_file() {
782 return Err(Error::Message(format!(
783 "join: missing data.mdb: {}",
784 src.display()
785 )));
786 }
787 }
788 {
789 let mut preview = std::collections::BTreeMap::new();
790 for (si, src) in sources.iter().enumerate() {
791 let c = ConCorpus::open_readonly(src)?;
792 for fk in c.list_frame_keys()? {
793 preview_traj(&mut preview, fk.traj_id, si as u64)?;
794 }
795 c.close();
796 }
797 }
798 let n = (|| -> Result<u32> {
799 let out = ConCorpus::open(dst)?;
800 let mut n = 0u32;
801 let mut seen = std::collections::BTreeSet::new();
802 for src in sources {
803 let c = ConCorpus::open_readonly(src)?;
804 let keys = c.list_frame_keys()?;
805 let mut by_traj: std::collections::BTreeMap<u64, Vec<FrameKey>> =
806 std::collections::BTreeMap::new();
807 for fk in keys {
808 by_traj.entry(fk.traj_id).or_default().push(fk);
809 }
810 for (tid, mut fks) in by_traj {
811 if !seen.insert(tid) {
812 return Err(Error::Message(format!(
813 "duplicate traj_id {tid} across join sources"
814 )));
815 }
816 fks.sort();
817 let mut concat = String::new();
818 for fk in &fks {
819 concat.push_str(&c.get_frame_text(*fk)?);
820 }
821 n += out.append_trajectory_str(tid, &concat, src.display().to_string())?;
822 }
823 c.close();
824 }
825 out.close();
826 Ok(n)
827 })();
828 rollback_new_dest(dst, n)
829}
830
831fn rollback_new_dest<T>(dst: &Path, r: Result<T>) -> Result<T> {
832 if r.is_err() {
833 let _ = std::fs::remove_dir_all(dst);
834 }
835 r
836}
837
838fn preview_traj(
841 preview: &mut std::collections::BTreeMap<u64, u64>,
842 traj_id: u64,
843 owner: u64,
844) -> Result<()> {
845 match preview.entry(traj_id) {
846 std::collections::btree_map::Entry::Occupied(e) if *e.get() != owner => {
847 Err(Error::Message(format!(
848 "traj_id {traj_id} appears in multiple shards or join sources"
849 )))
850 }
851 std::collections::btree_map::Entry::Occupied(_) => Ok(()),
852 std::collections::btree_map::Entry::Vacant(e) => {
853 e.insert(owner);
854 Ok(())
855 }
856 }
857}
858
859pub fn open_single_env_for_export(src: impl AsRef<Path>) -> Result<ConCorpus> {
862 let src = src.as_ref();
863 if src.join(MANIFEST).is_file() {
864 return Err(Error::Message(
865 "compact-export-extxyz: sharded root needs --sharded".into(),
866 ));
867 }
868 ConCorpus::open_readonly(src)
869}
870
871#[cfg(test)]
872mod compaction_tests {
873 use super::*;
874 use crate::keys::FrameKey;
875 use crate::select::Select;
876
877 fn fixture(name: &str) -> PathBuf {
878 PathBuf::from(env!("CARGO_MANIFEST_DIR"))
879 .join("resources/test")
880 .join(name)
881 }
882
883 #[test]
884 fn join_split_reversible_membership() {
885 let dir = tempfile::tempdir().unwrap();
886 let sharded_root = dir.path().join("sharded");
887 let con_text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
888 {
889 let mut s = ShardedConCorpus::open(&sharded_root, 4).unwrap();
890 for tid in [0u64, 1, 2, 3] {
891 s.append_trajectory_str(tid, &con_text, "t").unwrap();
892 }
893 }
894 let mut s = ShardedConCorpus::open(&sharded_root, 4).unwrap();
895 let before = s.select(&Select::new()).unwrap();
896 assert_eq!(before.len(), 4);
897
898 let joined = dir.path().join("joined");
899 let n = s.join_to_single_env(&joined).unwrap();
900 assert_eq!(n, 4);
901 assert!(s.join_to_single_env(&joined).is_err());
902 let joined_c = ConCorpus::open(&joined).unwrap();
903 let mid = joined_c.select(&Select::new()).unwrap();
904 assert_eq!(mid.len(), 4);
905
906 let split_root = dir.path().join("split_again");
907 let n2 = ShardedConCorpus::split_single_to_sharded(&joined_c, &split_root, 4).unwrap();
908 assert_eq!(n2, 4);
909 let mut s2 = ShardedConCorpus::open(&split_root, 4).unwrap();
910 let after = s2.select(&Select::new()).unwrap();
911 assert_eq!(after.len(), before.len());
912 let mut bt: Vec<_> = before.iter().map(|k| k.traj_id).collect();
914 let mut at: Vec<_> = after.iter().map(|k| k.traj_id).collect();
915 bt.sort();
916 at.sort();
917 assert_eq!(bt, at);
918 }
919
920 #[test]
921 fn join_drained_roots_keeps_disjoint_trajs_same_shard() {
922 let dir = tempfile::tempdir().unwrap();
923 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
924 let n_shards = 2u32;
925 let node_a = dir.path().join("node_a");
926 let node_b = dir.path().join("node_b");
927 ShardedConCorpus::open(&node_a, n_shards).unwrap();
928 ShardedConCorpus::open(&node_b, n_shards).unwrap();
929 ShardedConCorpus::open_shard(&node_a, 0)
931 .unwrap()
932 .append_trajectory_str(0, &text, "a")
933 .unwrap();
934 ShardedConCorpus::open_shard(&node_b, 0)
935 .unwrap()
936 .append_trajectory_str(2, &text, "b")
937 .unwrap();
938 let dest_a = dir.path().join("dest_a");
939 let dest_b = dir.path().join("dest_b");
940 assert_eq!(ShardedConCorpus::drain_to(&node_a, &dest_a).unwrap(), 1);
941 assert_eq!(ShardedConCorpus::drain_to(&node_b, &dest_b).unwrap(), 1);
942 assert!(ShardedConCorpus::drain_to(&node_a, &dest_a).is_err());
943 let joined = dir.path().join("joined");
944 let n = join_drained_roots(&[dest_a.clone(), dest_b.clone()], &joined).unwrap();
945 assert!(n >= 2);
946 let db = ConCorpus::open(&joined).unwrap();
947 let keys = db.select(&Select::new()).unwrap();
948 let mut tids: Vec<u64> = keys.iter().map(|k| k.traj_id).collect();
949 tids.sort_unstable();
950 tids.dedup();
951 assert_eq!(tids, vec![0, 2]);
952 assert!(join_drained_roots(&[dest_a, dest_b], &joined).is_err());
953 }
954
955 #[test]
956 fn join_drained_roots_refuses_existing_dest() {
957 let dir = tempfile::tempdir().unwrap();
958 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
959 let node = dir.path().join("node");
960 ShardedConCorpus::open(&node, 2).unwrap();
961 ShardedConCorpus::open_shard(&node, 0)
962 .unwrap()
963 .append_trajectory_str(0, &text, "a")
964 .unwrap();
965 let dest = dir.path().join("dest");
966 assert_eq!(ShardedConCorpus::drain_to(&node, &dest).unwrap(), 1);
967 let joined = dir.path().join("joined");
968 join_drained_roots(&[dest.clone()], &joined).unwrap();
969 let err = join_drained_roots(&[dest], &joined).unwrap_err();
970 assert!(
971 err.to_string().contains("join-drained dest exists"),
972 "{err}"
973 );
974 }
975
976 #[test]
977 fn export_kinds_documented() {
978 assert_eq!(CorpusExportKind::ShardedLmdb.as_str(), "sharded-lmdb");
979 assert_eq!(CorpusExportKind::SingleEnvLmdb.as_str(), "single-env-lmdb");
980 assert_eq!(CorpusExportKind::ExtXyz.as_str(), "extxyz");
981 }
982
983 #[test]
984 fn select_skips_empty_shard_dir_without_data() {
985 let dir = tempfile::tempdir().unwrap();
986 let root = dir.path().join("sharded");
987 let mut s = ShardedConCorpus::open(&root, 4).unwrap();
988 std::fs::create_dir_all(root.join("shard_0001")).unwrap();
989 let keys = s.select(&Select::new()).unwrap();
990 assert!(keys.is_empty());
991 assert!(!root.join("shard_0001").join("data.mdb").exists());
992 let err = s
993 .get_frame_text(FrameKey {
994 traj_id: 1,
995 frame_idx: 0,
996 })
997 .unwrap_err();
998 assert!(
999 err.to_string().contains("is not a corpus directory"),
1000 "{err}"
1001 );
1002 assert!(!root.join("shard_0001").join("data.mdb").exists());
1003 }
1004
1005 #[test]
1006 fn select_skips_missing_shard_dirs() {
1007 let dir = tempfile::tempdir().unwrap();
1008 let root = dir.path().join("sharded");
1009 let mut s = ShardedConCorpus::open(&root, 8).unwrap();
1010 s.append_trajectory_path(0, fixture("tiny_cuh2.con"))
1011 .unwrap();
1012 drop(s);
1013 let mut s = ShardedConCorpus::open_existing(&root).unwrap();
1014 let keys = s.select(&Select::new()).unwrap();
1015 assert_eq!(keys.len(), 1);
1016 let present: Vec<u32> = (0..8)
1017 .filter(|&i| root.join(format!("shard_{i:04}")).is_dir())
1018 .collect();
1019 assert_eq!(present, vec![0]);
1020 }
1021
1022 #[test]
1023 fn get_frame_text_missing_shard_does_not_mint() {
1024 let dir = tempfile::tempdir().unwrap();
1025 let root = dir.path().join("sharded");
1026 let mut s = ShardedConCorpus::open(&root, 4).unwrap();
1027 let err = s
1028 .get_frame_text(FrameKey {
1029 traj_id: 1,
1030 frame_idx: 0,
1031 })
1032 .unwrap_err();
1033 assert!(
1034 err.to_string().contains("is not a corpus directory"),
1035 "{err}"
1036 );
1037 assert!(!root.join("shard_0001").exists());
1038 }
1039
1040 #[test]
1041 fn join_drained_roots_keeps_multi_frame_traj() {
1042 let dir = tempfile::tempdir().unwrap();
1043 let text = std::fs::read_to_string(fixture("tiny_multi_cuh2.con")).unwrap();
1044 let node = dir.path().join("node");
1045 ShardedConCorpus::open(&node, 2).unwrap();
1046 ShardedConCorpus::open_shard(&node, 0)
1047 .unwrap()
1048 .append_trajectory_str(0, &text, "a")
1049 .unwrap();
1050 let dest = dir.path().join("dest");
1051 assert_eq!(ShardedConCorpus::drain_to(&node, &dest).unwrap(), 1);
1052 let joined = dir.path().join("joined");
1053 let n = join_drained_roots(&[dest], &joined).unwrap();
1054 assert!(n >= 2, "joined frames={n}");
1055 let db = ConCorpus::open(&joined).unwrap();
1056 let p0 = db
1057 .get_positions(crate::keys::FrameKey {
1058 traj_id: 0,
1059 frame_idx: 0,
1060 })
1061 .unwrap();
1062 let p1 = db
1063 .get_positions(crate::keys::FrameKey {
1064 traj_id: 0,
1065 frame_idx: 1,
1066 })
1067 .unwrap();
1068 assert!((p0[0][0] - 0.6394).abs() < 1e-4);
1069 assert!((p1[2][0] - 8.8549).abs() < 1e-4);
1070 }
1071
1072 #[test]
1073 fn join_corpus_dirs_keeps_multi_frame_traj() {
1074 let dir = tempfile::tempdir().unwrap();
1075 let a = dir.path().join("a");
1076 ConCorpus::open(&a)
1077 .unwrap()
1078 .append_trajectory_path(1, fixture("tiny_multi_cuh2.con"))
1079 .unwrap();
1080 let dest = dir.path().join("out");
1081 let n = join_corpus_dirs(&[a], &dest).unwrap();
1082 assert!(n >= 2, "joined frames={n}");
1083 let db = ConCorpus::open(&dest).unwrap();
1084 let p0 = db
1085 .get_positions(crate::keys::FrameKey {
1086 traj_id: 1,
1087 frame_idx: 0,
1088 })
1089 .unwrap();
1090 let p1 = db
1091 .get_positions(crate::keys::FrameKey {
1092 traj_id: 1,
1093 frame_idx: 1,
1094 })
1095 .unwrap();
1096 assert!((p0[0][0] - 0.6394).abs() < 1e-4);
1097 assert!((p1[2][0] - 8.8549).abs() < 1e-4);
1098 }
1099
1100 #[test]
1101 fn drain_to_n_shards_mismatch_refuses() {
1102 let dir = tempfile::tempdir().unwrap();
1103 let src = dir.path().join("src");
1104 let dest = dir.path().join("dest");
1105 ShardedConCorpus::open(&src, 2).unwrap();
1106 ShardedConCorpus::open(&dest, 4).unwrap();
1107 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
1108 ShardedConCorpus::open_shard(&src, 0)
1109 .unwrap()
1110 .append_trajectory_str(0, &text, "a")
1111 .unwrap();
1112 let err = ShardedConCorpus::drain_to(&src, &dest).unwrap_err();
1113 assert!(err.to_string().contains("n_shards"), "{err}");
1114 let man: ShardManifest =
1115 serde_json::from_str(&std::fs::read_to_string(dest.join("shards.json")).unwrap())
1116 .unwrap();
1117 assert_eq!(man.n_shards, 4);
1118 assert!(!dest.join("shard_0000").join("data.mdb").is_file());
1119 }
1120
1121 #[test]
1122 fn drain_to_new_dest_rollback_on_bad_src() {
1123 let dir = tempfile::tempdir().unwrap();
1124 let src = dir.path().join("src");
1125 ShardedConCorpus::open(&src, 2).unwrap();
1126 let shard0 = src.join("shard_0000");
1127 std::fs::create_dir_all(&shard0).unwrap();
1128 std::fs::write(shard0.join("data.mdb"), b"not-lmdb").unwrap();
1129 let dest = dir.path().join("dest");
1130 assert!(!dest.exists());
1131 assert!(ShardedConCorpus::drain_to(&src, &dest).is_err());
1132 assert!(!dest.exists(), "dest_was_new rollback must remove dest");
1133 }
1134
1135 #[test]
1136 fn drain_to_new_dest_rollback_on_empty_src_manifest() {
1137 let dir = tempfile::tempdir().unwrap();
1138 let src = dir.path().join("src");
1139 std::fs::create_dir_all(&src).unwrap();
1140 std::fs::write(src.join("shards.json"), b"").unwrap();
1141 let dest = dir.path().join("dest");
1142 assert!(!dest.exists());
1143 assert!(ShardedConCorpus::drain_to(&src, &dest).is_err());
1144 assert!(
1145 !dest.exists(),
1146 "dest leftover empty dest_man parse must remove dest"
1147 );
1148 }
1149
1150 #[test]
1151 fn drain_to_new_dest_rollback_on_dest_man_copy() {
1152 let dir = tempfile::tempdir().unwrap();
1153 let src = dir.path().join("src");
1154 ShardedConCorpus::open(&src, 2).unwrap();
1155 let dest = dir.path().join("dest");
1156 std::fs::create_dir_all(dest.join("shards.json")).unwrap();
1157 assert!(ShardedConCorpus::drain_to(&src, &dest).is_err());
1158 assert!(
1159 dest.exists(),
1160 "existing dest stays; dest_man copy failed because dest_man is a dir"
1161 );
1162 assert!(dest.join("shards.json").is_dir());
1163 assert!(!dest.join("shards.json").is_file());
1164 }
1165
1166 #[test]
1167 fn drain_to_new_dest_rollback_on_dest_man_copy_fail() {
1168 let dir = tempfile::tempdir().unwrap();
1169 let src = dir.path().join("src");
1170 ShardedConCorpus::open(&src, 2).unwrap();
1171 let dest = dir.path().join("dest");
1172 assert!(!dest.exists());
1173 FAIL_DEST_MAN_COPY.with(|f| f.set(true));
1174 assert!(ShardedConCorpus::drain_to(&src, &dest).is_err());
1175 assert!(
1176 !dest.exists(),
1177 "dest leftover dest_man copy dest_was_new must remove dest"
1178 );
1179 }
1180
1181 #[test]
1182 fn drain_to_failure_removes_shards_created_this_call() {
1183 let dir = tempfile::tempdir().unwrap();
1184 let src = dir.path().join("src");
1185 ShardedConCorpus::open(&src, 2).unwrap();
1186 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
1187 ShardedConCorpus::open_shard(&src, 0)
1188 .unwrap()
1189 .append_trajectory_str(0, &text, "a")
1190 .unwrap();
1191 ShardedConCorpus::open_shard(&src, 1)
1192 .unwrap()
1193 .append_trajectory_str(1, &text, "b")
1194 .unwrap();
1195 let dest = dir.path().join("dest");
1196 std::fs::create_dir_all(&dest).unwrap();
1197 std::fs::write(dest.join("shard_0001"), b"not-a-dir").unwrap();
1198 assert!(ShardedConCorpus::drain_to(&src, &dest).is_err());
1199 assert!(!dest.join("shard_0000").exists());
1200 }
1201
1202 #[test]
1203 fn rollback_new_dest_removes_dir() {
1204 let dir = tempfile::tempdir().unwrap();
1205 let dest = dir.path().join("out");
1206 std::fs::create_dir_all(&dest).unwrap();
1207 std::fs::write(dest.join("x"), b"y").unwrap();
1208 let err = rollback_new_dest::<u32>(&dest, Err(Error::Message("boom".into()))).unwrap_err();
1209 assert!(err.to_string().contains("boom"), "{err}");
1210 assert!(!dest.exists());
1211 }
1212
1213 #[test]
1214 fn join_to_single_env_keeps_multi_frame_traj() {
1215 let dir = tempfile::tempdir().unwrap();
1216 let root = dir.path().join("sharded");
1217 let mut s = ShardedConCorpus::open(&root, 2).unwrap();
1218 s.append_trajectory_path(0, fixture("tiny_multi_cuh2.con"))
1219 .unwrap();
1220 let dest = dir.path().join("out");
1221 let n = s.join_to_single_env(&dest).unwrap();
1222 assert!(n >= 2, "joined frames={n}");
1223 let db = ConCorpus::open(&dest).unwrap();
1224 let keys = db.select(&Select::new()).unwrap();
1225 assert!(keys.len() >= 2);
1226 assert!(keys.iter().all(|k| k.traj_id == 0));
1227 let p0 = db
1228 .get_positions(crate::keys::FrameKey {
1229 traj_id: 0,
1230 frame_idx: 0,
1231 })
1232 .unwrap();
1233 let p1 = db
1234 .get_positions(crate::keys::FrameKey {
1235 traj_id: 0,
1236 frame_idx: 1,
1237 })
1238 .unwrap();
1239 assert!((p0[0][0] - 0.6394).abs() < 1e-4);
1240 assert!((p1[2][0] - 8.8549).abs() < 1e-4);
1241 }
1242
1243 #[test]
1244 fn join_to_single_env_duplicate_traj_does_not_create_dest() {
1245 let dir = tempfile::tempdir().unwrap();
1246 let root = dir.path().join("sharded");
1247 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
1248 let mut s = ShardedConCorpus::open(&root, 2).unwrap();
1249 s.shard_mut(0)
1250 .unwrap()
1251 .append_trajectory_str(0, &text, "a")
1252 .unwrap();
1253 s.shard_mut(1)
1254 .unwrap()
1255 .append_trajectory_str(0, &text, "b")
1256 .unwrap();
1257 let dest = dir.path().join("out");
1258 let err = s.join_to_single_env(&dest).unwrap_err();
1259 assert!(err.to_string().contains("traj_id"), "{err}");
1260 assert!(!dest.exists());
1261 }
1262
1263 #[test]
1264 fn join_corpus_dirs_duplicate_does_not_create_dest() {
1265 let dir = tempfile::tempdir().unwrap();
1266 let text = std::fs::read_to_string(fixture("tiny_cuh2.con")).unwrap();
1267 let a = dir.path().join("a");
1268 let b = dir.path().join("b");
1269 ConCorpus::open(&a)
1270 .unwrap()
1271 .append_trajectory_str(1, &text, "a")
1272 .unwrap();
1273 ConCorpus::open(&b)
1274 .unwrap()
1275 .append_trajectory_str(1, &text, "b")
1276 .unwrap();
1277 let dest = dir.path().join("out");
1278 let err = join_corpus_dirs(&[a, b], &dest).unwrap_err();
1279 assert!(
1280 err.to_string().contains("traj_id") || err.to_string().contains("duplicate"),
1281 "{err}"
1282 );
1283 assert!(!dest.exists());
1284 }
1285
1286 #[test]
1287 fn join_corpus_dirs_missing_mdb_does_not_create_dest() {
1288 let dir = tempfile::tempdir().unwrap();
1289 let missing = dir.path().join("nope");
1290 std::fs::create_dir_all(&missing).unwrap();
1291 let dest = dir.path().join("out");
1292 let err = join_corpus_dirs(&[missing], &dest).unwrap_err();
1293 assert!(err.to_string().contains("missing data.mdb"), "{err}");
1294 assert!(!dest.exists());
1295 }
1296
1297 #[test]
1298 fn open_single_env_for_export_refuses_sharded_root() {
1299 let dir = tempfile::tempdir().unwrap();
1300 let root = dir.path().join("sharded");
1301 ShardedConCorpus::open(&root, 2).unwrap();
1302 assert!(!root.join("data.mdb").exists());
1303 match open_single_env_for_export(&root) {
1304 Ok(_) => panic!("expected sharded-root refuse"),
1305 Err(err) => assert!(err.to_string().contains("--sharded"), "{err}"),
1306 }
1307 assert!(!root.join("data.mdb").exists());
1308 }
1309
1310 #[test]
1311 fn export_extxyz_removes_temp_join() {
1312 let dir = tempfile::tempdir().unwrap();
1313 let root = dir.path().join("sharded");
1314 let mut s = ShardedConCorpus::open(&root, 2).unwrap();
1315 s.append_trajectory_path(0, fixture("tiny_cuh2.con"))
1316 .unwrap();
1317 let out = dir.path().join("out.xyz");
1318 let n = s.export_extxyz(&Select::new(), &out, "energy").unwrap();
1319 assert!(n >= 1);
1320 assert!(out.is_file());
1321 let prefix = format!("readcon_db_join_{}_", std::process::id());
1322 let leftovers: Vec<_> = std::env::temp_dir()
1323 .read_dir()
1324 .unwrap()
1325 .filter_map(|e| e.ok())
1326 .filter(|e| e.file_name().to_string_lossy().starts_with(&prefix))
1327 .collect();
1328 assert!(leftovers.is_empty(), "{leftovers:?}");
1329 }
1330
1331 #[test]
1332 fn open_existing_error_is_missing_shards_json() {
1333 let dir = tempfile::tempdir().unwrap();
1334 match ShardedConCorpus::open_existing(dir.path().join("nope")) {
1335 Ok(_) => panic!("expected missing shards.json"),
1336 Err(err) => {
1337 assert!(err.to_string().starts_with("missing shards.json:"), "{err}");
1338 assert!(!err.to_string().contains("join-drained"));
1339 }
1340 }
1341 }
1342}