Add checkpoint bucket usage + compactors early termination

Author: Guilucand

crates/io/src/concurrent/temp_reads/creads_utils.rs

@@ -272,6 +272,7 @@ pub mod helpers {
                 $reader:expr,
                 $read_thread:expr,
                 $with_multiplicity:expr,
+                |$checkpoint_data:ident| $c:expr,
                 |$data: ident, $extra_buffer: ident| $f:expr
             );
 
@@ -293,7 +294,8 @@ pub mod helpers {
                         $BucketMode,
                         WithMultiplicity,
                     >, false>(read_thread, Vec::new(), <$E>::new_temp_buffer());
-                while let Some((items, _)) = items.get_next_checkpoint() {
+                while let Some((items, $checkpoint_data)) = items.get_next_checkpoint() {
+                    $c
                     while let Some(($data, $extra_buffer)) = items.next() {
                         $f
                     }

crates/kmers_transform/src/reader.rs

@@ -426,6 +426,7 @@ impl<F: KmersTransformExecutorFactory> KmersTransformReader<F> {
             }
 
             let data_format: MinimizerBucketMode = reader.get_data_format_info().unwrap();
+            let mut checkpoint_rewrite_bucket = None;
 
             creads_helper! {
                 helper_read_bucket_with_opt_multiplicity::<
@@ -436,17 +437,19 @@ impl<F: KmersTransformExecutorFactory> KmersTransformReader<F> {
                     reader,
                     async_reader_thread.clone(),
                     matches!(data_format, MinimizerBucketMode::Compacted),
+                    |checkpoint_data| { checkpoint_rewrite_bucket = checkpoint_data.map(|d| d.target_subbucket); } ,
                     |read_info, extra_buffer| {
                         let bucket = if has_single_addr {
                             0
                         } else {
-                            let orig_bucket = F::PreprocessorType::get_rewrite_bucket(
+                            let orig_bucket = checkpoint_rewrite_bucket
+                            .unwrap_or_else(|| F::PreprocessorType::get_rewrite_bucket(
                                 global_extra_data.get_k(),
                                 global_extra_data.get_m(),
                                 &read_info,
                                 bucket_info.used_hash_bits,
                                 bucket_info.second_buckets_log_max,
-                            ) as usize;
+                            )) as usize;
 
                             bucket_info.buckets_remapping[orig_bucket]
                         };

crates/minimizer_bucketing/src/compactor.rs

@@ -237,6 +237,7 @@ impl<E: MinimizerBucketingExecutorFactory + Sync + Send + 'static> AsyncExecutor
                     processed_buckets += 1;
 
                     let format_data: MinimizerBucketMode = reader.get_data_format_info().unwrap();
+                    let mut checkpoint_rewrite_bucket = None;
                     creads_helper! {
                         helper_read_bucket_with_opt_multiplicity::<
                             E::ExtraData,
@@ -246,14 +247,16 @@ impl<E: MinimizerBucketingExecutorFactory + Sync + Send + 'static> AsyncExecutor
                             &reader,
                             read_thread.clone(),
                             matches!(format_data, MinimizerBucketMode::Compacted),
+                            |checkpoint_data| { checkpoint_rewrite_bucket = checkpoint_data.map(|d| d.target_subbucket); } ,
                             |data, _extra_buffer| {
 
-                                let rewrite_bucket = E::RewriteBucketCompute::get_rewrite_bucket(global_params.common.k,
+                                let rewrite_bucket = checkpoint_rewrite_bucket
+                                .unwrap_or_else(|| E::RewriteBucketCompute::get_rewrite_bucket(global_params.common.k,
                                     global_params.common.m,
                                     &data,
                                     used_hash_bits,
                                     second_buckets_log_max,
-                                );
+                                ));
                                 sequences_deltas[rewrite_bucket as usize] += 1;
 
                                 let (flags, _, _extra, read, multiplicity) = data;
@@ -281,7 +284,9 @@ impl<E: MinimizerBucketingExecutorFactory + Sync + Send + 'static> AsyncExecutor
                     }
 
                     // Do not process more buckets if it will increase the maximum number of allowed sequences
-                    if processed_buckets >= 2 && total_sequences > MAXIMUM_SEQUENCES {
+                    if !global_params.common.is_active.load(Ordering::Relaxed)
+                        || processed_buckets >= 2 && total_sequences > MAXIMUM_SEQUENCES
+                    {
                         break;
                     }
                 }

crates/minimizer_bucketing/src/lib.rs

@@ -52,7 +52,7 @@ use std::marker::PhantomData;
 use std::ops::Deref;
 use std::ops::Range;
 use std::path::{Path, PathBuf};
-use std::sync::atomic::{AtomicI64, AtomicU64, AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, AtomicUsize, Ordering};
 use std::sync::Arc;
 
 #[derive(Serialize, Deserialize, Clone, Copy, Debug)]
@@ -167,6 +167,7 @@ pub struct MinimizerBucketingCommonData<GlobalData> {
     pub global_counters: Vec<Vec<AtomicU64>>,
     pub compaction_offsets: Vec<AtomicI64>,
     pub global_data: GlobalData,
+    pub is_active: AtomicBool,
 }
 
 impl<GlobalData> MinimizerBucketingCommonData<GlobalData> {
@@ -197,6 +198,7 @@ impl<GlobalData> MinimizerBucketingCommonData<GlobalData> {
                 .collect(),
             compaction_offsets: (0..buckets_count).map(|_| AtomicI64::new(0)).collect(),
             global_data,
+            is_active: AtomicBool::new(true),
         }
     }
 }
@@ -607,6 +609,14 @@ impl GenericMinimizerBucketing {
             global_context.executor_group_address.write().take();
 
             execution_context.wait_for_completion(writer_executors);
+
+            // Let compactors know that the phase is finishing,
+            // so they can shortcut and avoid processing other buckets
+            global_context
+                .common
+                .is_active
+                .store(false, Ordering::Relaxed);
+
             execution_context.wait_for_completion(compactor_executors);
 
             execution_context.join_all();