1use crate::discontinuity::{current_open_file_count, open_file_limit};
9use crate::dna::{Base, ascii_base_bits, valid_ascii_base_bits};
10use crate::params::BuildParams;
11use crate::partition::WeakSuperKmer;
12use std::collections::BTreeMap;
13use std::fs::{self, File, OpenOptions};
14use std::io::{BufRead, BufReader, Read, Seek, SeekFrom, Write};
15use std::os::unix::fs::FileExt;
16use std::path::{Path, PathBuf};
17use std::sync::{
18 Arc, Mutex,
19 atomic::{AtomicU64, AtomicUsize, Ordering},
20};
21use std::time::{Duration, Instant};
22
23const MAGIC: &[u8; 8] = b"CF3WSK1\0";
24const RECORD_COUNT_OFFSET: u64 = 34;
25const HEADER_LEN: u64 = 42;
26const COMPRESSED_BLOCK_HEADER_LEN: usize = 12;
27const MAX_SOURCE_ID: u32 = 0x1f_ffff;
28const MAX_OPEN_BUCKET_WRITERS: usize = 512;
29const MAX_RECORD_BYTES: usize = 4 + 4 * 8;
31const MAX_PENDING_BUCKET_BYTES: usize = 1024 * 1024;
32const MAX_TOTAL_PENDING_BYTES: usize = 128 * 1024 * 1024;
37const ATLAS_GRAPH_COUNT: usize = 128;
38const SUBGRAPH_CHUNK_BYTES: usize = 64 * 1024;
39
40const DEFAULT_BUCKET_SEGMENT_BYTES: u64 = 256 * 1024;
55
56const RESERVED_NON_BUCKET_DESCRIPTORS: usize = 384;
60
61fn bucket_segment_bytes() -> u64 {
62 static BYTES: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
63 *BYTES.get_or_init(|| {
64 std::env::var("CF3_RS_BUCKET_SEGMENT_BYTES")
65 .ok()
66 .and_then(|value| value.parse::<u64>().ok())
67 .filter(|bytes| *bytes >= MAX_RECORD_BYTES as u64 && bytes % 4096 == 0)
68 .unwrap_or(DEFAULT_BUCKET_SEGMENT_BYTES)
69 })
70}
71
72#[derive(Debug)]
93pub struct BucketContainers {
94 files: Vec<BucketContainerFile>,
95 segment_bytes: u64,
96 punch_unsupported: std::sync::atomic::AtomicBool,
98}
99
100#[derive(Debug)]
101struct BucketContainerFile {
102 path: PathBuf,
103 file: File,
104 cursor: AtomicU64,
107}
108
109impl BucketContainers {
110 fn plan_container_count(atlas_count: usize) -> usize {
120 let budget = open_file_limit()
121 .saturating_sub(current_open_file_count())
122 .saturating_sub(RESERVED_NON_BUCKET_DESCRIPTORS)
126 / 2;
127 atlas_count.min(budget.max(1))
128 }
129
130 fn create(bucket_dir: &Path, container_count: usize) -> Result<Self, BucketError> {
131 let mut files = Vec::with_capacity(container_count);
132 for index in 0..container_count {
133 let path = bucket_dir.join(format!("{index:05}.wskc"));
134 let file = OpenOptions::new()
135 .create(true)
136 .truncate(true)
137 .read(true)
138 .write(true)
139 .open(&path)
140 .map_err(|source| BucketError::Io {
141 path: path.clone(),
142 source,
143 })?;
144 files.push(BucketContainerFile {
145 path,
146 file,
147 cursor: AtomicU64::new(0),
148 });
149 }
150 Ok(Self {
151 files,
152 segment_bytes: bucket_segment_bytes(),
153 punch_unsupported: std::sync::atomic::AtomicBool::new(false),
154 })
155 }
156
157 fn open(
159 bucket_dir: &Path,
160 container_count: usize,
161 segment_bytes: u64,
162 ) -> Result<Self, BucketError> {
163 let mut files = Vec::with_capacity(container_count);
164 for index in 0..container_count {
165 let path = bucket_dir.join(format!("{index:05}.wskc"));
166 let file = OpenOptions::new()
167 .read(true)
168 .write(true)
169 .open(&path)
170 .map_err(|source| BucketError::Io {
171 path: path.clone(),
172 source,
173 })?;
174 files.push(BucketContainerFile {
175 path,
176 file,
177 cursor: AtomicU64::new(0),
178 });
179 }
180 Ok(Self {
181 files,
182 segment_bytes,
183 punch_unsupported: std::sync::atomic::AtomicBool::new(false),
184 })
185 }
186
187 #[inline]
188 pub fn segment_bytes(&self) -> u64 {
189 self.segment_bytes
190 }
191
192 #[inline]
193 pub fn len(&self) -> usize {
194 self.files.len()
195 }
196
197 #[inline]
198 pub fn is_empty(&self) -> bool {
199 self.files.is_empty()
200 }
201
202 #[inline]
203 fn reserve_segment(&self, container: usize) -> u64 {
204 self.files[container]
205 .cursor
206 .fetch_add(self.segment_bytes, Ordering::Relaxed)
207 }
208
209 fn write_at(&self, container: usize, offset: u64, bytes: &[u8]) -> Result<(), BucketError> {
210 let file = &self.files[container];
211 file.file
212 .write_all_at(bytes, offset)
213 .map_err(|source| BucketError::Io {
214 path: file.path.clone(),
215 source,
216 })
217 }
218
219 fn read_at(&self, container: usize, offset: u64, buf: &mut [u8]) -> Result<(), BucketError> {
220 let file = &self.files[container];
221 file.file
222 .read_exact_at(buf, offset)
223 .map_err(|source| BucketError::Io {
224 path: file.path.clone(),
225 source,
226 })
227 }
228
229 pub fn release_segments(&self, container: usize, segments: &[u32]) {
248 if segments.is_empty() {
249 return;
250 }
251 let mut ordered = segments.to_vec();
252 ordered.sort_unstable();
253 let file = &self.files[container];
254 let mut start = u64::from(ordered[0]) * self.segment_bytes;
255 let mut end = start + self.segment_bytes;
256 for &index in &ordered[1..] {
257 let offset = u64::from(index) * self.segment_bytes;
258 if offset == end {
259 end += self.segment_bytes;
260 continue;
261 }
262 self.report_punch(punch_hole(&file.file, start, end - start));
263 start = offset;
264 end = offset + self.segment_bytes;
265 }
266 self.report_punch(punch_hole(&file.file, start, end - start));
267 }
268
269 fn report_punch(&self, punched: bool) {
278 if punched || self.punch_unsupported.swap(true, Ordering::Relaxed) {
279 return;
280 }
281 eprintln!(
282 "cuttlefish: this filesystem will not punch holes, so consumed \
283 bucket space is held until the build ends; peak disk will be higher"
284 );
285 }
286
287 pub fn paths(&self) -> impl Iterator<Item = &Path> {
288 self.files.iter().map(|file| file.path.as_path())
289 }
290}
291
292#[cfg(target_os = "linux")]
311fn punch_hole(file: &File, offset: u64, len: u64) -> bool {
312 use std::os::fd::AsRawFd;
313 unsafe {
316 libc::fallocate(
317 file.as_raw_fd(),
318 libc::FALLOC_FL_KEEP_SIZE | libc::FALLOC_FL_PUNCH_HOLE,
319 offset as libc::off_t,
320 len as libc::off_t,
321 ) == 0
322 }
323}
324
325#[cfg(target_vendor = "apple")]
326fn punch_hole(file: &File, offset: u64, len: u64) -> bool {
327 use std::os::fd::AsRawFd;
328 let punch = libc::fpunchhole_t {
329 fp_flags: 0,
330 reserved: 0,
331 fp_offset: offset as libc::off_t,
332 fp_length: len as libc::off_t,
333 };
334 unsafe { libc::fcntl(file.as_raw_fd(), libc::F_PUNCHHOLE, &punch) == 0 }
337}
338
339#[cfg(not(any(target_os = "linux", target_vendor = "apple")))]
340fn punch_hole(_file: &File, _offset: u64, _len: u64) -> bool {
341 false
342}
343
344#[derive(Debug)]
352struct SegmentChainReader {
353 containers: Arc<BucketContainers>,
354 container: usize,
355 segments: Vec<u64>,
356 len: u64,
357 pos: u64,
358}
359
360impl Read for SegmentChainReader {
361 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
362 if self.pos >= self.len || buf.is_empty() {
363 return Ok(0);
364 }
365 let segment_bytes = self.containers.segment_bytes;
366 let index = (self.pos / segment_bytes) as usize;
367 let within = self.pos % segment_bytes;
368 let room = (segment_bytes - within).min(self.len - self.pos);
369 let take = (buf.len() as u64).min(room) as usize;
370 let offset = self.segments[index] + within;
371 self.containers
372 .read_at(self.container, offset, &mut buf[..take])
373 .map_err(std::io::Error::other)?;
374 self.pos += take as u64;
375 Ok(take)
376 }
377}
378
379#[derive(Debug, Clone, PartialEq, Eq)]
380pub struct BucketEmitStats {
381 pub bucket_dir: PathBuf,
382 pub bucket_files: usize,
383 pub bytes_written: u64,
384}
385
386pub struct BucketEmitter {
387 bucket_dir: PathBuf,
388 k: u16,
389 minimizer_len: u16,
390 graph_count: usize,
391 colored: bool,
392 compress_buckets: bool,
393 label_words: usize,
394 files: BTreeMap<usize, BucketFileMeta>,
395 writers: BTreeMap<usize, BucketFile>,
396 pending: Vec<PendingBucket>,
397 pending_bytes: usize,
398 scratch: CompressionScratch,
399}
400
401pub struct SharedBucketSink {
402 bucket_dir: PathBuf,
403 containers: BucketContainers,
404 k: u16,
405 minimizer_len: u16,
406 graph_count: usize,
407 colored: bool,
408 compress_buckets: bool,
409 label_words: usize,
410 workers: usize,
411 atlases: Vec<Mutex<SharedBucketAtlas>>,
412 flush_calls: AtomicU64,
413 flush_nanos: AtomicU64,
414}
415
416pub struct SharedBucketEmitter {
417 sink: Arc<SharedBucketSink>,
418 pending: Vec<PendingBucket>,
419 uncolored_pending: Vec<PendingColoredAtlas>,
420 colored_pending: Vec<PendingColoredAtlas>,
421 pending_bytes: usize,
422 deferred_uncolored: bool,
423}
424
425#[derive(Debug, Clone, PartialEq, Eq)]
426struct BucketFileMeta {
427 path: PathBuf,
428 records: u64,
429 bytes_written: u64,
430}
431
432#[derive(Debug, Clone, PartialEq, Eq, Default)]
433struct PendingBucket {
434 records: u64,
435 bytes: Vec<u8>,
436}
437
438#[derive(Debug, Clone, PartialEq, Eq, Default)]
439struct PendingColoredAtlas {
440 graph_ids: Vec<u16>,
441 bytes: Vec<u8>,
442}
443
444#[derive(Default)]
445struct SharedBucketFileMeta {
446 buffer: Vec<u8>,
447 buffer_records: u64,
448 total_records: u64,
449 written_records: u64,
450 bytes_written: u64,
451 segments: Vec<u32>,
453 segment_used: u64,
457}
458
459struct SharedBucketAtlas {
460 first_graph_id: usize,
461 files: Vec<SharedBucketFileMeta>,
462 buffered_bytes: usize,
463 scratch: CompressionScratch,
466}
467
468#[derive(Default)]
469struct SharedBucketFlushStats {
470 calls: u64,
471}
472
473#[derive(Debug, Clone, PartialEq, Eq)]
480pub enum BucketLocation {
481 File(PathBuf),
483 Container {
487 container: usize,
488 segments: Vec<u32>,
490 bytes: u64,
493 },
494}
495
496#[derive(Debug, Clone, PartialEq, Eq)]
497pub struct BucketManifestEntry {
499 pub graph_id: usize,
500 pub records: u64,
501 pub location: BucketLocation,
502}
503
504impl BucketManifestEntry {
505 pub fn stored_bytes(&self) -> Result<u64, BucketError> {
510 match &self.location {
511 BucketLocation::File(path) => {
512 fs::metadata(path)
513 .map(|meta| meta.len())
514 .map_err(|source| BucketError::Io {
515 path: path.clone(),
516 source,
517 })
518 }
519 BucketLocation::Container { bytes, .. } => Ok(*bytes),
520 }
521 }
522
523 pub fn file_path(&self) -> Option<&Path> {
525 match &self.location {
526 BucketLocation::File(path) => Some(path.as_path()),
527 BucketLocation::Container { .. } => None,
528 }
529 }
530}
531
532#[derive(Debug, Clone, PartialEq, Eq)]
533pub struct BucketHeader {
535 pub k: u16,
536 pub minimizer_len: u16,
537 pub graph_count: usize,
538 pub graph_id: usize,
539 pub colored: bool,
540 pub compressed: bool,
541 pub interleaved_compression: bool,
542 pub label_words: usize,
543 pub records: u64,
544}
545
546#[derive(Debug, Clone, PartialEq, Eq, Default)]
547pub struct BucketRecord {
549 pub graph_id: usize,
550 pub len: usize,
551 pub source_id: Option<u32>,
552 pub left_discontinuous: bool,
553 pub right_discontinuous: bool,
554 pub label: Vec<u8>,
555}
556
557#[derive(Debug, Clone, PartialEq, Eq, Default)]
558pub struct BucketPackedRecord {
560 pub graph_id: usize,
561 pub len: usize,
562 pub source_id: Option<u32>,
563 pub left_discontinuous: bool,
564 pub right_discontinuous: bool,
565 pub words: Vec<u64>,
566}
567
568pub(crate) struct BorrowedBucketPackedRecord<'a> {
569 pub graph_id: usize,
570 pub len: usize,
571 pub source_id: Option<u32>,
572 pub left_discontinuous: bool,
573 pub right_discontinuous: bool,
574 pub words: &'a [u64],
575}
576
577enum BucketSource {
583 File(File),
584 Chain(SegmentChainReader),
585}
586
587impl Read for BucketSource {
588 #[inline]
589 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
590 match self {
591 Self::File(file) => file.read(buf),
592 Self::Chain(chain) => chain.read(buf),
593 }
594 }
595}
596
597pub struct BucketStore {
604 containers: Option<Arc<BucketContainers>>,
605 header: Option<ContainerManifestHeader>,
606}
607
608impl BucketStore {
609 pub fn open_dir(
611 bucket_dir: impl AsRef<Path>,
612 ) -> Result<(Self, Vec<BucketManifestEntry>), BucketError> {
613 let bucket_dir = bucket_dir.as_ref();
614 if let Some((header, entries)) = read_container_manifest(bucket_dir)? {
615 let containers =
616 BucketContainers::open(bucket_dir, header.container_count, header.segment_bytes)?;
617 return Ok((
618 Self {
619 containers: Some(Arc::new(containers)),
620 header: Some(header),
621 },
622 entries,
623 ));
624 }
625 let entries = read_manifest(bucket_dir)?;
626 Ok((
627 Self {
628 containers: None,
629 header: None,
630 },
631 entries,
632 ))
633 }
634
635 pub fn files_only() -> Self {
637 Self {
638 containers: None,
639 header: None,
640 }
641 }
642
643 pub fn containers(&self) -> Option<&Arc<BucketContainers>> {
644 self.containers.as_ref()
645 }
646
647 pub fn reader(&self, entry: &BucketManifestEntry) -> Result<BucketReader, BucketError> {
649 match &entry.location {
650 BucketLocation::File(path) => BucketReader::open(path),
651 BucketLocation::Container {
652 container,
653 segments,
654 bytes,
655 } => {
656 let (Some(containers), Some(header)) = (&self.containers, &self.header) else {
657 return Err(BucketError::MalformedRecord);
658 };
659 BucketReader::open_chain(
660 Arc::clone(containers),
661 *container,
662 segments.clone(),
663 *bytes,
664 header.bucket_header(entry.graph_id, entry.records),
665 )
666 }
667 }
668 }
669}
670
671pub struct BucketReader {
673 path: PathBuf,
674 file: BufReader<BucketSource>,
675 header: BucketHeader,
676 records_read: u64,
677 block_attrs: Vec<u8>,
678 block_labels: Vec<u8>,
679 block_interleaved: Vec<u8>,
680 compressed_block: Vec<u8>,
681 block_records: usize,
682 block_record: usize,
683}
684
685impl BucketReader {
686 pub fn open(path: impl AsRef<Path>) -> Result<Self, BucketError> {
687 let path = path.as_ref().to_path_buf();
688 let file = File::open(&path).map_err(|source| BucketError::Io {
689 path: path.clone(),
690 source,
691 })?;
692 let mut file = BufReader::with_capacity(1024 * 1024, BucketSource::File(file));
693 let header = read_header(&mut file, &path)?;
694
695 Ok(Self {
696 path,
697 file,
698 header,
699 records_read: 0,
700 block_attrs: Vec::new(),
701 block_labels: Vec::new(),
702 block_interleaved: Vec::new(),
703 compressed_block: Vec::new(),
704 block_records: 0,
705 block_record: 0,
706 })
707 }
708
709 fn open_chain(
715 containers: Arc<BucketContainers>,
716 container: usize,
717 segments: Vec<u32>,
718 bytes: u64,
719 header: BucketHeader,
720 ) -> Result<Self, BucketError> {
721 let segment_bytes = containers.segment_bytes();
722 let path = containers.files[container].path.clone();
723 let chain = SegmentChainReader {
724 containers,
725 container,
726 segments: segments
727 .iter()
728 .map(|s| u64::from(*s) * segment_bytes)
729 .collect(),
730 len: bytes,
731 pos: 0,
732 };
733 Ok(Self {
734 path,
735 file: BufReader::with_capacity(1024 * 1024, BucketSource::Chain(chain)),
736 header,
737 records_read: 0,
738 block_attrs: Vec::new(),
739 block_labels: Vec::new(),
740 block_interleaved: Vec::new(),
741 compressed_block: Vec::new(),
742 block_records: 0,
743 block_record: 0,
744 })
745 }
746
747 #[inline]
748 pub fn header(&self) -> &BucketHeader {
749 &self.header
750 }
751
752 pub fn next_record(&mut self) -> Result<Option<BucketRecord>, BucketError> {
753 let mut record = BucketRecord::default();
754 if self.next_record_into(&mut record)? {
755 Ok(Some(record))
756 } else {
757 Ok(None)
758 }
759 }
760
761 pub fn next_record_into(&mut self, record: &mut BucketRecord) -> Result<bool, BucketError> {
762 let mut packed = BucketPackedRecord::default();
763 if !self.next_packed_record_into(&mut packed)? {
764 return Ok(false);
765 }
766
767 record.graph_id = packed.graph_id;
768 record.len = packed.len;
769 record.source_id = packed.source_id;
770 record.left_discontinuous = packed.left_discontinuous;
771 record.right_discontinuous = packed.right_discontinuous;
772 decode_label_into(&packed.words, packed.len, &mut record.label)?;
773 Ok(true)
774 }
775
776 pub fn next_packed_record_into(
777 &mut self,
778 record: &mut BucketPackedRecord,
779 ) -> Result<bool, BucketError> {
780 if self.records_read == self.header.records {
781 return Ok(false);
782 }
783
784 let mut fixed = [0u8; 4];
785 let fixed_len = if self.header.colored { 4 } else { 2 };
786 if self.header.compressed {
787 if self.block_record == self.block_records {
788 self.read_compressed_block()?;
789 }
790 let start = self.block_record * fixed_len;
791 if self.header.interleaved_compression {
792 let record_len = record_size(self.header.colored, self.header.label_words);
793 let start = self.block_record * record_len;
794 fixed[..fixed_len]
795 .copy_from_slice(&self.block_interleaved[start..start + fixed_len]);
796 } else {
797 fixed[..fixed_len].copy_from_slice(&self.block_attrs[start..start + fixed_len]);
798 }
799 } else {
800 self.file
801 .read_exact(&mut fixed[..fixed_len])
802 .map_err(|source| BucketError::Io {
803 path: self.path.clone(),
804 source,
805 })?;
806 }
807 let packed_attr = if self.header.colored {
808 u32::from_le_bytes(fixed[0..4].try_into().unwrap())
809 } else {
810 u16::from_le_bytes(fixed[0..2].try_into().unwrap()) as u32
811 };
812
813 record.words.clear();
814 record.words.resize(self.header.label_words, 0);
815 for (word_idx, word) in record.words.iter_mut().enumerate() {
816 if self.header.compressed {
817 if self.header.interleaved_compression {
818 let record_len = record_size(self.header.colored, self.header.label_words);
819 let start = self.block_record * record_len + fixed_len + word_idx * 8;
820 *word = u64::from_le_bytes(
821 self.block_interleaved[start..start + 8].try_into().unwrap(),
822 );
823 } else {
824 let start = (self.block_record * self.header.label_words + word_idx) * 8;
825 *word =
826 u64::from_le_bytes(self.block_labels[start..start + 8].try_into().unwrap());
827 }
828 } else {
829 *word = read_u64(&mut self.file, &self.path)?;
830 }
831 }
832
833 let len = (packed_attr & 0xff) as usize;
834 let source_id = self
835 .header
836 .colored
837 .then_some((packed_attr >> 10) & MAX_SOURCE_ID);
838 record.graph_id = self.header.graph_id;
839 record.len = len;
840 record.source_id = source_id;
841 record.left_discontinuous = (packed_attr & (1 << 8)) != 0;
842 record.right_discontinuous = (packed_attr & (1 << 9)) != 0;
843
844 self.records_read += 1;
845 if self.header.compressed {
846 self.block_record += 1;
847 }
848 Ok(true)
849 }
850
851 pub fn try_for_each_packed_record<E, F>(
852 &mut self,
853 record: &mut BucketPackedRecord,
854 mut f: F,
855 ) -> Result<(), E>
856 where
857 E: From<BucketError>,
858 F: FnMut(&BucketPackedRecord) -> Result<(), E>,
859 {
860 let remaining = self.header.records.saturating_sub(self.records_read);
861 if remaining == 0 {
862 return Ok(());
863 }
864 if self.header.compressed {
865 while self.next_packed_record_into(record).map_err(E::from)? {
866 f(record)?;
867 }
868 return Ok(());
869 }
870 let record_size = record_size(self.header.colored, self.header.label_words);
871 let payload_bytes = remaining
872 .checked_mul(record_size as u64)
873 .ok_or(BucketError::TooManyRecords)?;
874 let mut payload = vec![0u8; payload_bytes as usize];
875 self.file
876 .read_exact(&mut payload)
877 .map_err(|source| BucketError::Io {
878 path: self.path.clone(),
879 source,
880 })?;
881
882 for chunk in payload.chunks_exact(record_size) {
883 let (packed_attr, words_start) = if self.header.colored {
884 (u32::from_le_bytes(chunk[0..4].try_into().unwrap()), 4)
885 } else {
886 (
887 u16::from_le_bytes(chunk[0..2].try_into().unwrap()) as u32,
888 2,
889 )
890 };
891
892 record.words.clear();
893 record.words.reserve(self.header.label_words);
894 for word_idx in 0..self.header.label_words {
895 let start = words_start + word_idx * 8;
896 record.words.push(u64::from_le_bytes(
897 chunk[start..start + 8].try_into().unwrap(),
898 ));
899 }
900
901 let len = (packed_attr & 0xff) as usize;
902 let source_id = self
903 .header
904 .colored
905 .then_some((packed_attr >> 10) & MAX_SOURCE_ID);
906 record.graph_id = self.header.graph_id;
907 record.len = len;
908 record.source_id = source_id;
909 record.left_discontinuous = (packed_attr & (1 << 8)) != 0;
910 record.right_discontinuous = (packed_attr & (1 << 9)) != 0;
911
912 self.records_read += 1;
913 f(record)?;
914 }
915 Ok(())
916 }
917
918 pub(crate) fn try_for_each_borrowed_packed_record<E, F>(&mut self, mut f: F) -> Result<(), E>
919 where
920 E: From<BucketError>,
921 F: FnMut(BorrowedBucketPackedRecord<'_>) -> Result<(), E>,
922 {
923 if !self.header.compressed {
924 let mut record = BucketPackedRecord::default();
925 return self.try_for_each_packed_record(&mut record, |record| {
926 f(BorrowedBucketPackedRecord {
927 graph_id: record.graph_id,
928 len: record.len,
929 source_id: record.source_id,
930 left_discontinuous: record.left_discontinuous,
931 right_discontinuous: record.right_discontinuous,
932 words: &record.words,
933 })
934 });
935 }
936 if self.header.label_words > 4 {
937 return Err(E::from(BucketError::MalformedRecord));
938 }
939
940 let fixed_len = if self.header.colored { 4 } else { 2 };
941 while self.records_read < self.header.records {
942 if self.block_record == self.block_records {
943 self.read_compressed_block().map_err(E::from)?;
944 }
945 let record_index = self.block_record;
946 let (packed_attr, words_bytes) = if self.header.interleaved_compression {
947 let record_len = record_size(self.header.colored, self.header.label_words);
948 let start = record_index * record_len;
949 let attr = if self.header.colored {
950 u32::from_le_bytes(self.block_interleaved[start..start + 4].try_into().unwrap())
951 } else {
952 u16::from_le_bytes(self.block_interleaved[start..start + 2].try_into().unwrap())
953 as u32
954 };
955 (
956 attr,
957 &self.block_interleaved[start + fixed_len..start + record_len],
958 )
959 } else {
960 let attr_start = record_index * fixed_len;
961 let attr = if self.header.colored {
962 u32::from_le_bytes(
963 self.block_attrs[attr_start..attr_start + 4]
964 .try_into()
965 .unwrap(),
966 )
967 } else {
968 u16::from_le_bytes(
969 self.block_attrs[attr_start..attr_start + 2]
970 .try_into()
971 .unwrap(),
972 ) as u32
973 };
974 let words_start = record_index * self.header.label_words * 8;
975 (
976 attr,
977 &self.block_labels[words_start..words_start + self.header.label_words * 8],
978 )
979 };
980 let mut words = [0u64; 4];
981 for (word, bytes) in words[..self.header.label_words]
982 .iter_mut()
983 .zip(words_bytes.chunks_exact(8))
984 {
985 *word = u64::from_le_bytes(bytes.try_into().unwrap());
986 }
987 self.records_read += 1;
988 self.block_record += 1;
989 f(BorrowedBucketPackedRecord {
990 graph_id: self.header.graph_id,
991 len: (packed_attr & 0xff) as usize,
992 source_id: self
993 .header
994 .colored
995 .then_some((packed_attr >> 10) & MAX_SOURCE_ID),
996 left_discontinuous: packed_attr & (1 << 8) != 0,
997 right_discontinuous: packed_attr & (1 << 9) != 0,
998 words: &words[..self.header.label_words],
999 })?;
1000 }
1001 Ok(())
1002 }
1003
1004 pub fn records(self) -> BucketRecords {
1005 BucketRecords { reader: self }
1006 }
1007
1008 fn read_compressed_block(&mut self) -> Result<(), BucketError> {
1009 let mut header = [0u8; COMPRESSED_BLOCK_HEADER_LEN];
1010 self.file
1011 .read_exact(&mut header)
1012 .map_err(|source| BucketError::Io {
1013 path: self.path.clone(),
1014 source,
1015 })?;
1016 let records = u32::from_le_bytes(header[0..4].try_into().unwrap()) as usize;
1017 let attr_bytes = u32::from_le_bytes(header[4..8].try_into().unwrap()) as usize;
1018 let label_bytes = u32::from_le_bytes(header[8..12].try_into().unwrap()) as usize;
1019 if records == 0
1020 || attr_bytes == 0
1021 || (!self.header.interleaved_compression && label_bytes == 0)
1022 || (self.header.interleaved_compression && label_bytes != 0)
1023 {
1024 return Err(BucketError::MalformedRecord);
1025 }
1026 let remaining = usize::try_from(self.header.records - self.records_read)
1027 .map_err(|_| BucketError::TooManyRecords)?;
1028 if records > remaining {
1029 return Err(BucketError::MalformedRecord);
1030 }
1031 self.compressed_block.resize(attr_bytes + label_bytes, 0);
1032 self.file
1033 .read_exact(&mut self.compressed_block)
1034 .map_err(|source| BucketError::Io {
1035 path: self.path.clone(),
1036 source,
1037 })?;
1038 let fixed_len = if self.header.colored { 4 } else { 2 };
1039 if self.header.interleaved_compression {
1040 self.block_interleaved.resize(
1041 records * record_size(self.header.colored, self.header.label_words),
1042 0,
1043 );
1044 let decoded = lz4_flex::block::decompress_into(
1045 &self.compressed_block[..attr_bytes],
1046 &mut self.block_interleaved,
1047 )
1048 .map_err(|_| BucketError::MalformedRecord)?;
1049 if decoded != self.block_interleaved.len() {
1050 return Err(BucketError::MalformedRecord);
1051 }
1052 self.block_records = records;
1053 self.block_record = 0;
1054 return Ok(());
1055 }
1056 self.block_attrs.resize(records * fixed_len, 0);
1057 self.block_labels
1058 .resize(records * self.header.label_words * 8, 0);
1059 let decoded_attrs = lz4_flex::block::decompress_into(
1060 &self.compressed_block[..attr_bytes],
1061 &mut self.block_attrs,
1062 )
1063 .map_err(|_| BucketError::MalformedRecord)?;
1064 let decoded_labels = lz4_flex::block::decompress_into(
1065 &self.compressed_block[attr_bytes..],
1066 &mut self.block_labels,
1067 )
1068 .map_err(|_| BucketError::MalformedRecord)?;
1069 if decoded_attrs != self.block_attrs.len() || decoded_labels != self.block_labels.len() {
1070 return Err(BucketError::MalformedRecord);
1071 }
1072 self.block_records = records;
1073 self.block_record = 0;
1074 Ok(())
1075 }
1076}
1077
1078pub struct BucketRecords {
1080 reader: BucketReader,
1081}
1082
1083impl Iterator for BucketRecords {
1084 type Item = Result<BucketRecord, BucketError>;
1085
1086 fn next(&mut self) -> Option<Self::Item> {
1087 self.reader.next_record().transpose()
1088 }
1089}
1090
1091const CONTAINER_MANIFEST_MAGIC: &[u8; 8] = b"CF3WSKC1";
1092const CONTAINER_MANIFEST_NAME: &str = "manifest.bin";
1093
1094#[derive(Debug, Clone, PartialEq, Eq)]
1102pub struct ContainerManifestHeader {
1103 pub k: u16,
1104 pub minimizer_len: u16,
1105 pub graph_count: usize,
1106 pub colored: bool,
1107 pub label_words: usize,
1108 pub compressed: bool,
1109 pub interleaved_compression: bool,
1110 pub segment_bytes: u64,
1111 pub container_count: usize,
1112}
1113
1114impl ContainerManifestHeader {
1115 fn compression_code(&self) -> u8 {
1116 if self.interleaved_compression {
1117 2
1118 } else {
1119 u8::from(self.compressed)
1120 }
1121 }
1122
1123 fn bucket_header(&self, graph_id: usize, records: u64) -> BucketHeader {
1125 BucketHeader {
1126 k: self.k,
1127 minimizer_len: self.minimizer_len,
1128 graph_count: self.graph_count,
1129 graph_id,
1130 colored: self.colored,
1131 label_words: self.label_words,
1132 compressed: self.compressed,
1133 interleaved_compression: self.interleaved_compression,
1134 records,
1135 }
1136 }
1137}
1138
1139fn write_u32_to(out: &mut Vec<u8>, value: u32) {
1140 out.extend_from_slice(&value.to_le_bytes());
1141}
1142
1143fn write_u64_to(out: &mut Vec<u8>, value: u64) {
1144 out.extend_from_slice(&value.to_le_bytes());
1145}
1146
1147pub fn write_container_manifest(
1149 bucket_dir: &Path,
1150 header: &ContainerManifestHeader,
1151 entries: &[BucketManifestEntry],
1152) -> Result<(), BucketError> {
1153 let path = bucket_dir.join(CONTAINER_MANIFEST_NAME);
1154 let mut out = Vec::with_capacity(64 + entries.len() * 32);
1155 out.extend_from_slice(CONTAINER_MANIFEST_MAGIC);
1156 out.extend_from_slice(&header.k.to_le_bytes());
1157 out.extend_from_slice(&header.minimizer_len.to_le_bytes());
1158 write_u64_to(&mut out, header.graph_count as u64);
1159 out.push(u8::from(header.colored));
1160 out.push(header.label_words as u8);
1161 out.push(header.compression_code());
1162 out.push(0);
1163 write_u64_to(&mut out, header.segment_bytes);
1164 write_u64_to(&mut out, header.container_count as u64);
1165 write_u64_to(&mut out, entries.len() as u64);
1166 for entry in entries {
1167 let BucketLocation::Container {
1168 container,
1169 segments,
1170 bytes,
1171 } = &entry.location
1172 else {
1173 return Err(BucketError::MalformedManifest(path.clone()));
1174 };
1175 write_u64_to(&mut out, entry.graph_id as u64);
1176 write_u64_to(&mut out, entry.records);
1177 write_u64_to(&mut out, *bytes);
1178 write_u32_to(&mut out, *container as u32);
1179 write_u32_to(&mut out, segments.len() as u32);
1180 for segment in segments {
1181 write_u32_to(&mut out, *segment);
1182 }
1183 }
1184 fs::write(&path, &out).map_err(|source| BucketError::Io {
1185 path: path.clone(),
1186 source,
1187 })
1188}
1189
1190struct ManifestCursor<'a> {
1191 bytes: &'a [u8],
1192 pos: usize,
1193 path: &'a Path,
1194}
1195
1196impl<'a> ManifestCursor<'a> {
1197 fn take(&mut self, len: usize) -> Result<&'a [u8], BucketError> {
1198 let end = self
1199 .pos
1200 .checked_add(len)
1201 .filter(|end| *end <= self.bytes.len())
1202 .ok_or_else(|| BucketError::MalformedManifest(self.path.to_path_buf()))?;
1203 let slice = &self.bytes[self.pos..end];
1204 self.pos = end;
1205 Ok(slice)
1206 }
1207
1208 fn u16(&mut self) -> Result<u16, BucketError> {
1209 Ok(u16::from_le_bytes(self.take(2)?.try_into().unwrap()))
1210 }
1211
1212 fn u32(&mut self) -> Result<u32, BucketError> {
1213 Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap()))
1214 }
1215
1216 fn u64(&mut self) -> Result<u64, BucketError> {
1217 Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap()))
1218 }
1219}
1220
1221pub fn read_container_manifest(
1223 bucket_dir: impl AsRef<Path>,
1224) -> Result<Option<(ContainerManifestHeader, Vec<BucketManifestEntry>)>, BucketError> {
1225 let path = bucket_dir.as_ref().join(CONTAINER_MANIFEST_NAME);
1226 let bytes = match fs::read(&path) {
1227 Ok(bytes) => bytes,
1228 Err(source) if source.kind() == std::io::ErrorKind::NotFound => return Ok(None),
1229 Err(source) => {
1230 return Err(BucketError::Io {
1231 path: path.clone(),
1232 source,
1233 });
1234 }
1235 };
1236 let mut cursor = ManifestCursor {
1237 bytes: &bytes,
1238 pos: 0,
1239 path: &path,
1240 };
1241 if cursor.take(8)? != CONTAINER_MANIFEST_MAGIC {
1242 return Err(BucketError::MalformedManifest(path.clone()));
1243 }
1244 let k = cursor.u16()?;
1245 let minimizer_len = cursor.u16()?;
1246 let graph_count = cursor.u64()? as usize;
1247 let flags = cursor.take(4)?;
1248 let (colored, label_words, compression) = (flags[0], flags[1], flags[2]);
1249 if colored > 1 || compression > 2 || flags[3] != 0 {
1250 return Err(BucketError::MalformedManifest(path.clone()));
1251 }
1252 let header = ContainerManifestHeader {
1253 k,
1254 minimizer_len,
1255 graph_count,
1256 colored: colored == 1,
1257 label_words: label_words as usize,
1258 compressed: compression != 0,
1259 interleaved_compression: compression == 2,
1260 segment_bytes: cursor.u64()?,
1261 container_count: cursor.u64()? as usize,
1262 };
1263 if header.segment_bytes == 0
1264 || header.label_words as usize != label_word_count(k, minimizer_len)
1265 {
1266 return Err(BucketError::MalformedManifest(path.clone()));
1267 }
1268 let bucket_count = cursor.u64()? as usize;
1269 let mut entries = Vec::with_capacity(bucket_count);
1270 for _ in 0..bucket_count {
1271 let graph_id = cursor.u64()? as usize;
1272 let records = cursor.u64()?;
1273 let bytes_len = cursor.u64()?;
1274 let container = cursor.u32()? as usize;
1275 let segment_count = cursor.u32()? as usize;
1276 if graph_id >= graph_count || container >= header.container_count {
1277 return Err(BucketError::MalformedManifest(path.clone()));
1278 }
1279 let mut segments = Vec::with_capacity(segment_count);
1280 for _ in 0..segment_count {
1281 segments.push(cursor.u32()?);
1282 }
1283 let capacity = segment_count as u64 * header.segment_bytes;
1286 if bytes_len > capacity || capacity.saturating_sub(bytes_len) >= header.segment_bytes {
1287 return Err(BucketError::MalformedManifest(path.clone()));
1288 }
1289 entries.push(BucketManifestEntry {
1290 graph_id,
1291 records,
1292 location: BucketLocation::Container {
1293 container,
1294 segments,
1295 bytes: bytes_len,
1296 },
1297 });
1298 }
1299 Ok(Some((header, entries)))
1300}
1301
1302pub fn read_manifest(
1304 bucket_dir: impl AsRef<Path>,
1305) -> Result<Vec<BucketManifestEntry>, BucketError> {
1306 let bucket_dir = bucket_dir.as_ref();
1307 let path = bucket_dir.join("manifest.tsv");
1308 let file = File::open(&path).map_err(|source| BucketError::Io {
1309 path: path.clone(),
1310 source,
1311 })?;
1312 let mut out = Vec::new();
1313
1314 for (line_no, line) in BufReader::new(file).lines().enumerate() {
1315 let line = line.map_err(|source| BucketError::Io {
1316 path: path.clone(),
1317 source,
1318 })?;
1319 if line_no == 0 {
1320 if line != "graph_id\trecords\tpath" {
1321 return Err(BucketError::MalformedManifest(path.clone()));
1322 }
1323 continue;
1324 }
1325 if line.trim().is_empty() {
1326 continue;
1327 }
1328
1329 let mut fields = line.splitn(3, '\t');
1330 let graph_id = fields
1331 .next()
1332 .ok_or_else(|| BucketError::MalformedManifest(path.clone()))?
1333 .parse()
1334 .map_err(|_| BucketError::MalformedManifest(path.clone()))?;
1335 let records = fields
1336 .next()
1337 .ok_or_else(|| BucketError::MalformedManifest(path.clone()))?
1338 .parse()
1339 .map_err(|_| BucketError::MalformedManifest(path.clone()))?;
1340 let bucket_path = fields
1341 .next()
1342 .ok_or_else(|| BucketError::MalformedManifest(path.clone()))?;
1343 out.push(BucketManifestEntry {
1344 graph_id,
1345 records,
1346 location: BucketLocation::File(PathBuf::from(bucket_path)),
1347 });
1348 }
1349
1350 Ok(out)
1351}
1352
1353pub fn coalesce_bucket_manifest(
1354 bucket_dir: &Path,
1355 entries: &[BucketManifestEntry],
1356) -> Result<BucketEmitStats, BucketError> {
1357 coalesce_bucket_manifest_with_threads(bucket_dir, entries, 1)
1358}
1359
1360pub fn coalesce_bucket_manifest_with_threads(
1361 bucket_dir: &Path,
1362 entries: &[BucketManifestEntry],
1363 threads: usize,
1364) -> Result<BucketEmitStats, BucketError> {
1365 let mut by_graph = BTreeMap::<usize, Vec<BucketManifestEntry>>::new();
1366 for entry in entries {
1367 by_graph
1368 .entry(entry.graph_id)
1369 .or_default()
1370 .push(entry.clone());
1371 }
1372
1373 let groups = by_graph.into_iter().collect::<Vec<_>>();
1374 let workers = threads.max(1).min(groups.len().max(1));
1375 let mut manifest = if workers == 1 {
1376 let mut manifest = Vec::with_capacity(groups.len());
1377 for group in &groups {
1378 manifest.push(coalesce_bucket_group(bucket_dir, group)?);
1379 }
1380 manifest
1381 } else {
1382 let next_group = AtomicUsize::new(0);
1383 let mut manifest = std::thread::scope(|scope| {
1384 let mut handles = Vec::new();
1385 for _ in 0..workers {
1386 let next_group = &next_group;
1387 let groups = &groups;
1388 handles.push(scope.spawn(move || {
1389 let mut local = Vec::new();
1390 loop {
1391 let group_idx = next_group.fetch_add(1, Ordering::Relaxed);
1392 let Some(group) = groups.get(group_idx) else {
1393 break;
1394 };
1395 local.push(coalesce_bucket_group(bucket_dir, group)?);
1396 }
1397 Ok::<_, BucketError>(local)
1398 }));
1399 }
1400
1401 let mut manifest = Vec::with_capacity(groups.len());
1402 for handle in handles {
1403 manifest.extend(handle.join().map_err(|_| BucketError::WorkerPanic)??);
1404 }
1405 Ok::<_, BucketError>(manifest)
1406 })?;
1407 manifest.sort_by_key(|(graph_id, _, path, _)| (*graph_id, path.clone()));
1408 manifest
1409 };
1410 let total_bytes = manifest.iter().map(|(_, _, _, bytes)| *bytes).sum();
1411 let public_manifest = manifest
1412 .drain(..)
1413 .map(|(graph_id, records, path, _)| (graph_id, records, path))
1414 .collect::<Vec<_>>();
1415
1416 write_manifest(bucket_dir, &public_manifest)?;
1417
1418 Ok(BucketEmitStats {
1419 bucket_dir: bucket_dir.to_path_buf(),
1420 bucket_files: public_manifest.len(),
1421 bytes_written: total_bytes,
1422 })
1423}
1424
1425fn coalesce_bucket_group(
1426 bucket_dir: &Path,
1427 group: &(usize, Vec<BucketManifestEntry>),
1428) -> Result<(usize, u64, PathBuf, u64), BucketError> {
1429 let (graph_id, graph_entries) = group;
1430 let entry_path = |entry: &BucketManifestEntry| {
1433 entry
1434 .file_path()
1435 .map(Path::to_path_buf)
1436 .ok_or_else(|| BucketError::MalformedManifest(bucket_dir.to_path_buf()))
1437 };
1438 let first_header = read_bucket_header(&entry_path(&graph_entries[0])?)?;
1439 let mut writer = BucketFile::create(
1440 bucket_dir,
1441 first_header.k,
1442 first_header.minimizer_len,
1443 first_header.graph_count,
1444 *graph_id,
1445 first_header.colored,
1446 first_header.label_words,
1447 first_header.compressed,
1448 )?;
1449
1450 for entry in graph_entries {
1451 let path = entry_path(entry)?;
1452 let copied_records = copy_bucket_payload(&path, &first_header, *graph_id, &mut writer)?;
1453 if copied_records != entry.records {
1454 return Err(BucketError::MalformedHeader(path));
1455 }
1456 }
1457
1458 writer.finish()?;
1459 Ok((
1460 *graph_id,
1461 writer.records,
1462 writer.path.clone(),
1463 writer.bytes_written,
1464 ))
1465}
1466
1467fn read_bucket_header(path: &Path) -> Result<BucketHeader, BucketError> {
1468 let mut file = File::open(path).map_err(|source| BucketError::Io {
1469 path: path.to_path_buf(),
1470 source,
1471 })?;
1472 read_header(&mut file, path)
1473}
1474
1475fn copy_bucket_payload(
1476 path: &Path,
1477 expected: &BucketHeader,
1478 graph_id: usize,
1479 writer: &mut BucketFile,
1480) -> Result<u64, BucketError> {
1481 let mut file = File::open(path).map_err(|source| BucketError::Io {
1482 path: path.to_path_buf(),
1483 source,
1484 })?;
1485 let header = read_header(&mut file, path)?;
1486 if header.k != expected.k
1487 || header.minimizer_len != expected.minimizer_len
1488 || header.graph_count != expected.graph_count
1489 || header.graph_id != graph_id
1490 || header.colored != expected.colored
1491 || header.label_words != expected.label_words
1492 {
1493 return Err(BucketError::MalformedHeader(path.to_path_buf()));
1494 }
1495
1496 let payload_bytes = header
1497 .records
1498 .checked_mul(record_size(header.colored, header.label_words) as u64)
1499 .ok_or(BucketError::TooManyRecords)?;
1500 let actual_len = file
1501 .metadata()
1502 .map_err(|source| BucketError::Io {
1503 path: path.to_path_buf(),
1504 source,
1505 })?
1506 .len();
1507 if actual_len != HEADER_LEN + payload_bytes {
1508 return Err(BucketError::MalformedHeader(path.to_path_buf()));
1509 }
1510
1511 file.seek(SeekFrom::Start(HEADER_LEN))
1512 .map_err(|source| BucketError::Io {
1513 path: path.to_path_buf(),
1514 source,
1515 })?;
1516 let copied =
1517 std::io::copy(&mut file.take(payload_bytes), &mut writer.file).map_err(|source| {
1518 BucketError::Io {
1519 path: path.to_path_buf(),
1520 source,
1521 }
1522 })?;
1523 if copied != payload_bytes {
1524 return Err(BucketError::MalformedHeader(path.to_path_buf()));
1525 }
1526 writer.records = writer
1527 .records
1528 .checked_add(header.records)
1529 .ok_or(BucketError::TooManyRecords)?;
1530 writer.bytes_written += copied;
1531 Ok(header.records)
1532}
1533
1534impl BucketEmitter {
1535 pub fn create(params: &BuildParams, graph_count: usize) -> Result<Self, BucketError> {
1536 Self::create_in_dir(params, graph_count, bucket_dir(params))
1537 }
1538
1539 pub fn create_in_dir(
1540 params: &BuildParams,
1541 graph_count: usize,
1542 bucket_dir: PathBuf,
1543 ) -> Result<Self, BucketError> {
1544 if graph_count > u64::MAX as usize {
1545 return Err(BucketError::GraphCountTooLarge(graph_count));
1546 }
1547
1548 if bucket_dir.exists() {
1549 fs::remove_dir_all(&bucket_dir).map_err(|source| BucketError::Io {
1550 path: bucket_dir.clone(),
1551 source,
1552 })?;
1553 }
1554 fs::create_dir_all(&bucket_dir).map_err(|source| BucketError::Io {
1555 path: bucket_dir.clone(),
1556 source,
1557 })?;
1558
1559 Ok(Self {
1560 bucket_dir,
1561 k: params.k,
1562 minimizer_len: params.minimizer_len,
1563 graph_count,
1564 colored: params.color,
1565 compress_buckets: params.color || params.compress_buckets,
1566 label_words: label_word_count(params.k, params.minimizer_len),
1567 files: BTreeMap::new(),
1568 scratch: CompressionScratch::default(),
1569 writers: BTreeMap::new(),
1570 pending: vec![PendingBucket::default(); graph_count],
1571 pending_bytes: 0,
1572 })
1573 }
1574
1575 pub fn add(&mut self, superkmer: &WeakSuperKmer, seq: &[u8]) -> Result<(), BucketError> {
1576 if superkmer.graph_id >= self.graph_count {
1577 return Err(BucketError::InvalidGraphId(superkmer.graph_id));
1578 }
1579 if seq.len() > u8::MAX as usize {
1580 return Err(BucketError::LabelTooLong(seq.len()));
1581 }
1582
1583 let attr = if self.colored {
1584 let source_id = superkmer.source_id.ok_or(BucketError::MissingSourceId)?;
1585 if source_id > MAX_SOURCE_ID {
1586 return Err(BucketError::SourceIdTooLarge(source_id));
1587 }
1588 pack_colored_attr(
1589 seq.len(),
1590 source_id,
1591 superkmer.left_discontinuous,
1592 superkmer.right_discontinuous,
1593 )
1594 } else {
1595 pack_uncolored_attr(
1596 seq.len(),
1597 superkmer.left_discontinuous,
1598 superkmer.right_discontinuous,
1599 )
1600 };
1601 let graph_id = superkmer.graph_id;
1602 let record_len = record_size(self.colored, self.label_words);
1603 let pending = &mut self.pending[graph_id];
1604 pending.records = pending
1605 .records
1606 .checked_add(1)
1607 .ok_or(BucketError::TooManyRecords)?;
1608 append_record(
1609 &mut pending.bytes,
1610 attr,
1611 graph_id,
1612 seq,
1613 self.label_words,
1614 self.colored,
1615 )?;
1616 self.pending_bytes += record_len;
1617
1618 if self.pending[graph_id].bytes.len() >= MAX_PENDING_BUCKET_BYTES {
1619 self.flush_pending_bucket(graph_id)?;
1620 } else if self.pending_bytes >= MAX_TOTAL_PENDING_BYTES {
1621 self.flush_largest_pending_bucket()?;
1622 }
1623 Ok(())
1624 }
1625
1626 fn flush_largest_pending_bucket(&mut self) -> Result<(), BucketError> {
1627 let Some((graph_id, _)) = self
1628 .pending
1629 .iter()
1630 .enumerate()
1631 .max_by_key(|(_, pending)| pending.bytes.len())
1632 else {
1633 return Ok(());
1634 };
1635 self.flush_pending_bucket(graph_id)
1636 }
1637
1638 fn flush_pending_bucket(&mut self, graph_id: usize) -> Result<(), BucketError> {
1639 if self.pending[graph_id].bytes.is_empty() {
1640 return Ok(());
1641 }
1642 let pending = std::mem::take(&mut self.pending[graph_id]);
1643 self.pending_bytes -= pending.bytes.len();
1644
1645 let (records, bytes_written) = {
1646 self.ensure_writer(graph_id)?;
1647 let Self {
1648 writers, scratch, ..
1649 } = self;
1650 let writer = writers.get_mut(&graph_id).expect("writer just ensured");
1651 writer.write_records(&pending.bytes, pending.records, scratch)?
1652 };
1653 let meta = self.files.get_mut(&graph_id).unwrap();
1654 meta.records = records;
1655 meta.bytes_written = bytes_written;
1656 Ok(())
1657 }
1658
1659 fn ensure_writer(&mut self, graph_id: usize) -> Result<&mut BucketFile, BucketError> {
1660 if !self.writers.contains_key(&graph_id) {
1661 self.evict_writer_if_needed(graph_id)?;
1662 let writer = match self.files.get(&graph_id) {
1663 Some(meta) => {
1664 BucketFile::open_existing(&meta.path, meta.records, meta.bytes_written)?
1665 }
1666 None => {
1667 let writer = BucketFile::create(
1668 &self.bucket_dir,
1669 self.k,
1670 self.minimizer_len,
1671 self.graph_count,
1672 graph_id,
1673 self.colored,
1674 self.label_words,
1675 self.compress_buckets,
1676 )?;
1677 self.files.insert(
1678 graph_id,
1679 BucketFileMeta {
1680 path: writer.path.clone(),
1681 records: writer.records,
1682 bytes_written: writer.bytes_written,
1683 },
1684 );
1685 writer
1686 }
1687 };
1688 self.writers.insert(graph_id, writer);
1689 }
1690
1691 Ok(self.writers.get_mut(&graph_id).unwrap())
1692 }
1693
1694 fn evict_writer_if_needed(&mut self, requested_graph_id: usize) -> Result<(), BucketError> {
1695 if self.writers.len() < MAX_OPEN_BUCKET_WRITERS {
1696 return Ok(());
1697 }
1698
1699 let evict_graph_id = self
1700 .writers
1701 .keys()
1702 .copied()
1703 .find(|&graph_id| graph_id != requested_graph_id)
1704 .unwrap_or(requested_graph_id);
1705 if let Some(mut writer) = self.writers.remove(&evict_graph_id) {
1706 writer.flush()?;
1707 }
1708 Ok(())
1709 }
1710
1711 pub fn finish(mut self) -> Result<BucketEmitStats, BucketError> {
1712 let mut manifest = Vec::new();
1713 for graph_id in 0..self.pending.len() {
1714 if !self.pending[graph_id].bytes.is_empty() {
1715 self.flush_pending_bucket(graph_id)?;
1716 }
1717 }
1718
1719 for (graph_id, meta) in &self.files {
1720 if let Some(writer) = self.writers.get_mut(graph_id) {
1721 writer.finish()?;
1722 } else {
1723 BucketFile::finish_closed(&meta.path, meta.records)?;
1724 }
1725 manifest.push((*graph_id, meta.records, meta.path.clone()));
1726 }
1727
1728 write_manifest(&self.bucket_dir, &manifest)?;
1729
1730 Ok(BucketEmitStats {
1731 bucket_dir: self.bucket_dir,
1732 bucket_files: manifest.len(),
1733 bytes_written: self.files.values().map(|meta| meta.bytes_written).sum(),
1734 })
1735 }
1736}
1737
1738impl SharedBucketSink {
1739 pub fn create(params: &BuildParams, graph_count: usize) -> Result<Arc<Self>, BucketError> {
1740 let bucket_dir = bucket_dir(params);
1741 if bucket_dir.exists() {
1742 fs::remove_dir_all(&bucket_dir).map_err(|source| BucketError::Io {
1743 path: bucket_dir.clone(),
1744 source,
1745 })?;
1746 }
1747 fs::create_dir_all(&bucket_dir).map_err(|source| BucketError::Io {
1748 path: bucket_dir.clone(),
1749 source,
1750 })?;
1751
1752 let atlas_count = graph_count.div_ceil(ATLAS_GRAPH_COUNT);
1757 let container_count = BucketContainers::plan_container_count(atlas_count);
1758 if container_count < atlas_count {
1759 eprintln!(
1760 "cuttlefish: descriptor budget allows {container_count} bucket container(s) rather than {atlas_count}"
1761 );
1762 }
1763 let containers = BucketContainers::create(&bucket_dir, container_count)?;
1764
1765 Ok(Arc::new(Self {
1766 bucket_dir,
1767 containers,
1768 k: params.k,
1769 minimizer_len: params.minimizer_len,
1770 graph_count,
1771 colored: params.color,
1772 compress_buckets: params.color || params.compress_buckets,
1773 label_words: label_word_count(params.k, params.minimizer_len),
1774 workers: params.threads.max(1),
1775 atlases: (0..graph_count.div_ceil(ATLAS_GRAPH_COUNT))
1776 .map(|atlas_id| {
1777 let first_graph_id = atlas_id * ATLAS_GRAPH_COUNT;
1778 let file_count = (graph_count - first_graph_id).min(ATLAS_GRAPH_COUNT);
1779 Mutex::new(SharedBucketAtlas {
1780 first_graph_id,
1781 files: (0..file_count)
1782 .map(|_| SharedBucketFileMeta::default())
1783 .collect(),
1784 buffered_bytes: 0,
1785 scratch: CompressionScratch::default(),
1786 })
1787 })
1788 .collect(),
1789 flush_calls: AtomicU64::new(0),
1790 flush_nanos: AtomicU64::new(0),
1791 }))
1792 }
1793
1794 pub fn emitter(self: &Arc<Self>) -> SharedBucketEmitter {
1795 self.make_emitter(false)
1796 }
1797
1798 pub fn deferred_uncolored_emitter(self: &Arc<Self>) -> SharedBucketEmitter {
1799 self.make_emitter(true)
1800 }
1801
1802 fn make_emitter(self: &Arc<Self>, deferred_uncolored: bool) -> SharedBucketEmitter {
1803 SharedBucketEmitter {
1804 sink: Arc::clone(self),
1805 pending: if self.colored || deferred_uncolored {
1806 Vec::new()
1807 } else {
1808 vec![PendingBucket::default(); self.graph_count]
1809 },
1810 uncolored_pending: if !self.colored && deferred_uncolored {
1811 (0..self.atlases.len())
1812 .map(|_| PendingColoredAtlas::default())
1813 .collect()
1814 } else {
1815 Vec::new()
1816 },
1817 colored_pending: if self.colored {
1818 (0..self.atlases.len())
1819 .map(|_| PendingColoredAtlas::default())
1820 .collect()
1821 } else {
1822 Vec::new()
1823 },
1824 pending_bytes: 0,
1825 deferred_uncolored,
1826 }
1827 }
1828
1829 pub fn flush_uncolored_emitters(
1830 &self,
1831 emitters: Vec<SharedBucketEmitter>,
1832 ) -> Result<(), BucketError> {
1833 let started = Instant::now();
1834 let mut by_atlas = (0..self.atlases.len())
1835 .map(|_| Mutex::new(Vec::<PendingColoredAtlas>::new()))
1836 .collect::<Vec<_>>();
1837 for mut emitter in emitters {
1838 for (atlas_id, pending) in emitter.uncolored_pending.drain(..).enumerate() {
1839 if !pending.bytes.is_empty() {
1840 by_atlas[atlas_id]
1841 .get_mut()
1842 .map_err(|_| BucketError::WorkerPanic)?
1843 .push(pending);
1844 }
1845 }
1846 }
1847 let next = AtomicUsize::new(0);
1848 let workers = self.workers.min(self.atlases.len().max(1));
1849 let calls = std::thread::scope(|scope| {
1850 let mut handles = Vec::with_capacity(workers);
1851 for _ in 0..workers {
1852 handles.push(scope.spawn(|| {
1853 let mut calls = 0;
1854 loop {
1855 let atlas_id = next.fetch_add(1, Ordering::Relaxed);
1856 let Some(pending) = by_atlas.get(atlas_id) else {
1857 break;
1858 };
1859 let pending = std::mem::take(
1860 &mut *pending.lock().map_err(|_| BucketError::WorkerPanic)?,
1861 );
1862 if pending.is_empty() {
1863 continue;
1864 }
1865 let mut stats = SharedBucketFlushStats::default();
1866 for chunk in pending {
1867 self.append_uncolored_atlas_inner(atlas_id, chunk, &mut stats)?;
1868 }
1869 let mut atlas = self.atlases[atlas_id]
1870 .lock()
1871 .map_err(|_| BucketError::WorkerPanic)?;
1872 atlas.flush_all(
1873 &self.containers,
1874 false,
1875 self.label_words,
1876 self.compress_buckets,
1877 &mut stats,
1878 )?;
1879 calls += stats.calls;
1880 }
1881 Ok::<_, BucketError>(calls)
1882 }));
1883 }
1884 let mut calls = 0;
1885 for handle in handles {
1886 calls += handle.join().map_err(|_| BucketError::WorkerPanic)??;
1887 }
1888 Ok::<_, BucketError>(calls)
1889 })?;
1890 self.record_flush_stats(SharedBucketFlushStats { calls }, started.elapsed());
1891 Ok(())
1892 }
1893
1894 fn append_uncolored_atlas(
1895 &self,
1896 atlas_id: usize,
1897 pending: PendingColoredAtlas,
1898 ) -> Result<(), BucketError> {
1899 let started = Instant::now();
1900 let mut stats = SharedBucketFlushStats::default();
1901 self.append_uncolored_atlas_inner(atlas_id, pending, &mut stats)?;
1902 self.record_flush_stats(stats, started.elapsed());
1903 Ok(())
1904 }
1905
1906 fn append_uncolored_atlas_inner(
1907 &self,
1908 atlas_id: usize,
1909 pending: PendingColoredAtlas,
1910 stats: &mut SharedBucketFlushStats,
1911 ) -> Result<(), BucketError> {
1912 if pending.bytes.is_empty() {
1913 return Ok(());
1914 }
1915 let record_len = record_size(false, self.label_words);
1916 if pending.graph_ids.len() * record_len != pending.bytes.len() {
1917 return Err(BucketError::MalformedRecord);
1918 }
1919 let mut atlas = self.atlases[atlas_id]
1920 .lock()
1921 .map_err(|_| BucketError::WorkerPanic)?;
1922 for (&graph_id, record) in pending
1923 .graph_ids
1924 .iter()
1925 .zip(pending.bytes.chunks_exact(record_len))
1926 {
1927 let graph_id = usize::from(graph_id);
1928 let Some(local_graph_id) = graph_id.checked_sub(atlas.first_graph_id) else {
1929 return Err(BucketError::InvalidGraphId(graph_id));
1930 };
1931 let Some(file) = atlas.files.get_mut(local_graph_id) else {
1932 return Err(BucketError::InvalidGraphId(graph_id));
1933 };
1934 file.total_records = file
1935 .total_records
1936 .checked_add(1)
1937 .ok_or(BucketError::TooManyRecords)?;
1938 file.buffer_records = file
1939 .buffer_records
1940 .checked_add(1)
1941 .ok_or(BucketError::TooManyRecords)?;
1942 file.buffer.extend_from_slice(record);
1943 }
1944 atlas.buffered_bytes += pending.bytes.len();
1945 for local_graph_id in 0..atlas.files.len() {
1946 if atlas.files[local_graph_id].buffer.len() >= SUBGRAPH_CHUNK_BYTES {
1947 atlas.flush_subgraph(
1948 local_graph_id,
1949 &self.containers,
1950 false,
1951 self.label_words,
1952 self.compress_buckets,
1953 stats,
1954 )?;
1955 }
1956 }
1957 Ok(())
1958 }
1959
1960 fn append_bucket(&self, graph_id: usize, pending: PendingBucket) -> Result<(), BucketError> {
1961 if pending.bytes.is_empty() {
1962 return Ok(());
1963 }
1964 let started = Instant::now();
1965 let atlas_id = graph_id / ATLAS_GRAPH_COUNT;
1966 let local_graph_id = graph_id % ATLAS_GRAPH_COUNT;
1967 let mut atlas = self.atlases[atlas_id]
1968 .lock()
1969 .map_err(|_| BucketError::WorkerPanic)?;
1970 atlas.append_bucket(local_graph_id, pending)?;
1971 if self.colored {
1972 return Ok(());
1973 }
1974 if atlas.files[local_graph_id].buffer.len() < SUBGRAPH_CHUNK_BYTES {
1975 return Ok(());
1976 }
1977
1978 let mut flush_stats = SharedBucketFlushStats::default();
1979 atlas.flush_subgraph(
1980 local_graph_id,
1981 &self.containers,
1982 self.colored,
1983 self.label_words,
1984 self.compress_buckets,
1985 &mut flush_stats,
1986 )?;
1987 self.record_flush_stats(flush_stats, started.elapsed());
1988 Ok(())
1989 }
1990
1991 fn append_colored_atlas(
1992 &self,
1993 atlas_id: usize,
1994 pending: PendingColoredAtlas,
1995 ) -> Result<(), BucketError> {
1996 if pending.bytes.is_empty() {
1997 return Ok(());
1998 }
1999 let record_len = record_size(true, self.label_words);
2000 if pending.graph_ids.len() * record_len != pending.bytes.len() {
2001 return Err(BucketError::MalformedRecord);
2002 }
2003 let started = Instant::now();
2004 let mut atlas = self.atlases[atlas_id]
2005 .lock()
2006 .map_err(|_| BucketError::WorkerPanic)?;
2007 for (&graph_id, record) in pending
2008 .graph_ids
2009 .iter()
2010 .zip(pending.bytes.chunks_exact(record_len))
2011 {
2012 let graph_id = usize::from(graph_id);
2013 let Some(local_graph_id) = graph_id.checked_sub(atlas.first_graph_id) else {
2014 return Err(BucketError::InvalidGraphId(graph_id));
2015 };
2016 let Some(file) = atlas.files.get_mut(local_graph_id) else {
2017 return Err(BucketError::InvalidGraphId(graph_id));
2018 };
2019 file.total_records = file
2020 .total_records
2021 .checked_add(1)
2022 .ok_or(BucketError::TooManyRecords)?;
2023 file.buffer_records = file
2024 .buffer_records
2025 .checked_add(1)
2026 .ok_or(BucketError::TooManyRecords)?;
2027 file.buffer.extend_from_slice(record);
2028 }
2029 atlas.buffered_bytes += pending.bytes.len();
2030 let mut flush_stats = SharedBucketFlushStats::default();
2031 for local_graph_id in 0..atlas.files.len() {
2032 if atlas.files[local_graph_id].buffer.len() >= SUBGRAPH_CHUNK_BYTES {
2033 atlas.flush_subgraph(
2034 local_graph_id,
2035 &self.containers,
2036 true,
2037 self.label_words,
2038 self.compress_buckets,
2039 &mut flush_stats,
2040 )?;
2041 }
2042 }
2043 self.record_flush_stats(flush_stats, started.elapsed());
2044 Ok(())
2045 }
2046
2047 pub fn flush_stats(&self) -> (u64, Duration) {
2048 (
2049 self.flush_calls.load(Ordering::Relaxed),
2050 Duration::from_nanos(self.flush_nanos.load(Ordering::Relaxed)),
2051 )
2052 }
2053
2054 pub fn flush_colored_window(
2055 &self,
2056 source_min: u32,
2057 source_max: u32,
2058 ) -> Result<(), BucketError> {
2059 if !self.colored || source_min > source_max {
2060 return Err(BucketError::MalformedRecord);
2061 }
2062 let started = Instant::now();
2063 let record_len = record_size(true, self.label_words);
2064 let next_atlas = AtomicUsize::new(0);
2065 let workers = self.workers.min(self.atlases.len().max(1));
2066 let flush_calls = std::thread::scope(|scope| {
2067 let mut handles = Vec::with_capacity(workers);
2068 for _ in 0..workers {
2069 handles.push(scope.spawn(|| {
2070 let mut stats = SharedBucketFlushStats::default();
2071 loop {
2072 let atlas_id = next_atlas.fetch_add(1, Ordering::Relaxed);
2073 let Some(atlas) = self.atlases.get(atlas_id) else {
2074 break;
2075 };
2076 let mut atlas = atlas.lock().map_err(|_| BucketError::WorkerPanic)?;
2077 for local_graph_id in 0..atlas.files.len() {
2078 sort_colored_payload_by_source(
2079 &mut atlas.files[local_graph_id].buffer,
2080 record_len,
2081 source_min,
2082 source_max,
2083 )?;
2084 atlas.flush_subgraph(
2085 local_graph_id,
2086 &self.containers,
2087 true,
2088 self.label_words,
2089 true,
2090 &mut stats,
2091 )?;
2092 }
2093 }
2094 Ok::<_, BucketError>(stats.calls)
2095 }));
2096 }
2097 let mut calls = 0u64;
2098 for handle in handles {
2099 calls += handle.join().map_err(|_| BucketError::WorkerPanic)??;
2100 }
2101 Ok::<_, BucketError>(calls)
2102 })?;
2103 self.record_flush_stats(
2104 SharedBucketFlushStats { calls: flush_calls },
2105 started.elapsed(),
2106 );
2107 Ok(())
2108 }
2109
2110 pub fn flush_colored_emitters(
2111 &self,
2112 emitters: Vec<SharedBucketEmitter>,
2113 ) -> Result<(), BucketError> {
2114 if !self.colored {
2115 return Err(BucketError::MalformedRecord);
2116 }
2117
2118 let emitter_count = emitters.len();
2119 let mut pending_by_atlas = (0..self.atlases.len())
2120 .map(|_| Vec::with_capacity(emitter_count))
2121 .collect::<Vec<Vec<PendingColoredAtlas>>>();
2122 for emitter in emitters {
2123 if !std::ptr::eq(Arc::as_ptr(&emitter.sink), self) {
2124 return Err(BucketError::MalformedRecord);
2125 }
2126 if !emitter.pending.is_empty()
2127 || emitter.colored_pending.len() != pending_by_atlas.len()
2128 {
2129 return Err(BucketError::MalformedRecord);
2130 }
2131 for (atlas_id, pending) in emitter.colored_pending.into_iter().enumerate() {
2132 if !pending.bytes.is_empty() {
2133 pending_by_atlas[atlas_id].push(pending);
2134 }
2135 }
2136 }
2137
2138 let started = Instant::now();
2139 let workers = self.workers.min(pending_by_atlas.len().max(1));
2140 let chunk_size = pending_by_atlas.len().div_ceil(workers);
2141 let flush_calls = std::thread::scope(|scope| {
2142 let mut handles = Vec::with_capacity(workers);
2143 for (chunk_id, atlas_work) in pending_by_atlas.chunks_mut(chunk_size).enumerate() {
2144 handles.push(scope.spawn(move || {
2145 let mut stats = SharedBucketFlushStats::default();
2146 let first_atlas = chunk_id * chunk_size;
2147 for (offset, worker_buckets) in atlas_work.iter_mut().enumerate() {
2148 let atlas_id = first_atlas + offset;
2149 let mut atlas = self.atlases[atlas_id]
2150 .lock()
2151 .map_err(|_| BucketError::WorkerPanic)?;
2152 let record_len = record_size(true, self.label_words);
2153 for pending in worker_buckets.iter() {
2154 if pending.graph_ids.len() * record_len != pending.bytes.len() {
2155 return Err(BucketError::MalformedRecord);
2156 }
2157 for (&graph_id, record) in pending
2158 .graph_ids
2159 .iter()
2160 .zip(pending.bytes.chunks_exact(record_len))
2161 {
2162 append_colored_atlas_record(&mut atlas, graph_id, record)?;
2163 }
2164 }
2165 atlas.buffered_bytes += worker_buckets
2166 .iter()
2167 .map(|pending| pending.bytes.len())
2168 .sum::<usize>();
2169 for local_graph_id in 0..atlas.files.len() {
2170 atlas.flush_subgraph(
2171 local_graph_id,
2172 &self.containers,
2173 true,
2174 self.label_words,
2175 true,
2176 &mut stats,
2177 )?;
2178 atlas.files[local_graph_id].buffer = Vec::new();
2183 }
2184 }
2185 Ok::<_, BucketError>(stats.calls)
2186 }));
2187 }
2188 let mut calls = 0u64;
2189 for handle in handles {
2190 calls += handle.join().map_err(|_| BucketError::WorkerPanic)??;
2191 }
2192 Ok::<_, BucketError>(calls)
2193 })?;
2194 self.record_flush_stats(
2195 SharedBucketFlushStats { calls: flush_calls },
2196 started.elapsed(),
2197 );
2198 Ok(())
2199 }
2200
2201 pub fn finish(&self) -> Result<BucketEmitStats, BucketError> {
2202 let mut entries = Vec::new();
2203 let mut finish_flush_stats = SharedBucketFlushStats::default();
2204 let finish_started = Instant::now();
2205 let mut bytes_written = 0u64;
2206 for atlas in &self.atlases {
2207 let mut atlas = atlas.lock().map_err(|_| BucketError::WorkerPanic)?;
2208 atlas.flush_all(
2209 &self.containers,
2210 self.colored,
2211 self.label_words,
2212 self.compress_buckets,
2213 &mut finish_flush_stats,
2214 )?;
2215 let first_graph_id = atlas.first_graph_id;
2216 let container = (first_graph_id / ATLAS_GRAPH_COUNT) % self.containers.len();
2217 for (local_graph_id, meta) in atlas.files.iter_mut().enumerate() {
2218 let meta = std::mem::take(meta);
2219 if meta.total_records == 0 {
2220 continue;
2221 }
2222 bytes_written += meta.bytes_written;
2223 entries.push(BucketManifestEntry {
2224 graph_id: first_graph_id + local_graph_id,
2225 records: meta.total_records,
2226 location: BucketLocation::Container {
2227 container,
2228 segments: meta.segments,
2229 bytes: meta.bytes_written,
2230 },
2231 });
2232 }
2233 }
2234 self.record_flush_stats(finish_flush_stats, finish_started.elapsed());
2235 entries.sort_unstable_by_key(|entry| entry.graph_id);
2236
2237 let header = ContainerManifestHeader {
2243 k: self.k,
2244 minimizer_len: self.minimizer_len,
2245 graph_count: self.graph_count,
2246 colored: self.colored,
2247 label_words: self.label_words,
2248 compressed: self.compress_buckets,
2249 interleaved_compression: self.compress_buckets
2250 && !self.colored
2251 && !force_split_compression(),
2252 segment_bytes: self.containers.segment_bytes(),
2253 container_count: self.containers.len(),
2256 };
2257 write_container_manifest(&self.bucket_dir, &header, &entries)?;
2258 Ok(BucketEmitStats {
2259 bucket_dir: self.bucket_dir.clone(),
2260 bucket_files: entries.len(),
2261 bytes_written,
2262 })
2263 }
2264
2265 fn record_flush_stats(&self, stats: SharedBucketFlushStats, elapsed: Duration) {
2266 if stats.calls == 0 {
2267 return;
2268 }
2269 self.flush_calls.fetch_add(stats.calls, Ordering::Relaxed);
2270 self.flush_nanos.fetch_add(
2271 u64::try_from(elapsed.as_nanos()).unwrap_or(u64::MAX),
2272 Ordering::Relaxed,
2273 );
2274 }
2275}
2276
2277fn append_colored_atlas_record(
2278 atlas: &mut SharedBucketAtlas,
2279 graph_id: u16,
2280 record: &[u8],
2281) -> Result<(), BucketError> {
2282 let graph_id = usize::from(graph_id);
2283 let local_graph_id = graph_id
2284 .checked_sub(atlas.first_graph_id)
2285 .ok_or(BucketError::InvalidGraphId(graph_id))?;
2286 let file = atlas
2287 .files
2288 .get_mut(local_graph_id)
2289 .ok_or(BucketError::InvalidGraphId(graph_id))?;
2290 file.total_records = file
2291 .total_records
2292 .checked_add(1)
2293 .ok_or(BucketError::TooManyRecords)?;
2294 file.buffer_records = file
2295 .buffer_records
2296 .checked_add(1)
2297 .ok_or(BucketError::TooManyRecords)?;
2298 file.buffer.extend_from_slice(record);
2299 Ok(())
2300}
2301
2302fn sort_colored_payload_by_source(
2303 payload: &mut Vec<u8>,
2304 record_len: usize,
2305 source_min: u32,
2306 source_max: u32,
2307) -> Result<(), BucketError> {
2308 if payload.is_empty() {
2309 return Ok(());
2310 }
2311 if payload.len() % record_len != 0 {
2312 return Err(BucketError::MalformedRecord);
2313 }
2314 let source_count =
2315 usize::try_from(source_max - source_min + 1).map_err(|_| BucketError::MalformedRecord)?;
2316 let mut offsets = vec![0usize; source_count];
2317 for record in payload.chunks_exact(record_len) {
2318 let attr = u32::from_le_bytes(record[..4].try_into().expect("colored attribute"));
2319 let source = attr >> 10;
2320 if source < source_min || source > source_max {
2321 return Err(BucketError::MalformedRecord);
2322 }
2323 offsets[(source - source_min) as usize] += 1;
2324 }
2325 let mut prefix = 0usize;
2326 for offset in &mut offsets {
2327 let count = *offset;
2328 *offset = prefix;
2329 prefix += count;
2330 }
2331 let mut sorted = vec![0u8; payload.len()];
2332 for record in payload.chunks_exact(record_len) {
2333 let attr = u32::from_le_bytes(record[..4].try_into().expect("colored attribute"));
2334 let source = ((attr >> 10) - source_min) as usize;
2335 let output = offsets[source] * record_len;
2336 sorted[output..output + record_len].copy_from_slice(record);
2337 offsets[source] += 1;
2338 }
2339 *payload = sorted;
2340 Ok(())
2341}
2342
2343impl SharedBucketAtlas {
2344 fn append_bucket(
2345 &mut self,
2346 local_graph_id: usize,
2347 pending: PendingBucket,
2348 ) -> Result<(), BucketError> {
2349 let file = &mut self.files[local_graph_id];
2350 file.total_records = file
2351 .total_records
2352 .checked_add(pending.records)
2353 .ok_or(BucketError::TooManyRecords)?;
2354 file.buffer_records = file
2355 .buffer_records
2356 .checked_add(pending.records)
2357 .ok_or(BucketError::TooManyRecords)?;
2358 self.buffered_bytes += pending.bytes.len();
2359 file.buffer.extend_from_slice(&pending.bytes);
2360 Ok(())
2361 }
2362
2363 fn flush_all(
2364 &mut self,
2365 containers: &BucketContainers,
2366 colored: bool,
2367 label_words: usize,
2368 compress_buckets: bool,
2369 stats: &mut SharedBucketFlushStats,
2370 ) -> Result<(), BucketError> {
2371 for local_graph_id in 0..self.files.len() {
2372 self.flush_subgraph(
2373 local_graph_id,
2374 containers,
2375 colored,
2376 label_words,
2377 compress_buckets,
2378 stats,
2379 )?;
2380 }
2381 Ok(())
2382 }
2383
2384 fn flush_subgraph(
2394 &mut self,
2395 local_graph_id: usize,
2396 containers: &BucketContainers,
2397 colored: bool,
2398 label_words: usize,
2399 compress_buckets: bool,
2400 stats: &mut SharedBucketFlushStats,
2401 ) -> Result<(), BucketError> {
2402 let container = (self.first_graph_id / ATLAS_GRAPH_COUNT) % containers.len();
2403 let Self { files, scratch, .. } = self;
2404 let file = &mut files[local_graph_id];
2405 if file.buffer.is_empty() {
2406 return Ok(());
2407 }
2408 let flushed_bytes = file.buffer.len();
2409 let record_size = record_size(colored, label_words) as u64;
2410
2411 let written = if compress_buckets {
2412 let interleaved = !colored && !force_split_compression();
2413 let len = encode_compressed_block(
2414 &file.buffer,
2415 file.buffer_records,
2416 record_size,
2417 label_words,
2418 interleaved,
2419 scratch,
2420 )?;
2421 let block = std::mem::take(&mut scratch.block);
2422 let written = append_to_chain(containers, container, file, &block)?;
2423 scratch.block = block;
2424 debug_assert_eq!(written, len as u64);
2425 written
2426 } else {
2427 let buffer = std::mem::take(&mut file.buffer);
2428 let written = append_to_chain(containers, container, file, &buffer);
2429 file.buffer = buffer;
2430 written?
2431 };
2432
2433 file.written_records += file.buffer_records;
2434 file.bytes_written += written;
2435 file.buffer.clear();
2436 file.buffer_records = 0;
2437 self.buffered_bytes -= flushed_bytes;
2438 stats.calls += 1;
2439 Ok(())
2440 }
2441}
2442
2443fn append_to_chain(
2451 containers: &BucketContainers,
2452 container: usize,
2453 file: &mut SharedBucketFileMeta,
2454 bytes: &[u8],
2455) -> Result<u64, BucketError> {
2456 let segment_bytes = containers.segment_bytes();
2457 let mut written = 0usize;
2458 while written < bytes.len() {
2459 if file.segments.is_empty() || file.segment_used == segment_bytes {
2460 let offset = containers.reserve_segment(container);
2461 let index =
2462 u32::try_from(offset / segment_bytes).map_err(|_| BucketError::TooManyRecords)?;
2463 file.segments.push(index);
2464 file.segment_used = 0;
2465 }
2466 let room = (segment_bytes - file.segment_used) as usize;
2467 let take = room.min(bytes.len() - written);
2468 let offset = u64::from(*file.segments.last().unwrap()) * segment_bytes + file.segment_used;
2469 containers.write_at(container, offset, &bytes[written..written + take])?;
2470 written += take;
2471 file.segment_used += take as u64;
2472 }
2473 Ok(written as u64)
2474}
2475
2476impl SharedBucketEmitter {
2477 pub fn flush_colored_worker_if_required(&mut self) -> Result<(), BucketError> {
2478 if !self.sink.colored {
2479 return Err(BucketError::MalformedRecord);
2480 }
2481 for atlas_id in 0..self.colored_pending.len() {
2482 if self.colored_pending[atlas_id].bytes.len() >= SUBGRAPH_CHUNK_BYTES {
2483 self.flush_pending_colored_atlas(atlas_id)?;
2484 }
2485 }
2486 Ok(())
2487 }
2488
2489 pub fn add(&mut self, superkmer: &WeakSuperKmer, seq: &[u8]) -> Result<(), BucketError> {
2490 self.add_impl(superkmer, seq, true)
2491 }
2492
2493 pub fn add_valid(&mut self, superkmer: &WeakSuperKmer, seq: &[u8]) -> Result<(), BucketError> {
2494 debug_assert!(seq.iter().all(|&base| ascii_base_bits(base).is_some()));
2495 self.add_impl(superkmer, seq, false)
2496 }
2497
2498 fn add_impl(
2499 &mut self,
2500 superkmer: &WeakSuperKmer,
2501 seq: &[u8],
2502 check_bases: bool,
2503 ) -> Result<(), BucketError> {
2504 if superkmer.graph_id >= self.sink.graph_count {
2505 return Err(BucketError::InvalidGraphId(superkmer.graph_id));
2506 }
2507 if seq.len() > u8::MAX as usize {
2508 return Err(BucketError::LabelTooLong(seq.len()));
2509 }
2510
2511 let attr = if self.sink.colored {
2512 let source_id = superkmer.source_id.ok_or(BucketError::MissingSourceId)?;
2513 if source_id > MAX_SOURCE_ID {
2514 return Err(BucketError::SourceIdTooLarge(source_id));
2515 }
2516 pack_colored_attr(
2517 seq.len(),
2518 source_id,
2519 superkmer.left_discontinuous,
2520 superkmer.right_discontinuous,
2521 )
2522 } else {
2523 pack_uncolored_attr(
2524 seq.len(),
2525 superkmer.left_discontinuous,
2526 superkmer.right_discontinuous,
2527 )
2528 };
2529 let graph_id = superkmer.graph_id;
2530 let record_len = record_size(self.sink.colored, self.sink.label_words);
2531 if self.sink.colored {
2532 let atlas_id = graph_id / ATLAS_GRAPH_COUNT;
2533 let pending = &mut self.colored_pending[atlas_id];
2534 pending.graph_ids.push(graph_id as u16);
2535 if check_bases {
2536 append_record(
2537 &mut pending.bytes,
2538 attr,
2539 graph_id,
2540 seq,
2541 self.sink.label_words,
2542 true,
2543 )?;
2544 } else {
2545 append_record_valid(
2546 &mut pending.bytes,
2547 attr,
2548 graph_id,
2549 seq,
2550 self.sink.label_words,
2551 true,
2552 )?;
2553 }
2554 } else if self.deferred_uncolored {
2555 let atlas_id = graph_id / ATLAS_GRAPH_COUNT;
2556 let pending = &mut self.uncolored_pending[atlas_id];
2557 pending.graph_ids.push(graph_id as u16);
2558 if !check_bases {
2559 append_uncolored_record_valid(
2560 &mut pending.bytes,
2561 attr as u16,
2562 graph_id as u16,
2563 seq,
2564 self.sink.label_words,
2565 )?;
2566 } else {
2567 append_record(
2568 &mut pending.bytes,
2569 attr,
2570 graph_id,
2571 seq,
2572 self.sink.label_words,
2573 false,
2574 )?;
2575 }
2576 } else {
2577 let pending = &mut self.pending[graph_id];
2578 pending.records = pending
2579 .records
2580 .checked_add(1)
2581 .ok_or(BucketError::TooManyRecords)?;
2582 if !check_bases {
2583 append_uncolored_record_valid(
2584 &mut pending.bytes,
2585 attr as u16,
2586 graph_id as u16,
2587 seq,
2588 self.sink.label_words,
2589 )?;
2590 } else {
2591 append_record(
2592 &mut pending.bytes,
2593 attr,
2594 graph_id,
2595 seq,
2596 self.sink.label_words,
2597 false,
2598 )?;
2599 }
2600 }
2601 self.pending_bytes += record_len;
2602
2603 if self.deferred_uncolored {
2606 let atlas_id = graph_id / ATLAS_GRAPH_COUNT;
2607 if self.uncolored_pending[atlas_id].bytes.len() >= SUBGRAPH_CHUNK_BYTES {
2608 self.flush_pending_uncolored_atlas(atlas_id)?;
2609 }
2610 } else if !self.sink.colored
2611 && !self.deferred_uncolored
2612 && self.pending[graph_id].bytes.len() >= MAX_PENDING_BUCKET_BYTES
2613 {
2614 self.flush_pending_bucket(graph_id)?;
2615 } else if !self.sink.colored
2616 && !self.deferred_uncolored
2617 && self.pending_bytes >= MAX_TOTAL_PENDING_BYTES
2618 {
2619 self.flush_largest_pending_bucket()?;
2620 }
2621 Ok(())
2622 }
2623
2624 fn flush_largest_pending_bucket(&mut self) -> Result<(), BucketError> {
2625 if self.sink.colored {
2626 let Some((atlas_id, _)) = self
2627 .colored_pending
2628 .iter()
2629 .enumerate()
2630 .max_by_key(|(_, pending)| pending.bytes.len())
2631 else {
2632 return Ok(());
2633 };
2634 return self.flush_pending_colored_atlas(atlas_id);
2635 }
2636 let Some((graph_id, _)) = self
2637 .pending
2638 .iter()
2639 .enumerate()
2640 .max_by_key(|(_, pending)| pending.bytes.len())
2641 else {
2642 return Ok(());
2643 };
2644 self.flush_pending_bucket(graph_id)
2645 }
2646
2647 fn flush_pending_bucket(&mut self, graph_id: usize) -> Result<(), BucketError> {
2648 if self.pending[graph_id].bytes.is_empty() {
2649 return Ok(());
2650 }
2651 let pending = std::mem::take(&mut self.pending[graph_id]);
2652 self.pending_bytes -= pending.bytes.len();
2653 self.sink.append_bucket(graph_id, pending)
2654 }
2655
2656 fn flush_pending_colored_atlas(&mut self, atlas_id: usize) -> Result<(), BucketError> {
2657 if self.colored_pending[atlas_id].bytes.is_empty() {
2658 return Ok(());
2659 }
2660 let pending = std::mem::take(&mut self.colored_pending[atlas_id]);
2661 self.pending_bytes -= pending.bytes.len();
2662 self.sink.append_colored_atlas(atlas_id, pending)
2663 }
2664
2665 fn flush_pending_uncolored_atlas(&mut self, atlas_id: usize) -> Result<(), BucketError> {
2666 if self.uncolored_pending[atlas_id].bytes.is_empty() {
2667 return Ok(());
2668 }
2669 let pending = std::mem::take(&mut self.uncolored_pending[atlas_id]);
2670 self.pending_bytes -= pending.bytes.len();
2671 self.sink.append_uncolored_atlas(atlas_id, pending)
2672 }
2673
2674 pub fn finish(mut self) -> Result<(), BucketError> {
2675 if self.sink.colored {
2676 for atlas_id in 0..self.colored_pending.len() {
2677 if !self.colored_pending[atlas_id].bytes.is_empty() {
2678 self.flush_pending_colored_atlas(atlas_id)?;
2679 }
2680 }
2681 } else if self.deferred_uncolored {
2682 for atlas_id in 0..self.uncolored_pending.len() {
2683 if !self.uncolored_pending[atlas_id].bytes.is_empty() {
2684 self.flush_pending_uncolored_atlas(atlas_id)?;
2685 }
2686 }
2687 } else {
2688 for graph_id in 0..self.pending.len() {
2689 if !self.pending[graph_id].bytes.is_empty() {
2690 self.flush_pending_bucket(graph_id)?;
2691 }
2692 }
2693 }
2694 Ok(())
2695 }
2696}
2697
2698pub fn bucket_dir(params: &BuildParams) -> PathBuf {
2699 let output_name = Path::new(¶ms.output_prefix)
2700 .file_name()
2701 .and_then(|s| s.to_str())
2702 .filter(|s| !s.is_empty())
2703 .unwrap_or("cuttlefish3");
2704 PathBuf::from(¶ms.work_dir).join(format!("{output_name}.cf3rs.wsk"))
2705}
2706
2707pub fn write_manifest(
2708 bucket_dir: &Path,
2709 entries: &[(usize, u64, PathBuf)],
2710) -> Result<(), BucketError> {
2711 let path = bucket_dir.join("manifest.tsv");
2712 let mut out = File::create(&path).map_err(|source| BucketError::Io {
2713 path: path.clone(),
2714 source,
2715 })?;
2716 writeln!(out, "graph_id\trecords\tpath").map_err(|source| BucketError::Io {
2717 path: path.clone(),
2718 source,
2719 })?;
2720 for (graph_id, records, bucket_path) in entries {
2721 writeln!(out, "{graph_id}\t{records}\t{}", bucket_path.display()).map_err(|source| {
2722 BucketError::Io {
2723 path: path.clone(),
2724 source,
2725 }
2726 })?;
2727 }
2728 Ok(())
2729}
2730
2731fn encode_compressed_block(
2737 bytes: &[u8],
2738 records: u64,
2739 record_size: u64,
2740 label_words: usize,
2741 interleaved: bool,
2742 scratch: &mut CompressionScratch,
2743) -> Result<usize, BucketError> {
2744 let records_u32 = u32::try_from(records).map_err(|_| BucketError::TooManyRecords)?;
2745 if records_u32 == 0 {
2746 scratch.block.clear();
2747 return Ok(0);
2748 }
2749 let (attr_len, label_len) = if interleaved {
2750 (CompressionScratch::encode(bytes, &mut scratch.encoded)?, 0)
2751 } else {
2752 let record_len = usize::try_from(record_size).unwrap();
2753 let fixed_len = record_len - label_words * 8;
2754 scratch.attrs.clear();
2755 scratch.labels.clear();
2756 scratch.attrs.reserve(records as usize * fixed_len);
2757 scratch.labels.reserve(records as usize * label_words * 8);
2758 for record in bytes.chunks_exact(record_len) {
2759 scratch.attrs.extend_from_slice(&record[..fixed_len]);
2760 scratch.labels.extend_from_slice(&record[fixed_len..]);
2761 }
2762 let attrs = std::mem::take(&mut scratch.attrs);
2763 let labels = std::mem::take(&mut scratch.labels);
2764 let attr_len = CompressionScratch::encode(&attrs, &mut scratch.encoded)?;
2765 let label_len = CompressionScratch::encode(&labels, &mut scratch.encoded_labels)?;
2766 scratch.attrs = attrs;
2767 scratch.labels = labels;
2768 (attr_len, label_len)
2769 };
2770 let attr_len_u32 = u32::try_from(attr_len).map_err(|_| BucketError::TooManyRecords)?;
2771 let label_len_u32 = u32::try_from(label_len).map_err(|_| BucketError::TooManyRecords)?;
2772
2773 scratch.block.clear();
2774 scratch.block.extend_from_slice(&records_u32.to_le_bytes());
2775 scratch.block.extend_from_slice(&attr_len_u32.to_le_bytes());
2776 scratch
2777 .block
2778 .extend_from_slice(&label_len_u32.to_le_bytes());
2779 scratch
2780 .block
2781 .extend_from_slice(&scratch.encoded[..attr_len]);
2782 if label_len != 0 {
2783 scratch
2784 .block
2785 .extend_from_slice(&scratch.encoded_labels[..label_len]);
2786 }
2787 Ok(scratch.block.len())
2788}
2789
2790fn append_record(
2791 out: &mut Vec<u8>,
2792 packed_attr: u32,
2793 _graph_id: usize,
2794 seq: &[u8],
2795 label_words: usize,
2796 colored: bool,
2797) -> Result<(), BucketError> {
2798 let mut words = [0u64; 4];
2799 if label_words > words.len() {
2800 return Err(BucketError::MalformedRecord);
2801 }
2802 for (idx, &ch) in seq.iter().enumerate() {
2803 let base_bits = ascii_base_bits(ch).ok_or(BucketError::InvalidBase(ch))?;
2804 let word_idx = idx / 32;
2805 let shift = 2 * (31 - (idx % 32));
2806 words[word_idx] |= (base_bits as u64) << shift;
2807 }
2808
2809 if colored {
2810 out.extend_from_slice(&packed_attr.to_le_bytes());
2811 } else {
2812 out.extend_from_slice(&(packed_attr as u16).to_le_bytes());
2813 }
2814 for &word in &words[..label_words] {
2815 out.extend_from_slice(&word.to_le_bytes());
2816 }
2817 Ok(())
2818}
2819
2820fn append_record_valid(
2821 out: &mut Vec<u8>,
2822 packed_attr: u32,
2823 _graph_id: usize,
2824 seq: &[u8],
2825 label_words: usize,
2826 colored: bool,
2827) -> Result<(), BucketError> {
2828 let words = pack_valid_label(seq, label_words)?;
2829
2830 let mut record = [0u8; MAX_RECORD_BYTES];
2831 let attr_len = if colored {
2832 record[..4].copy_from_slice(&packed_attr.to_le_bytes());
2833 4
2834 } else {
2835 record[..2].copy_from_slice(&(packed_attr as u16).to_le_bytes());
2836 2
2837 };
2838 for (idx, &word) in words[..label_words].iter().enumerate() {
2839 let at = attr_len + idx * 8;
2840 record[at..at + 8].copy_from_slice(&word.to_le_bytes());
2841 }
2842 out.extend_from_slice(&record[..attr_len + label_words * 8]);
2843 Ok(())
2844}
2845
2846#[inline]
2847fn append_uncolored_record_valid(
2848 out: &mut Vec<u8>,
2849 packed_attr: u16,
2850 _graph_id: u16,
2851 seq: &[u8],
2852 label_words: usize,
2853) -> Result<(), BucketError> {
2854 let words = pack_valid_label(seq, label_words)?;
2855
2856 let mut record = [0u8; MAX_RECORD_BYTES];
2861 record[..2].copy_from_slice(&packed_attr.to_le_bytes());
2862 for (idx, &word) in words[..label_words].iter().enumerate() {
2863 let at = 2 + idx * 8;
2864 record[at..at + 8].copy_from_slice(&word.to_le_bytes());
2865 }
2866 out.extend_from_slice(&record[..2 + label_words * 8]);
2867 Ok(())
2868}
2869
2870#[inline(always)]
2871fn pack_valid_word_32(seq: &[u8]) -> u64 {
2879 debug_assert!(seq.len() >= 32);
2880 #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))]
2881 {
2882 let mut word = 0u64;
2883 for chunk in 0..4 {
2884 let bytes = u64::from_le_bytes(
2885 seq[chunk * 8..chunk * 8 + 8]
2886 .try_into()
2887 .expect("eight bases"),
2888 );
2889 let codes = ((bytes >> 2) ^ (bytes >> 1)) & 0x0303_0303_0303_0303;
2891 let packed = unsafe { core::arch::x86_64::_pext_u64(codes, 0x0303_0303_0303_0303) };
2893 let reversed = reverse_2bit_groups_16(packed as u16);
2897 word = (word << 16) | u64::from(reversed);
2898 }
2899 word
2900 }
2901 #[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))]
2902 {
2903 let mut word = 0u64;
2904 for &base in &seq[..32] {
2905 word = (word << 2) | u64::from(valid_ascii_base_bits(base));
2906 }
2907 word
2908 }
2909}
2910
2911#[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))]
2913#[inline]
2914fn reverse_2bit_groups_16(value: u16) -> u16 {
2915 let v = value as u32;
2916 let v = ((v & 0x3333) << 2) | ((v >> 2) & 0x3333);
2918 let v = ((v & 0x0f0f) << 4) | ((v >> 4) & 0x0f0f);
2919 (((v & 0x00ff) << 8) | ((v >> 8) & 0x00ff)) as u16
2920}
2921
2922#[inline]
2923fn pack_valid_label(seq: &[u8], label_words: usize) -> Result<[u64; 4], BucketError> {
2924 let mut words = [0u64; 4];
2925 if label_words > words.len() || seq.len() > label_words * 32 {
2926 return Err(BucketError::MalformedRecord);
2927 }
2928
2929 let full_words = seq.len() / 32;
2930 for word_idx in 0..full_words {
2931 words[word_idx] = pack_valid_word_32(&seq[word_idx * 32..]);
2932 }
2933 let tail = &seq[full_words * 32..];
2934 if !tail.is_empty() {
2935 let mut word = 0u64;
2936 for &base in tail {
2937 word = (word << 2) | u64::from(valid_ascii_base_bits(base));
2938 }
2939 words[full_words] = word << (2 * (32 - tail.len()));
2940 }
2941 Ok(words)
2942}
2943
2944struct BucketFile {
2945 path: PathBuf,
2946 file: File,
2947 records: u64,
2948 bytes_written: u64,
2949 record_size: u64,
2950 compressed: bool,
2951 interleaved_compression: bool,
2952 label_words: usize,
2953}
2954
2955#[derive(Default)]
2960pub(crate) struct CompressionScratch {
2961 attrs: Vec<u8>,
2963 labels: Vec<u8>,
2965 encoded: Vec<u8>,
2967 encoded_labels: Vec<u8>,
2969 block: Vec<u8>,
2971}
2972
2973impl CompressionScratch {
2974 fn encode(input: &[u8], out: &mut Vec<u8>) -> Result<usize, BucketError> {
2979 let bound = lz4_flex::block::get_maximum_output_size(input.len());
2980 if out.len() < bound {
2981 out.resize(bound, 0);
2982 }
2983 lz4_flex::block::compress_into(input, &mut out[..bound])
2984 .map_err(|_| BucketError::MalformedRecord)
2985 }
2986}
2987
2988fn force_split_compression() -> bool {
2994 static SPLIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2995 *SPLIT.get_or_init(|| std::env::var_os("CF3_RS_SPLIT_COMPRESSION").is_some())
2996}
2997
2998impl BucketFile {
2999 #[allow(clippy::too_many_arguments)]
3000 fn create(
3001 bucket_dir: &Path,
3002 k: u16,
3003 minimizer_len: u16,
3004 graph_count: usize,
3005 graph_id: usize,
3006 colored: bool,
3007 label_words: usize,
3008 compressed: bool,
3009 ) -> Result<Self, BucketError> {
3010 let path = bucket_dir.join(format!("{graph_id:05}.wsk"));
3011 let mut file = File::create(&path).map_err(|source| BucketError::Io {
3012 path: path.clone(),
3013 source,
3014 })?;
3015
3016 file.write_all(MAGIC).map_err(|source| BucketError::Io {
3017 path: path.clone(),
3018 source,
3019 })?;
3020 write_u16(&mut file, &path, k)?;
3021 write_u16(&mut file, &path, minimizer_len)?;
3022 write_u64(&mut file, &path, graph_count as u64)?;
3023 write_u64(&mut file, &path, graph_id as u64)?;
3024 let interleaved_compression = compressed && !colored && !force_split_compression();
3025 file.write_all(&[
3026 u8::from(colored),
3027 label_words as u8,
3028 if interleaved_compression {
3029 2
3030 } else {
3031 u8::from(compressed)
3032 },
3033 0,
3034 0,
3035 0,
3036 ])
3037 .map_err(|source| BucketError::Io {
3038 path: path.clone(),
3039 source,
3040 })?;
3041 write_u64(&mut file, &path, 0)?;
3042
3043 Ok(Self {
3044 path,
3045 file,
3046 records: 0,
3047 bytes_written: HEADER_LEN,
3048 record_size: record_size(colored, label_words) as u64,
3049 compressed,
3050 interleaved_compression,
3051 label_words,
3052 })
3053 }
3054
3055 fn open_existing(path: &Path, records: u64, bytes_written: u64) -> Result<Self, BucketError> {
3056 let mut file = OpenOptions::new()
3057 .read(true)
3058 .write(true)
3059 .open(path)
3060 .map_err(|source| BucketError::Io {
3061 path: path.to_path_buf(),
3062 source,
3063 })?;
3064 let header = read_header(&mut file, path)?;
3065 file.seek(SeekFrom::End(0))
3066 .map_err(|source| BucketError::Io {
3067 path: path.to_path_buf(),
3068 source,
3069 })?;
3070
3071 Ok(Self {
3072 path: path.to_path_buf(),
3073 file,
3074 records,
3075 bytes_written,
3076 record_size: record_size(header.colored, header.label_words) as u64,
3077 compressed: header.compressed,
3078 interleaved_compression: header.interleaved_compression,
3079 label_words: header.label_words,
3080 })
3081 }
3082
3083 fn write_records(
3084 &mut self,
3085 bytes: &[u8],
3086 records: u64,
3087 scratch: &mut CompressionScratch,
3088 ) -> Result<(u64, u64), BucketError> {
3089 debug_assert_eq!(bytes.len() as u64, records * self.record_size);
3090 let written = if self.compressed {
3091 self.write_compressed_block(bytes, records, scratch)?
3092 } else {
3093 self.file
3094 .write_all(bytes)
3095 .map_err(|source| BucketError::Io {
3096 path: self.path.clone(),
3097 source,
3098 })?;
3099 bytes.len() as u64
3100 };
3101 self.records = self
3102 .records
3103 .checked_add(records)
3104 .ok_or(BucketError::TooManyRecords)?;
3105 self.bytes_written += written;
3106 Ok((self.records, self.bytes_written))
3107 }
3108
3109 fn write_compressed_block(
3110 &mut self,
3111 bytes: &[u8],
3112 records: u64,
3113 scratch: &mut CompressionScratch,
3114 ) -> Result<u64, BucketError> {
3115 let len = encode_compressed_block(
3116 bytes,
3117 records,
3118 self.record_size,
3119 self.label_words,
3120 self.interleaved_compression,
3121 scratch,
3122 )?;
3123 if len == 0 {
3124 return Ok(0);
3125 }
3126 self.file
3127 .write_all(&scratch.block)
3128 .map_err(|source| BucketError::Io {
3129 path: self.path.clone(),
3130 source,
3131 })?;
3132 Ok(len as u64)
3133 }
3134
3135 fn flush(&mut self) -> Result<(), BucketError> {
3136 self.file.flush().map_err(|source| BucketError::Io {
3137 path: self.path.clone(),
3138 source,
3139 })
3140 }
3141
3142 fn finish(&mut self) -> Result<(), BucketError> {
3143 write_record_count(&mut self.file, &self.path, self.records)?;
3144 self.file
3145 .seek(SeekFrom::End(0))
3146 .map_err(|source| BucketError::Io {
3147 path: self.path.clone(),
3148 source,
3149 })?;
3150 self.file.flush().map_err(|source| BucketError::Io {
3151 path: self.path.clone(),
3152 source,
3153 })
3154 }
3155
3156 fn finish_closed(path: &Path, records: u64) -> Result<(), BucketError> {
3157 let mut file = OpenOptions::new()
3158 .read(true)
3159 .write(true)
3160 .open(path)
3161 .map_err(|source| BucketError::Io {
3162 path: path.to_path_buf(),
3163 source,
3164 })?;
3165 write_record_count(&mut file, path, records)?;
3166 file.flush().map_err(|source| BucketError::Io {
3167 path: path.to_path_buf(),
3168 source,
3169 })
3170 }
3171}
3172
3173fn write_record_count(file: &mut File, path: &Path, records: u64) -> Result<(), BucketError> {
3174 file.seek(SeekFrom::Start(RECORD_COUNT_OFFSET))
3175 .map_err(|source| BucketError::Io {
3176 path: path.to_path_buf(),
3177 source,
3178 })?;
3179 write_u64(file, path, records)
3180}
3181
3182fn label_word_count(k: u16, minimizer_len: u16) -> usize {
3183 let max_weak_superkmer_len = 2 * (usize::from(k) - 1) - usize::from(minimizer_len) + 2;
3184 max_weak_superkmer_len.div_ceil(32)
3185}
3186
3187fn record_size(colored: bool, label_words: usize) -> usize {
3188 (if colored { 4 } else { 2 }) + label_words * 8
3189}
3190
3191fn pack_uncolored_attr(len: usize, left_discontinuous: bool, right_discontinuous: bool) -> u32 {
3192 (len as u32) | ((left_discontinuous as u32) << 8) | ((right_discontinuous as u32) << 9)
3193}
3194
3195fn pack_colored_attr(
3196 len: usize,
3197 source_id: u32,
3198 left_discontinuous: bool,
3199 right_discontinuous: bool,
3200) -> u32 {
3201 pack_uncolored_attr(len, left_discontinuous, right_discontinuous) | (source_id << 10)
3202}
3203
3204fn decode_label_into(words: &[u64], len: usize, seq: &mut Vec<u8>) -> Result<(), BucketError> {
3205 if len > words.len() * 32 {
3206 return Err(BucketError::MalformedRecord);
3207 }
3208
3209 seq.clear();
3210 seq.reserve(len);
3211 for idx in 0..len {
3212 let word_idx = idx / 32;
3213 let shift = 2 * (31 - (idx % 32));
3214 let base = match ((words[word_idx] >> shift) & 0b11) as u8 {
3215 0 => Base::A,
3216 1 => Base::C,
3217 2 => Base::G,
3218 3 => Base::T,
3219 _ => unreachable!(),
3220 };
3221 seq.push(base.to_ascii());
3222 }
3223 Ok(())
3224}
3225
3226fn read_header(file: &mut impl Read, path: &Path) -> Result<BucketHeader, BucketError> {
3227 let mut magic = [0u8; 8];
3228 file.read_exact(&mut magic)
3229 .map_err(|source| BucketError::Io {
3230 path: path.to_path_buf(),
3231 source,
3232 })?;
3233 if &magic != MAGIC {
3234 return Err(BucketError::BadMagic(path.to_path_buf()));
3235 }
3236
3237 let k = read_u16(file, path)?;
3238 let minimizer_len = read_u16(file, path)?;
3239 let graph_count = usize::try_from(read_u64(file, path)?)
3240 .map_err(|_| BucketError::MalformedHeader(path.to_path_buf()))?;
3241 let graph_id = usize::try_from(read_u64(file, path)?)
3242 .map_err(|_| BucketError::MalformedHeader(path.to_path_buf()))?;
3243
3244 let mut flags = [0u8; 6];
3245 file.read_exact(&mut flags)
3246 .map_err(|source| BucketError::Io {
3247 path: path.to_path_buf(),
3248 source,
3249 })?;
3250 let colored = match flags[0] {
3251 0 => false,
3252 1 => true,
3253 _ => return Err(BucketError::MalformedHeader(path.to_path_buf())),
3254 };
3255 let label_words = flags[1] as usize;
3256 if label_words == 0 || label_words != label_word_count(k, minimizer_len) {
3257 return Err(BucketError::MalformedHeader(path.to_path_buf()));
3258 }
3259 let (compressed, interleaved_compression) = match flags[2] {
3260 0 => (false, false),
3261 1 => (true, false),
3262 2 if !colored => (true, true),
3263 _ => return Err(BucketError::MalformedHeader(path.to_path_buf())),
3264 };
3265 if flags[3..].iter().any(|&b| b != 0) {
3266 return Err(BucketError::MalformedHeader(path.to_path_buf()));
3267 }
3268 let records = read_u64(file, path)?;
3269
3270 if graph_id >= graph_count {
3271 return Err(BucketError::MalformedHeader(path.to_path_buf()));
3272 }
3273
3274 Ok(BucketHeader {
3275 k,
3276 minimizer_len,
3277 graph_count,
3278 graph_id,
3279 colored,
3280 compressed,
3281 interleaved_compression,
3282 label_words,
3283 records,
3284 })
3285}
3286
3287fn read_u16(file: &mut impl Read, path: &Path) -> Result<u16, BucketError> {
3288 let mut bytes = [0u8; 2];
3289 file.read_exact(&mut bytes)
3290 .map_err(|source| BucketError::Io {
3291 path: path.to_path_buf(),
3292 source,
3293 })?;
3294 Ok(u16::from_le_bytes(bytes))
3295}
3296
3297fn read_u64(file: &mut impl Read, path: &Path) -> Result<u64, BucketError> {
3298 let mut bytes = [0u8; 8];
3299 file.read_exact(&mut bytes)
3300 .map_err(|source| BucketError::Io {
3301 path: path.to_path_buf(),
3302 source,
3303 })?;
3304 Ok(u64::from_le_bytes(bytes))
3305}
3306
3307fn write_u16(file: &mut File, path: &Path, value: u16) -> Result<(), BucketError> {
3308 file.write_all(&value.to_le_bytes())
3309 .map_err(|source| BucketError::Io {
3310 path: path.to_path_buf(),
3311 source,
3312 })
3313}
3314
3315fn write_u64(file: &mut File, path: &Path, value: u64) -> Result<(), BucketError> {
3316 file.write_all(&value.to_le_bytes())
3317 .map_err(|source| BucketError::Io {
3318 path: path.to_path_buf(),
3319 source,
3320 })
3321}
3322
3323#[derive(Debug)]
3324pub enum BucketError {
3325 Io {
3326 path: PathBuf,
3327 source: std::io::Error,
3328 },
3329 GraphCountTooLarge(usize),
3330 InvalidGraphId(usize),
3331 LabelTooLong(usize),
3332 MissingSourceId,
3333 SourceIdTooLarge(u32),
3334 InvalidBase(u8),
3335 TooManyRecords,
3336 BadMagic(PathBuf),
3337 MalformedHeader(PathBuf),
3338 MalformedManifest(PathBuf),
3339 MalformedRecord,
3340 RecordGraphMismatch {
3341 expected: usize,
3342 got: usize,
3343 },
3344 WorkerPanic,
3345}
3346
3347impl std::fmt::Display for BucketError {
3348 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3349 match self {
3350 Self::Io { path, source } => write!(f, "{}: {source}", path.display()),
3351 Self::GraphCountTooLarge(count) => write!(f, "graph count is too large: {count}"),
3352 Self::InvalidGraphId(graph_id) => write!(f, "invalid graph id: {graph_id}"),
3353 Self::LabelTooLong(len) => write!(f, "weak super-kmer label is too long: {len}"),
3354 Self::MissingSourceId => write!(f, "colored bucket record is missing source id"),
3355 Self::SourceIdTooLarge(source_id) => {
3356 write!(
3357 f,
3358 "source id exceeds 21-bit colored bucket limit: {source_id}"
3359 )
3360 }
3361 Self::InvalidBase(b) => write!(f, "invalid base in bucket label: '{}'", *b as char),
3362 Self::TooManyRecords => write!(f, "too many weak super-kmer records"),
3363 Self::BadMagic(path) => write!(
3364 f,
3365 "not a CF3 Rust weak-superkmer bucket: {}",
3366 path.display()
3367 ),
3368 Self::MalformedHeader(path) => {
3369 write!(
3370 f,
3371 "malformed weak-superkmer bucket header: {}",
3372 path.display()
3373 )
3374 }
3375 Self::MalformedManifest(path) => {
3376 write!(
3377 f,
3378 "malformed weak-superkmer bucket manifest: {}",
3379 path.display()
3380 )
3381 }
3382 Self::MalformedRecord => write!(f, "malformed weak-superkmer bucket record"),
3383 Self::RecordGraphMismatch { expected, got } => {
3384 write!(
3385 f,
3386 "bucket record graph id mismatch: expected {expected}, got {got}"
3387 )
3388 }
3389 Self::WorkerPanic => write!(f, "bucket worker thread panicked"),
3390 }
3391 }
3392}
3393
3394impl std::error::Error for BucketError {}
3395
3396#[cfg(test)]
3397mod tests {
3398 #[test]
3409 fn releasing_segments_returns_blocks_to_the_filesystem() {
3410 use std::os::unix::fs::MetadataExt;
3411
3412 let dir = std::env::temp_dir().join(format!(
3413 "cf3-punch-{}-{:?}",
3414 std::process::id(),
3415 std::thread::current().id()
3416 ));
3417 let _ = fs::remove_dir_all(&dir);
3418 fs::create_dir_all(&dir).unwrap();
3419
3420 let containers = BucketContainers::create(&dir, 1).unwrap();
3421 let segment = containers.segment_bytes();
3422 let segments: Vec<u32> = (0..8).collect();
3423 let payload = vec![0xA5u8; segment as usize];
3424 for &index in &segments {
3425 containers
3426 .write_at(0, u64::from(index) * segment, &payload)
3427 .unwrap();
3428 }
3429
3430 let path = dir.join("00000.wskc");
3431 let allocated = || fs::metadata(&path).unwrap().blocks() * 512;
3432 let before = allocated();
3433 assert!(
3434 before >= segments.len() as u64 * segment,
3435 "expected {} bytes of blocks before punching, saw {before}",
3436 segments.len() as u64 * segment
3437 );
3438
3439 containers.release_segments(0, &segments);
3440 let after = allocated();
3441
3442 assert!(
3443 after < before / 2,
3444 "punching {} segments freed {} of {before} bytes; this filesystem may \
3445 not support hole punching, in which case consumed bucket space is \
3446 held until the build ends and peak disk is higher than documented",
3447 segments.len(),
3448 before - after
3449 );
3450 assert_eq!(
3452 fs::metadata(&path).unwrap().len(),
3453 segments.len() as u64 * segment
3454 );
3455
3456 fs::remove_dir_all(&dir).unwrap();
3457 }
3458
3459 #[test]
3460 fn packed_word_matches_scalar_reference() {
3461 fn scalar(seq: &[u8]) -> u64 {
3462 let mut word = 0u64;
3463 for &base in &seq[..32] {
3464 word = (word << 2) | u64::from(valid_ascii_base_bits(base));
3465 }
3466 word
3467 }
3468 let alphabet = b"ACGT";
3469 let mut state = 0x1234_5678_9abc_def0u64;
3471 for case in 0..2048 {
3472 let mut seq = [0u8; 32];
3473 for (i, slot) in seq.iter_mut().enumerate() {
3474 state = state
3475 .wrapping_mul(6364136223846793005)
3476 .wrapping_add(1442695040888963407);
3477 *slot = if case < 4 {
3478 alphabet[(case + i) % 4]
3479 } else {
3480 alphabet[(state >> 33) as usize % 4]
3481 };
3482 }
3483 assert_eq!(
3484 pack_valid_word_32(&seq),
3485 scalar(&seq),
3486 "mismatch for {:?}",
3487 std::str::from_utf8(&seq).unwrap()
3488 );
3489 }
3490 }
3491
3492 use super::*;
3493
3494 #[test]
3495 fn deferred_uncolored_atlas_chunks_preserve_graph_buckets() {
3496 let dir = std::env::temp_dir().join(format!(
3497 "cf3-uncolored-atlas-bucket-{}-{:?}",
3498 std::process::id(),
3499 std::thread::current().id()
3500 ));
3501 fs::create_dir_all(&dir).unwrap();
3502 let mut params = BuildParams::new(crate::GraphInput::References, "test".to_string());
3503 params.k = 31;
3504 params.minimizer_len = 15;
3505 params.vertex_partitions = 1;
3506 params.threads = 2;
3507 params.work_dir = dir.to_string_lossy().into_owned();
3508 let sink = SharedBucketSink::create(¶ms, ATLAS_GRAPH_COUNT + 1).unwrap();
3509 let expected = [(0, b'A'), (7, b'C'), (ATLAS_GRAPH_COUNT, b'G')];
3510 let mut emitters = Vec::new();
3511 for repeat in 0..2 {
3512 let mut emitter = sink.deferred_uncolored_emitter();
3513 for &(graph_id, base) in &expected {
3514 for _ in 0..(repeat + 1) {
3515 emitter
3516 .add_valid(
3517 &WeakSuperKmer {
3518 graph_id,
3519 offset: 0,
3520 len: 31,
3521 source_id: None,
3522 left_discontinuous: false,
3523 right_discontinuous: false,
3524 },
3525 &[base; 31],
3526 )
3527 .unwrap();
3528 }
3529 }
3530 emitters.push(emitter);
3531 }
3532 sink.flush_uncolored_emitters(emitters).unwrap();
3533 let stats = sink.finish().unwrap();
3534
3535 let (store, entries) = BucketStore::open_dir(&stats.bucket_dir).unwrap();
3536 for &(graph_id, base) in &expected {
3537 let entry = entries
3538 .iter()
3539 .find(|entry| entry.graph_id == graph_id)
3540 .expect("bucket in manifest");
3541 let mut reader = store.reader(entry).unwrap();
3542 let mut record = BucketRecord::default();
3543 let mut count = 0;
3544 while reader.next_record_into(&mut record).unwrap() {
3545 assert_eq!(record.graph_id, graph_id);
3546 assert_eq!(record.label, vec![base; 31]);
3547 count += 1;
3548 }
3549 assert_eq!(count, 3);
3550 }
3551 fs::remove_dir_all(dir).unwrap();
3552 }
3553
3554 #[test]
3555 fn colored_atlas_windows_preserve_global_source_order() {
3556 let dir = std::env::temp_dir().join(format!(
3557 "cf3-colored-bucket-{}-{:?}",
3558 std::process::id(),
3559 std::thread::current().id()
3560 ));
3561 fs::create_dir_all(&dir).unwrap();
3562 let mut params = BuildParams::new(crate::GraphInput::References, "test".to_string());
3563 params.color = true;
3564 params.k = 31;
3565 params.minimizer_len = 15;
3566 params.vertex_partitions = 1;
3567 params.threads = 1;
3568 params.work_dir = dir.to_string_lossy().into_owned();
3569 let sink = SharedBucketSink::create(¶ms, 1).unwrap();
3570
3571 for (source_min, source_max, sources) in
3572 [(1, 3, vec![3, 1, 3, 2, 1]), (4, 6, vec![6, 4, 5, 4])]
3573 {
3574 let mut emitter = sink.emitter();
3575 for source_id in sources {
3576 emitter
3577 .add_valid(
3578 &WeakSuperKmer {
3579 graph_id: 0,
3580 offset: 0,
3581 len: 31,
3582 source_id: Some(source_id),
3583 left_discontinuous: false,
3584 right_discontinuous: false,
3585 },
3586 &[b'A'; 31],
3587 )
3588 .unwrap();
3589 }
3590 emitter.finish().unwrap();
3591 sink.flush_colored_window(source_min, source_max).unwrap();
3592 }
3593 let stats = sink.finish().unwrap();
3594
3595 let (store, entries) = BucketStore::open_dir(&stats.bucket_dir).unwrap();
3596 let entry = entries
3597 .iter()
3598 .find(|entry| entry.graph_id == 0)
3599 .expect("bucket in manifest");
3600 let mut reader = store.reader(entry).unwrap();
3601 assert!(reader.header().compressed);
3602 let mut sources = Vec::new();
3603 let mut record = BucketPackedRecord::default();
3604 while reader.next_packed_record_into(&mut record).unwrap() {
3605 sources.push(record.source_id.unwrap());
3606 }
3607 assert_eq!(sources, [1, 1, 2, 3, 3, 4, 4, 5, 6]);
3608
3609 let container_path = stats.bucket_dir.join("00000.wskc");
3613 let len = fs::metadata(&container_path).unwrap().len();
3614 OpenOptions::new()
3615 .write(true)
3616 .open(&container_path)
3617 .unwrap()
3618 .set_len(len - 1)
3619 .unwrap();
3620 let (store, entries) = BucketStore::open_dir(&stats.bucket_dir).unwrap();
3621 let entry = entries
3622 .iter()
3623 .find(|entry| entry.graph_id == 0)
3624 .expect("bucket in manifest");
3625 let mut truncated = store.reader(entry).unwrap();
3626 let mut record = BucketPackedRecord::default();
3627 let mut failed = false;
3628 loop {
3629 match truncated.next_packed_record_into(&mut record) {
3630 Ok(true) => {}
3631 Ok(false) => break,
3632 Err(_) => {
3633 failed = true;
3634 break;
3635 }
3636 }
3637 }
3638 assert!(failed, "truncated compressed block must be rejected");
3639 fs::remove_dir_all(dir).unwrap();
3640 }
3641
3642 #[test]
3643 fn colored_worker_tails_preserve_cpp_worker_order() {
3644 let dir = std::env::temp_dir().join(format!(
3645 "cf3-colored-worker-tails-{}-{:?}",
3646 std::process::id(),
3647 std::thread::current().id()
3648 ));
3649 fs::create_dir_all(&dir).unwrap();
3650 let mut params = BuildParams::new(crate::GraphInput::References, "test".to_string());
3651 params.color = true;
3652 params.k = 31;
3653 params.minimizer_len = 15;
3654 params.threads = 2;
3655 params.work_dir = dir.to_string_lossy().into_owned();
3656 let sink = SharedBucketSink::create(¶ms, 1).unwrap();
3657
3658 let mut emitters = Vec::new();
3659 for sources in [[3, 1], [4, 2]] {
3660 let mut emitter = sink.emitter();
3661 for source_id in sources {
3662 emitter
3663 .add_valid(
3664 &WeakSuperKmer {
3665 graph_id: 0,
3666 offset: 0,
3667 len: 31,
3668 source_id: Some(source_id),
3669 left_discontinuous: false,
3670 right_discontinuous: false,
3671 },
3672 &[b'A'; 31],
3673 )
3674 .unwrap();
3675 }
3676 emitters.push(emitter);
3677 }
3678 sink.flush_colored_emitters(emitters).unwrap();
3679 let stats = sink.finish().unwrap();
3680
3681 let (store, entries) = BucketStore::open_dir(&stats.bucket_dir).unwrap();
3682 let entry = entries
3683 .iter()
3684 .find(|entry| entry.graph_id == 0)
3685 .expect("bucket in manifest");
3686 let mut reader = store.reader(entry).unwrap();
3687 let mut sources = Vec::new();
3688 let mut record = BucketPackedRecord::default();
3689 while reader.next_packed_record_into(&mut record).unwrap() {
3690 sources.push(record.source_id.unwrap());
3691 }
3692 assert_eq!(sources, [3, 1, 4, 2]);
3693 fs::remove_dir_all(dir).unwrap();
3694 }
3695}