1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
//! Two-pass directory compression.
//!
//! Pass 1 — BIG files (> slice_size): sequential chunked reads, metadata by re-reading file.
//! Pass 2 — SMALL files (≤ slice_size): read into slot, metadata from in-memory data.
//!
//! Arrow IPC index is written incrementally — each pass writes its batch as soon as it finishes.
//! No accumulation, no merge.
use anyhow::{Result, anyhow};
use crossbeam_channel::{bounded, unbounded};
use std::fs::File;
use std::io::{self, BufReader, Read};
use std::os::unix::fs::FileExt;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::thread;
use walkdir::WalkDir;
use znippy_common::common_config::CONFIG;
use znippy_common::index::{
FileExtMeta, build_arrow_metadata_for_config,
build_metadata_batch, compose_index_schema, should_skip_compression,
};
use znippy_common::meta::{BlobMeta, ChunkMeta};
use znippy_common::slotpool::Magazine;
use znippy_common::CompressionReport;
use znippy_common::{ArchiveMetaSink, ArrowIpcSink, GroupKey};
const SLOT_SIZE: usize = 200 * 1024 * 1024;
const NUM_SLOTS: usize = 8;
/// TEST SEAM: when true, `run_small_pass` skips io_uring init and takes the
/// graceful fallback path even on a kernel where io_uring is available. This
/// lets the inject-assert test force the fallback without making the kernel
/// fail. Production code never sets this; it defaults to `false`.
#[cfg(test)]
static FORCE_SMALL_FALLBACK: std::sync::atomic::AtomicBool =
std::sync::atomic::AtomicBool::new(false);
/// Returns whether the small-file pass should force the io_uring fallback.
/// In non-test builds this is a const `false` and optimises away entirely, so
/// there is no production-path overhead or behaviour change.
#[inline(always)]
fn force_small_fallback() -> bool {
#[cfg(test)]
{
FORCE_SMALL_FALLBACK.load(Ordering::Relaxed)
}
#[cfg(not(test))]
{
false
}
}
/// What the writer pwrites for one chunk.
enum Payload {
/// Compressed output owned by the worker; recycled (dropped) after pwrite.
Buf(Vec<u8>),
/// PERF (Law 1): zero-copy skip/incompressible path. The bytes live in the
/// slot; the writer pwrites straight from the slot via `Round::as_slice()`
/// and releases the slot AFTER the pwrite. No memcpy out of the slot.
Slot(znippy_common::slotpool::Round),
}
struct WriteJob {
payload: Payload,
on_disk_len: usize,
file_index: u64,
fdata_offset: u64,
chunk_seq: u32,
checksum: [u8; 32],
compressed: bool,
uncompressed_size: u64,
}
fn read_fully<R: Read>(r: &mut R, buf: &mut [u8]) -> io::Result<usize> {
let mut n = 0;
while n < buf.len() {
match r.read(&mut buf[n..])? {
0 => break,
k => n += k,
}
}
Ok(n)
}
pub fn compress_dir(
input_dir: &PathBuf,
output: &PathBuf,
no_skip: bool,
plugin: Option<&znippy_common::plugin::PluginRegistry>,
repo: Option<&str>,
sink_factory: Option<znippy_common::MetaSinkFactory>,
) -> Result<CompressionReport> {
let mut total_dirs = 0u64;
let all_files: Arc<Vec<PathBuf>> = Arc::new(
WalkDir::new(input_dir)
.into_iter()
.filter_map(|e| e.ok())
.filter_map(|e| {
if e.file_type().is_dir() {
total_dirs += 1;
None
} else if e.file_type().is_file() {
Some(e.into_path())
} else {
None
}
})
.collect(),
);
let total_files = all_files.len() as u64;
let ext_fields: Vec<znippy_common::arrow::datatypes::Field> =
plugin.map(|r| r.schema_fields()).unwrap_or_default();
let output_path = output.with_extension("znippy");
let file = Arc::new(File::create(&output_path)?);
let out_cursor = Arc::new(AtomicU64::new(0));
let num_workers = CONFIG.max_core_in_flight.max(1);
let slice_size = SLOT_SIZE / num_workers.max(1);
// ── PARTITION ────────────────────────────────────────────────────────────
let mut big_indices: Vec<usize> = Vec::new();
let mut small_indices: Vec<usize> = Vec::new();
for (i, path) in all_files.iter().enumerate() {
let size = path.metadata().map(|m| m.len()).unwrap_or(0);
if size > slice_size as u64 || size == 0 {
big_indices.push(i);
} else {
small_indices.push(i);
}
}
let mut ext_meta: Vec<FileExtMeta> = vec![None; all_files.len()];
let mut uncompressed_files = 0u64;
let mut uncompressed_bytes = 0u64;
let mut compressed_files = 0u64;
let mut compressed_bytes = 0u64;
let mut total_chunks = 0u64;
// Metadata index schema (shared across both passes). The batches are
// collected and handed to the metadata sink as one sub-index below.
let meta_map = build_arrow_metadata_for_config(&CONFIG);
let composed = compose_index_schema(&ext_fields);
let schema_with_meta =
arrow::datatypes::Schema::new_with_metadata(composed.fields().to_vec(), meta_map);
let mut index_batches: Vec<arrow::record_batch::RecordBatch> = Vec::new();
let input_dir_for_paths = input_dir.clone();
let all_files_for_paths = Arc::clone(&all_files);
// ══════════════════════════════════════════════════════════════════════════
// PASS 1: BIG FILES
// ══════════════════════════════════════════════════════════════════════════
if !big_indices.is_empty() {
let (uf, ub, cf, cb, blobs, meta) = run_big_pass(
&all_files, input_dir, &big_indices, no_skip, plugin,
&file, &out_cursor, num_workers,
)?;
uncompressed_files += uf; uncompressed_bytes += ub;
compressed_files += cf; compressed_bytes += cb;
for (idx, m) in meta { if idx < ext_meta.len() { ext_meta[idx] = m; } }
total_chunks += blobs.len() as u64;
let all_f = Arc::clone(&all_files_for_paths);
let inp = input_dir_for_paths.clone();
let resolver = |file_index: u64| {
let idx = file_index as usize;
all_f[idx].strip_prefix(&inp).unwrap_or(&all_f[idx])
.to_string_lossy().to_string()
};
let batch = build_metadata_batch(&blobs, resolver, &ext_meta, &ext_fields)
.map_err(|e| anyhow!("big index batch: {e}"))?;
index_batches.push(batch);
}
// ══════════════════════════════════════════════════════════════════════════
// PASS 2: SMALL FILES
// ══════════════════════════════════════════════════════════════════════════
if !small_indices.is_empty() {
let (uf, ub, cf, cb, blobs, meta) = run_small_pass(
&all_files, input_dir, &small_indices, no_skip, plugin,
&file, &out_cursor, num_workers, slice_size,
)?;
uncompressed_files += uf; uncompressed_bytes += ub;
compressed_files += cf; compressed_bytes += cb;
for (idx, m) in meta { if idx < ext_meta.len() { ext_meta[idx] = m; } }
total_chunks += blobs.len() as u64;
let all_f = Arc::clone(&all_files_for_paths);
let inp = input_dir_for_paths.clone();
let resolver = |file_index: u64| {
let idx = file_index as usize;
all_f[idx].strip_prefix(&inp).unwrap_or(&all_f[idx])
.to_string_lossy().to_string()
};
let batch = build_metadata_batch(&blobs, resolver, &ext_meta, &ext_fields)
.map_err(|e| anyhow!("small index batch: {e}"))?;
index_batches.push(batch);
}
// ══════════════════════════════════════════════════════════════════════════
// FINALIZE: write the metadata layer (one sub-index of all batches) via the sink
// ══════════════════════════════════════════════════════════════════════════
let index_offset = out_cursor.load(Ordering::Relaxed);
let blob_bytes = index_offset;
let pkg_type_val: i8 = plugin.and_then(|r| r.type_id()).unwrap_or(0);
let mut sink: Box<dyn ArchiveMetaSink> = match sink_factory {
Some(make) => make(Arc::clone(&file), blob_bytes),
None => Box::new(ArrowIpcSink::new(Arc::clone(&file), blob_bytes)),
};
sink.push_subindex(
&schema_with_meta,
&index_batches,
GroupKey {
pkg_type: pkg_type_val,
repo: repo.unwrap_or("").to_string(),
module_name: String::new(),
},
)?;
let total_bytes_out = sink.finish()?;
Ok(CompressionReport {
total_files,
compressed_files,
uncompressed_files,
chunks: total_chunks,
total_dirs,
total_bytes_in: compressed_bytes + uncompressed_bytes,
total_bytes_out,
compressed_bytes,
uncompressed_bytes,
compression_ratio: if uncompressed_bytes > 0 {
(compressed_bytes as f32 / blob_bytes.max(1) as f32) * 100.0
} else {
0.0
},
})
}
// ─────────────────────────────────────────────────────────────────────────────
// PASS 1: big files — sequential chunked reads, re-read for metadata
// ─────────────────────────────────────────────────────────────────────────────
fn run_big_pass(
all_files: &Arc<Vec<PathBuf>>,
input_dir: &PathBuf,
big_indices: &[usize],
no_skip: bool,
plugin: Option<&znippy_common::plugin::PluginRegistry>,
file: &Arc<File>,
out_cursor: &Arc<AtomicU64>,
num_workers: usize,
) -> Result<(u64, u64, u64, u64, Vec<BlobMeta>, Vec<(usize, FileExtMeta)>)> {
let pool = Magazine::new(NUM_SLOTS, SLOT_SIZE, num_workers);
let returner = pool.returner();
let (tx_slice, rx_slice) = bounded(NUM_SLOTS * 4);
let (tx_write, rx_write) = unbounded::<WriteJob>();
let (tx_meta, rx_meta) = unbounded::<(usize, FileExtMeta)>();
let plugin_addr: usize = plugin.map(|p| p as *const _ as usize).unwrap_or(0);
let reader = {
let all_files = Arc::clone(all_files);
let input_dir = input_dir.clone();
let big_indices = big_indices.to_vec();
let tx_meta = tx_meta.clone();
thread::spawn(move || -> (u64, u64, u64, u64) {
let plugin_ref: Option<&znippy_common::plugin::PluginRegistry> =
if plugin_addr != 0 { Some(unsafe { &*(plugin_addr as *const _) }) } else { None };
let mut uf = 0u64; let mut ub = 0u64;
let mut cf = 0u64; let mut cb = 0u64;
let mut cur = None;
let ss = pool.slice_size();
// PERF (Law 2): one buffer reused for the big-file metadata re-read,
// instead of a fresh `std::fs::read` allocation per matching file.
let mut meta_buf: Vec<u8> = Vec::new();
for &file_index in &big_indices {
let path = &all_files[file_index];
let file_size = path.metadata().map(|m| m.len()).unwrap_or(0);
let skip = !no_skip && should_skip_compression(path);
if skip { uf += 1; ub += file_size; } else { cf += 1; cb += file_size; }
if file_size == 0 {
ensure_room(&pool, &tx_slice, &mut cur, 0);
cur.as_mut().unwrap().commit_slice(0, skip, file_index as u64, 0, 0);
continue;
}
let f = match File::open(path) {
Ok(f) => f,
Err(e) => { log::warn!("[big] open {}: {}", path.display(), e); continue; }
};
let mut rdr = BufReader::new(f);
let mut fdata_offset = 0u64;
let mut chunk_seq = 0u32;
let mut remaining = file_size;
while remaining > 0 {
let want = ss.min(remaining as usize);
ensure_room(&pool, &tx_slice, &mut cur, want);
let fill = cur.as_mut().unwrap();
let buf = fill.writable(want);
let got = match read_fully(&mut rdr, buf) {
Ok(g) => g,
Err(e) => { log::warn!("[big] read {}: {}", path.display(), e); break; }
};
if got == 0 { break; }
fill.commit_slice(got, skip, file_index as u64, fdata_offset, chunk_seq);
fdata_offset += got as u64;
chunk_seq += 1;
remaining = remaining.saturating_sub(got as u64);
if got < want { break; }
}
// Big file metadata: re-read the file (acceptable for large files).
if let Some(reg) = plugin_ref {
let rel = path.strip_prefix(&input_dir).unwrap_or(path).to_string_lossy();
if reg.matches(&rel) {
// Reuse meta_buf: truncate then read the whole file into it.
meta_buf.clear();
match File::open(path)
.and_then(|mut f| f.read_to_end(&mut meta_buf))
{
Ok(_) => {
if let Some((tid, row)) = reg.extract_typed(&rel, &meta_buf) {
tx_meta.send((file_index, Some((tid, row)))).ok();
}
}
Err(e) => log::warn!("[big] meta re-read {}: {}", path.display(), e),
}
}
}
}
if let Some(fill) = cur.take() {
for s in fill.publish() { tx_slice.send(s).ok(); }
}
// Drain: reclaim all slots to prove workers/writer are done with slot memory.
for _ in 0..NUM_SLOTS {
if pool.claim().is_none() { break; }
}
drop(tx_slice);
drop(tx_meta);
drop(pool);
(uf, ub, cf, cb)
})
};
let workers = spawn_workers(num_workers, rx_slice.clone(), tx_write.clone(), returner.clone());
let writer = spawn_writer(Arc::clone(file), Arc::clone(out_cursor), returner.clone(), rx_write);
drop(tx_write); drop(rx_slice); drop(tx_meta);
let (uf, ub, cf, cb) = reader.join().map_err(|_| anyhow!("big reader panicked"))?;
for w in workers { w.join().map_err(|_| anyhow!("big worker panicked"))??; }
let blobs = writer.join().map_err(|_| anyhow!("big writer panicked"))??;
let mut meta = Vec::new();
while let Ok(m) = rx_meta.try_recv() { meta.push(m); }
Ok((uf, ub, cf, cb, blobs, meta))
}
// ─────────────────────────────────────────────────────────────────────────────
// PASS 2: small files — read into slot, metadata from in-memory data
// ─────────────────────────────────────────────────────────────────────────────
fn run_small_pass(
all_files: &Arc<Vec<PathBuf>>,
input_dir: &PathBuf,
small_indices: &[usize],
no_skip: bool,
plugin: Option<&znippy_common::plugin::PluginRegistry>,
file: &Arc<File>,
out_cursor: &Arc<AtomicU64>,
num_workers: usize,
_slice_size: usize,
) -> Result<(u64, u64, u64, u64, Vec<BlobMeta>, Vec<(usize, FileExtMeta)>)> {
let pool = Magazine::new(NUM_SLOTS, SLOT_SIZE, num_workers);
let returner = pool.returner();
let (tx_slice, rx_slice) = bounded(NUM_SLOTS * 4);
let (tx_write, rx_write) = unbounded::<WriteJob>();
let (tx_meta, rx_meta) = unbounded::<(usize, FileExtMeta)>();
let plugin_addr: usize = plugin.map(|p| p as *const _ as usize).unwrap_or(0);
let reader = {
let all_files = Arc::clone(all_files);
let input_dir = input_dir.clone();
let small_indices = small_indices.to_vec();
let tx_meta = tx_meta.clone();
thread::spawn(move || -> (u64, u64, u64, u64) {
let plugin_ref: Option<&znippy_common::plugin::PluginRegistry> =
if plugin_addr != 0 { Some(unsafe { &*(plugin_addr as *const _) }) } else { None };
let mut uf = 0u64; let mut ub = 0u64;
let mut cf = 0u64; let mut cb = 0u64;
let mut cur = None;
// io_uring may be unavailable (old kernel < 5.1, seccomp SCMP_ACT_KILL,
// Docker with default seccomp). Try to initialise and fall back to
// std::fs::read if creation fails.
// io_uring may be unavailable. Also, the test seam can force the
// fallback by leaving `ring` as None even when init would succeed.
let mut ring: Option<io_uring::IoUring> = if force_small_fallback() {
None
} else {
io_uring::IoUring::new(256).ok()
};
let mut idx = 0usize;
let n = small_indices.len();
// PERF (Law 2): per-batch scratch buffers hoisted out of the loop and
// cleared each batch — no per-128-file reallocation of these four Vecs.
let mut batch: Vec<(usize, usize, bool)> = Vec::with_capacity(128); // (file_index, size, skip)
let mut cstrings: Vec<std::ffi::CString> = Vec::with_capacity(128);
let mut fds: Vec<i32> = Vec::with_capacity(128);
let mut offsets: Vec<usize> = Vec::with_capacity(128);
let mut read_results: Vec<usize> = Vec::with_capacity(128);
while idx < n {
// Collect a batch that fits in current slot
if cur.is_none() { cur = pool.claim(); }
batch.clear();
let mut batch_total = 0usize;
while idx < n && batch.len() < 128 {
let file_index = small_indices[idx];
let path = &all_files[file_index];
let file_size = path.metadata().map(|m| m.len()).unwrap_or(0) as usize;
let skip = !no_skip && should_skip_compression(path);
let remaining = cur.as_ref().unwrap().remaining();
if batch_total + file_size > remaining {
if batch_total == 0 {
// Slot is too full for even one file — publish and get new slot
if let Some(fill) = cur.take() {
for s in fill.publish() { tx_slice.send(s).ok(); }
}
cur = pool.claim();
continue; // retry with new slot
}
break; // process what we have
}
if skip { uf += 1; ub += file_size as u64; }
else { cf += 1; cb += file_size as u64; }
batch.push((file_index, file_size, skip));
batch_total += file_size;
idx += 1;
}
if batch.is_empty() { continue; }
let blen = batch.len();
// Phase 2: batch read directly into slot via writable()
// writable() gives a slice at cursor without advancing — use it for
// the entire batch, then commit each file.
let fill = cur.as_mut().unwrap();
let slot_buf = fill.writable(batch_total);
let slot_ptr = slot_buf.as_mut_ptr();
offsets.clear();
let mut off = 0usize;
for &(_, size, _) in &batch {
offsets.push(off);
off += size;
}
read_results.clear();
read_results.resize(blen, 0);
if let Some(ref mut ring) = ring {
// Phase 1 (io_uring path): batch open
cstrings.clear();
for &(fi, _, _) in &batch {
let p = all_files[fi].as_os_str().as_encoded_bytes();
cstrings.push(unsafe { std::ffi::CString::from_vec_unchecked(p.to_vec()) });
}
fds.clear();
fds.resize(blen, -1);
// PERF: step over the range in 256-wide windows instead of
// materializing a throwaway index Vec just to `.chunks(256)`.
let mut start = 0usize;
while start < blen {
let end = (start + 256).min(blen);
for i in start..end {
let open_e = io_uring::opcode::OpenAt::new(
io_uring::types::Fd(libc::AT_FDCWD),
cstrings[i].as_ptr(),
)
.flags(libc::O_RDONLY | libc::O_CLOEXEC)
.build()
.user_data(i as u64);
unsafe { ring.submission().push(&open_e).ok(); }
}
let want = end - start;
ring.submit_and_wait(want).ok();
let mut got = 0;
while got < want {
if let Some(cqe) = ring.completion().next() {
fds[cqe.user_data() as usize] = cqe.result();
got += 1;
}
}
start = end;
}
// Phase 2 (io_uring path): batch read
// PERF: 256-wide windows via step, no throwaway index Vec.
let mut start = 0usize;
while start < blen {
let end = (start + 256).min(blen);
let mut to_submit = 0;
for i in start..end {
if fds[i] < 0 { continue; }
let (_, size, _) = batch[i];
let dst = unsafe { slot_ptr.add(offsets[i]) };
let read_e = io_uring::opcode::Read::new(
io_uring::types::Fd(fds[i]),
dst,
size as u32,
)
.build()
.user_data(i as u64);
unsafe { ring.submission().push(&read_e).ok(); }
to_submit += 1;
}
if to_submit > 0 {
ring.submit_and_wait(to_submit).ok();
let mut got = 0;
while got < to_submit {
if let Some(cqe) = ring.completion().next() {
let i = cqe.user_data() as usize;
read_results[i] = if cqe.result() > 0 { cqe.result() as usize } else { 0 };
got += 1;
}
}
}
start = end;
}
// Phase 3 (io_uring path): close fds
for &fd in &fds {
if fd >= 0 { unsafe { libc::close(fd); } }
}
} else {
// Fallback path: io_uring unavailable — use std::fs::read per file.
for i in 0..blen {
let (fi, size, _) = batch[i];
let path = &all_files[fi];
let dst = unsafe { std::slice::from_raw_parts_mut(slot_ptr.add(offsets[i]), size) };
if let Ok(mut f) = std::fs::File::open(path) {
// `size` comes from metadata().len(); read the whole file
// (looping past short reads) and record the true byte
// count, so the slot holds exactly `got` valid bytes —
// symmetric with the io_uring Read path.
read_results[i] = match read_fully(&mut f, dst) {
Ok(g) => g,
Err(e) => {
log::warn!("[small] read {}: {}", path.display(), e);
0
}
};
}
}
}
// Phase 4: plugin extraction + commit each file
let fill = cur.as_mut().unwrap();
for i in 0..batch.len() {
let (file_index, _size, skip) = batch[i];
let got = read_results[i];
if let Some(reg) = plugin_ref {
if got > 0 {
let path = &all_files[file_index];
let rel = path.strip_prefix(&input_dir).unwrap_or(path).to_string_lossy();
if reg.matches(&rel) {
let data = unsafe {
std::slice::from_raw_parts(slot_ptr.add(offsets[i]), got)
};
if let Some((tid, row)) = reg.extract_typed(&rel, data) {
tx_meta.send((file_index, Some((tid, row)))).ok();
}
}
}
}
fill.commit_slice(got, skip, file_index as u64, 0, 0);
}
}
if let Some(fill) = cur.take() {
for s in fill.publish() { tx_slice.send(s).ok(); }
}
for _ in 0..NUM_SLOTS {
if pool.claim().is_none() { break; }
}
drop(tx_slice);
drop(tx_meta);
drop(pool);
(uf, ub, cf, cb)
})
};
let workers = spawn_workers(num_workers, rx_slice.clone(), tx_write.clone(), returner.clone());
let writer = spawn_writer(Arc::clone(file), Arc::clone(out_cursor), returner.clone(), rx_write);
drop(tx_write); drop(rx_slice); drop(tx_meta);
let (uf, ub, cf, cb) = reader.join().map_err(|_| anyhow!("small reader panicked"))?;
for w in workers { w.join().map_err(|_| anyhow!("small worker panicked"))??; }
let blobs = writer.join().map_err(|_| anyhow!("small writer panicked"))??;
let mut meta = Vec::new();
while let Ok(m) = rx_meta.try_recv() { meta.push(m); }
Ok((uf, ub, cf, cb, blobs, meta))
}
// ─────────────────────────────────────────────────────────────────────────────
// Shared helpers
// ─────────────────────────────────────────────────────────────────────────────
fn spawn_workers(
num_workers: usize,
rx_slice: crossbeam_channel::Receiver<znippy_common::slotpool::Round>,
tx_write: crossbeam_channel::Sender<WriteJob>,
returner: znippy_common::slotpool::Ejector,
) -> Vec<thread::JoinHandle<Result<()>>> {
let level = CONFIG.compression_level;
(0..num_workers).map(|_| {
let rx = rx_slice.clone();
let tw = tx_write.clone();
let ret = returner.clone();
thread::spawn(move || -> Result<()> {
let mut cctx = znippy_common::codec::CompressCtx::new(level)?;
let mut reuse_buf: Vec<u8> = Vec::new();
while let Ok(slice) = rx.recv() {
let src = unsafe { slice.as_slice() };
let checksum = *blake3::hash(src).as_bytes();
let len = src.len();
// Snapshot the scalar index fields before any move of `slice`.
let (file_index, fdata_offset, chunk_seq) =
(slice.file_index, slice.fdata_offset, slice.chunk_seq);
if slice.skip {
// PERF (Law 1): zero-copy skip path. Hand the slot `Round`
// straight to the writer — it pwrites from the slot bytes and
// releases the slot afterward. No memcpy out of the slot, which
// is the full-I/O-speed stored-raw common case (.jar/.gz/...).
tw.send(WriteJob {
payload: Payload::Slot(slice), on_disk_len: len,
file_index, fdata_offset,
chunk_seq, checksum, compressed: false,
uncompressed_size: len as u64,
}).ok();
} else {
let mut buf = std::mem::take(&mut reuse_buf);
let n = cctx.compress_into(src, &mut buf)?;
let usz = len as u64;
if n >= len {
// Incompressible: the codec frame is no smaller than the
// input. Store the raw bytes — never write a blob bigger
// than the source, and skip the decode cost.
// PERF (Law 1): the worker keeps `buf` (recycled next loop)
// and hands the slot `Round` to the writer for a zero-copy
// pwrite, instead of memcpying src into buf.
reuse_buf = buf; // recycle the just-allocated compress buf
tw.send(WriteJob {
payload: Payload::Slot(slice), on_disk_len: len,
file_index, fdata_offset,
chunk_seq, checksum, compressed: false,
uncompressed_size: usz,
}).ok();
} else {
// Compressed output is owned by `buf`; release the slot now
// since the source bytes are no longer needed.
ret.release_one(slice.slot_id);
tw.send(WriteJob {
payload: Payload::Buf(buf), on_disk_len: n,
file_index, fdata_offset,
chunk_seq, checksum, compressed: true,
uncompressed_size: usz,
}).ok();
}
}
}
Ok(())
})
}).collect()
}
fn spawn_writer(
file: Arc<File>,
out_cursor: Arc<AtomicU64>,
returner: znippy_common::slotpool::Ejector,
rx_write: crossbeam_channel::Receiver<WriteJob>,
) -> thread::JoinHandle<Result<Vec<BlobMeta>>> {
thread::spawn(move || -> Result<Vec<BlobMeta>> {
let mut blobs = Vec::new();
while let Ok(job) = rx_write.recv() {
let off = out_cursor.fetch_add(job.on_disk_len as u64, Ordering::Relaxed);
match &job.payload {
Payload::Buf(buf) => {
file.write_all_at(&buf[..job.on_disk_len], off)?;
}
Payload::Slot(round) => {
// PERF (Law 1): pwrite straight from the slot, then release it
// so the reader can reuse it. The slot stayed alive through the
// worker's hand-off precisely so this copy never happened.
let src = unsafe { round.as_slice() };
file.write_all_at(&src[..job.on_disk_len], off)?;
returner.release_one(round.slot_id);
}
}
blobs.push(BlobMeta {
chunk_meta: ChunkMeta {
fdata_offset: job.fdata_offset, file_index: job.file_index,
chunk_seq: job.chunk_seq, checksum: job.checksum,
compressed: job.compressed, uncompressed_size: job.uncompressed_size,
compressed_size: job.on_disk_len as u64,
},
blob_offset: off, blob_size: job.on_disk_len as u64,
});
}
Ok(blobs)
})
}
fn ensure_room<'p>(
pool: &'p Magazine,
tx_slice: &crossbeam_channel::Sender<znippy_common::slotpool::Round>,
cur: &mut Option<znippy_common::slotpool::Clip<'p>>,
need: usize,
) {
loop {
if cur.is_none() {
*cur = pool.claim();
if cur.is_none() { return; }
}
if cur.as_ref().unwrap().remaining() >= need { return; }
let slices = cur.take().unwrap().publish();
for s in slices { tx_slice.send(s).ok(); }
}
}
// ─────────────────────────────────────────────────────────────────────────────
// Tests
// ─────────────────────────────────────────────────────────────────────────────
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
use std::io::Write;
/// Unique scratch dir under the system temp dir (no tempfile dev-dep here).
/// Cleaned up by the caller on success; left in place on panic for triage.
fn scratch(tag: &str) -> PathBuf {
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos();
let dir = std::env::temp_dir().join(format!(
"znippy_fallback_test_{tag}_{}_{nanos}",
std::process::id()
));
fs::create_dir_all(&dir).unwrap();
dir
}
/// Walk `root` and collect (relative-path, bytes) for every file, so two
/// extracted trees can be compared byte-for-byte regardless of walk order.
fn read_tree(root: &PathBuf) -> std::collections::BTreeMap<String, Vec<u8>> {
let mut out = std::collections::BTreeMap::new();
for e in WalkDir::new(root).into_iter().filter_map(|e| e.ok()) {
if e.file_type().is_file() {
let rel = e
.path()
.strip_prefix(root)
.unwrap()
.to_string_lossy()
.to_string();
out.insert(rel, fs::read(e.path()).unwrap());
}
}
out
}
/// INJECT-ASSERT (audit HIGH): the io_uring small-file path and the graceful
/// fallback path MUST produce byte-identical archive output. We compress the
/// SAME directory of several small files twice — once with io_uring, once
/// with the fallback forced via the `FORCE_SMALL_FALLBACK` test seam — and
/// assert that:
/// 1. the per-chunk index checksums match exactly (the bytes the codec saw
/// were identical, i.e. the fallback read the same bytes into the slot),
/// 2. the decompressed file trees are byte-for-byte identical,
/// 3. both equal the original injected input bytes.
///
/// This is a real input → real output assertion, not a "didn't panic" smoke
/// test: every file carries distinct, deliberately-chosen bytes (including a
/// multi-batch fan-out, an empty file, an already-compressed `.gz` that is
/// stored raw, and a highly-compressible file) so a wrong byte, a short read,
/// or a wrong length would change a checksum and fail the assertion.
#[test]
fn fallback_byte_identical_to_io_uring() {
let input = scratch("input");
// Inject a directory of several small files with distinct real bytes.
// Many files so the small-pass batches them and the read loop runs hot.
fs::write(input.join("alpha.txt"), b"the quick brown fox").unwrap();
fs::write(input.join("beta.bin"), (0u8..=255).collect::<Vec<u8>>()).unwrap();
fs::create_dir_all(input.join("nested/deep")).unwrap();
fs::write(input.join("nested/gamma.txt"), b"nested content here").unwrap();
fs::write(input.join("nested/deep/delta.dat"), vec![0xABu8; 4096]).unwrap();
// Empty file: read must record 0 bytes, not short-read into garbage.
fs::write(input.join("empty.txt"), b"").unwrap();
// Already-compressed extension: stored raw at full I/O speed.
fs::write(input.join("packed.gz"), vec![0x1F, 0x8B, 0x08, 0x00, 0x99, 0x42]).unwrap();
// Highly compressible file: exercises the real codec path.
fs::write(input.join("zeros.log"), vec![0u8; 9000]).unwrap();
// A spray of small files to force multiple read iterations in one batch.
for i in 0..40 {
let mut f = fs::File::create(input.join(format!("frag_{i:03}.txt"))).unwrap();
// distinct content per file so any swap/short-read flips a checksum
writeln!(f, "fragment number {i} :: payload {}", "x".repeat(i)).unwrap();
}
let input = input; // PathBuf
// ── Run A: io_uring path (seam OFF) ─────────────────────────────────
FORCE_SMALL_FALLBACK.store(false, Ordering::SeqCst);
let out_a_dir = scratch("out_a");
let arc_a = out_a_dir.join("a.znippy");
let report_a = compress_dir(&input, &arc_a, false, None, None, None)
.expect("io_uring compress");
let arc_a = arc_a.with_extension("znippy");
// ── Run B: forced fallback path (seam ON) ───────────────────────────
FORCE_SMALL_FALLBACK.store(true, Ordering::SeqCst);
let out_b_dir = scratch("out_b");
let arc_b = out_b_dir.join("b.znippy");
let report_b = compress_dir(&input, &arc_b, false, None, None, None)
.expect("fallback compress");
let arc_b = arc_b.with_extension("znippy");
FORCE_SMALL_FALLBACK.store(false, Ordering::SeqCst);
// ── Assert 1: same file/chunk accounting ────────────────────────────
assert_eq!(report_a.total_files, report_b.total_files, "file count differs");
assert_eq!(report_a.chunks, report_b.chunks, "chunk count differs");
assert_eq!(
report_a.total_bytes_in, report_b.total_bytes_in,
"input byte total differs"
);
// ── Assert 2: per-chunk index checksums are identical ───────────────
// Read each archive's index and collect the set of (relative_path,
// chunk_seq, checksum, uncompressed_size, compressed_size). If the
// fallback read a single wrong/short byte, a blake3 checksum here would
// differ and this set comparison would fail.
let sig_a = chunk_signatures(&arc_a);
let sig_b = chunk_signatures(&arc_b);
assert_eq!(
sig_a, sig_b,
"io_uring vs fallback chunk checksums differ — NOT byte-identical"
);
// ── Assert 3: decompressed trees byte-identical, and == original ────
let dec_a = scratch("dec_a");
let dec_b = scratch("dec_b");
let va = znippy_common::decompress_archive(&arc_a, true, &dec_a)
.expect("decompress A");
let vb = znippy_common::decompress_archive(&arc_b, true, &dec_b)
.expect("decompress B");
assert_eq!(va.corrupt_files, 0, "io_uring archive had corrupt files");
assert_eq!(vb.corrupt_files, 0, "fallback archive had corrupt files");
let tree_orig = read_tree(&input);
let tree_a = read_tree(&dec_a);
let tree_b = read_tree(&dec_b);
assert_eq!(tree_a, tree_b, "decompressed trees differ between paths");
assert_eq!(tree_a, tree_orig, "io_uring round-trip != original bytes");
assert_eq!(tree_b, tree_orig, "fallback round-trip != original bytes");
// Cleanup on success (left behind on panic for triage).
for d in [&input, &out_a_dir, &out_b_dir, &dec_a, &dec_b] {
let _ = fs::remove_dir_all(d);
}
}
/// Collect a sorted, order-independent signature of every chunk in an
/// archive's index: (relative_path, chunk_seq, checksum, uncompressed_size,
/// blob_size, compressed). Blob offsets are deliberately excluded — they are
/// assigned in writer/thread completion order and so are legitimately
/// non-deterministic, but the bytes (hence checksums) and sizes must not be.
fn chunk_signatures(
archive: &PathBuf,
) -> Vec<(String, u32, [u8; 32], u64, u64, bool)> {
use arrow::array::{
BooleanArray, FixedSizeBinaryArray, StringArray, UInt32Array, UInt64Array,
};
let (_schema, batches) = znippy_common::index::read_znippy_index_filtered(
archive,
&znippy_common::index::IndexFilter::default(),
)
.expect("read index");
let mut sigs = Vec::new();
for batch in &batches {
let paths = batch
.column_by_name("relative_path")
.unwrap()
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
let seqs = batch
.column_by_name("chunk_seq")
.unwrap()
.as_any()
.downcast_ref::<UInt32Array>()
.unwrap();
let sums = batch
.column_by_name("checksum")
.unwrap()
.as_any()
.downcast_ref::<FixedSizeBinaryArray>()
.unwrap();
let usz = batch
.column_by_name("uncompressed_size")
.unwrap()
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap();
let bsz = batch
.column_by_name("blob_size")
.unwrap()
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap();
let comp = batch
.column_by_name("compressed")
.unwrap()
.as_any()
.downcast_ref::<BooleanArray>()
.unwrap();
for i in 0..batch.num_rows() {
let mut sum = [0u8; 32];
sum.copy_from_slice(&sums.value(i)[..32]);
sigs.push((
paths.value(i).to_string(),
seqs.value(i),
sum,
usz.value(i),
bsz.value(i),
comp.value(i),
));
}
}
sigs.sort();
sigs
}
}