use std::{
error::Error,
sync::{atomic::AtomicBool, Arc},
};
use bytemuck::Pod;
use crossbeam_channel::{bounded, unbounded, Receiver, Select, Sender};
use log::debug;
use crate::{
chunk::{FileChunk, FileChunkDir, MemChunk},
merge::merge_chunks_with_binary_heap,
Result,
};
pub struct SortConfig {
pub(crate) max_chunk_bytes: usize,
pub(crate) concurrency: usize,
pub(crate) merge_k: usize,
pub(crate) canceled: Arc<AtomicBool>,
}
impl Default for SortConfig {
fn default() -> Self {
Self {
max_chunk_bytes: 512 * 1024 * 1024,
concurrency: 8,
merge_k: 16,
canceled: Arc::new(AtomicBool::new(false)),
}
}
}
impl SortConfig {
pub fn new() -> Self {
Default::default()
}
pub fn set_cancel_flag(self, canceled: Arc<AtomicBool>) -> Self {
Self { canceled, ..self }
}
pub fn get_cancel_flag(&self) -> Arc<AtomicBool> {
self.canceled.clone()
}
pub(crate) fn ensure_not_canceled<E>(&self) -> Result<(), E> {
if self.canceled.load(std::sync::atomic::Ordering::Relaxed) {
Err(crate::Error::Canceled)
} else {
Ok(())
}
}
pub fn max_chunk_bytes(self, max_chunk_bytes: usize) -> Self {
assert!(
max_chunk_bytes > 0,
"max_chunk_bytes must be greater than 0"
);
Self {
max_chunk_bytes,
..self
}
}
pub fn concurrency(self, concurrency: usize) -> Self {
assert!(concurrency > 0, "concurrency must be greater than 0");
Self {
concurrency,
..self
}
}
pub fn merge_k(self, merge_k: usize) -> Self {
assert!(merge_k >= 2, "merge_k must not be less than 2");
Self { merge_k, ..self }
}
}
pub fn sort<K, E>(
source: impl Iterator<Item = std::result::Result<(K, Vec<u8>), E>> + Send,
config: SortConfig,
) -> SortedIter<K, E>
where
K: Ord + Pod + Copy + Send + Sync,
E: Error + Send + 'static,
{
let (output_tx, output_rx) = bounded(config.concurrency * 16);
let chunk_dir = match FileChunkDir::<K>::new() {
Ok(chunk_dir) => Arc::new(chunk_dir),
Err(e) => {
let _ = output_tx.send(Err(e));
return SortedIter::new(output_rx, None);
}
};
let (file_chunk_tx, file_chunk_rx) = unbounded();
let chunk_dir = chunk_dir.clone();
{
let chunk_dir = chunk_dir.clone();
rayon::ThreadPoolBuilder::new()
.num_threads(config.concurrency + 1)
.build()
.unwrap()
.install(|| {
start_sorting_stage(&config, source, chunk_dir.clone(), file_chunk_tx);
start_merging_stage(&config, file_chunk_rx, chunk_dir.clone(), output_tx);
});
}
SortedIter::new(output_rx, Some(chunk_dir))
}
pub struct SortedIter<K: Pod, E> {
output_rx: Receiver<Result<(K, Vec<u8>), E>>,
done: bool,
#[allow(dead_code)]
chunk_dir: Option<Arc<FileChunkDir<K>>>,
}
impl<K: Pod, E> SortedIter<K, E> {
fn new(
output_rx: Receiver<Result<(K, Vec<u8>), E>>,
chunk_dir: Option<Arc<FileChunkDir<K>>>,
) -> Self {
SortedIter {
output_rx,
chunk_dir,
done: false,
}
}
}
impl<K: Pod, E> Iterator for SortedIter<K, E> {
type Item = Result<(K, Vec<u8>), E>;
fn next(&mut self) -> Option<Self::Item> {
if self.done {
return None;
}
match self.output_rx.recv() {
Ok(Ok(v)) => Some(Ok(v)),
Ok(Err(e)) => {
self.done = true;
Some(Err(e))
}
Err(_) => {
self.done = true;
None
}
}
}
}
fn mem_to_file_chunk<K: Pod + Ord, E>(
buffer: Vec<(K, Vec<u8>)>,
chunk_dir: Arc<FileChunkDir<K>>,
) -> Result<FileChunk<K>, E> {
let mem_chunk = MemChunk::from_unsorted(buffer);
let mut file_chunk = chunk_dir.add_chunk()?;
mem_chunk.write_to_file(&mut file_chunk)?;
Ok(file_chunk.finalize())
}
fn start_sorting_stage<K, E>(
config: &SortConfig,
source: impl Iterator<Item = std::result::Result<(K, Vec<u8>), E>> + Send,
chunk_dir: Arc<FileChunkDir<K>>,
chunk_tx: Sender<Result<FileChunk<K>, E>>,
) where
K: Ord + Pod + Copy + Send + Sync,
E: Send + 'static,
{
debug!("Sorting stage started.");
let item_header_size = std::mem::size_of::<Vec<u8>>();
let mut chunk_size = 0;
let mut buffer = Vec::new();
for res in source {
match res {
Ok((key, value)) => {
let item_size = item_header_size + value.len();
if chunk_size + item_size >= config.max_chunk_bytes {
let buffer = std::mem::take(&mut buffer);
let chunk_dir = chunk_dir.clone();
let chunk_tx = chunk_tx.clone();
if let Err(e) = config.ensure_not_canceled() {
let _ = chunk_tx.send(Err(e));
return;
}
rayon::spawn(move || {
let _ = chunk_tx.send(mem_to_file_chunk(buffer, chunk_dir));
});
chunk_size = 0;
}
chunk_size += item_size;
buffer.push((key, value));
}
Err(e) => {
let _ = chunk_tx.send(Err(crate::Error::Source(e)));
}
}
}
if !buffer.is_empty() {
if let Err(e) = config.ensure_not_canceled() {
let _ = chunk_tx.send(Err(e));
return;
}
rayon::spawn(move || {
let _ = chunk_tx.send(mem_to_file_chunk(buffer, chunk_dir));
});
}
}
fn start_merging_stage<K, E>(
config: &SortConfig,
chunk_rx: Receiver<Result<FileChunk<K>, E>>,
chunk_dir: Arc<FileChunkDir<K>>,
output_tx: Sender<Result<(K, Vec<u8>), E>>,
) where
K: Ord + Pod + Copy + Send + Sync,
E: Send + 'static,
{
debug!("Merging stage started.");
let (merged_tx, merged_rx) = unbounded::<Result<FileChunk<K>, E>>();
let mut pending = Vec::new();
let mut source_finished = false;
let mut num_running_merges = 0;
let mut recv_select = Select::new();
recv_select.recv(&chunk_rx); recv_select.recv(&merged_rx);
loop {
let idx = recv_select.ready();
match idx {
0 => match chunk_rx.try_recv() {
Ok(Ok(chunk)) => {
debug!("Received chunk: items={}", chunk.len());
pending.push(chunk)
}
Ok(Err(e)) => {
let _ = output_tx.send(Err(e));
break;
}
Err(_) => {
debug!("All chunks received from the sorting stage");
source_finished = true;
recv_select.remove(0);
}
},
1 => match merged_rx.try_recv() {
Ok(Ok(chunk)) => {
debug!("Received merged chunk: items={}", chunk.len());
num_running_merges -= 1;
pending.push(chunk)
}
Ok(Err(e)) => {
let _ = output_tx.send(Err(e));
break;
}
Err(_) => {
panic!("merged_rx should not be closed at this point")
}
},
_ => unreachable!(),
}
if let Err(e) = config.ensure_not_canceled() {
let _ = output_tx.send(Err(e));
}
let total_chunks = pending.len() + num_running_merges;
let num_merge = if source_finished {
if pending.len() > config.merge_k {
(total_chunks - config.merge_k + 1).min(config.merge_k)
} else if num_running_merges == 0 {
break;
} else {
continue;
}
} else if total_chunks >= config.merge_k * 2 - 1 {
if pending.len() >= config.merge_k {
config.merge_k
} else {
continue;
}
} else {
continue;
};
pending.sort_by_key(|chunk| chunk.len());
let remaining = pending.split_off(num_merge.min(pending.len()));
let merging = std::mem::replace(&mut pending, remaining);
let merged_tx = merged_tx.clone();
let mut chunk_writer = match chunk_dir.add_chunk() {
Ok(chunk_writer) => chunk_writer,
Err(e) => {
let _ = output_tx.send(Err(e));
break;
}
};
debug!("Start merging {} chunks", merging.len());
num_running_merges += 1;
let canceled = config.canceled.clone();
rayon::spawn(move || {
match merge_chunks_with_binary_heap(canceled, merging, |(key, value)| {
chunk_writer.push(&key, &value)
}) {
Ok(()) => {
let _ = merged_tx.send(Ok(chunk_writer.finalize()));
}
Err(e) => {
let _ = merged_tx.send(Err(e));
}
}
});
}
let canceled = config.canceled.clone();
if canceled.load(std::sync::atomic::Ordering::Relaxed) {
return;
}
debug!("Start iteration by merging {} chunks", pending.len());
rayon::spawn(move || {
if let Err(e) = merge_chunks_with_binary_heap(canceled, pending, |(key, value)| {
let _ = output_tx.send(Ok((key, value)));
Ok(())
}) {
let _ = merged_tx.send(Err(e));
}
drop(chunk_dir);
});
}