use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::mpsc::{Receiver, SyncSender, TryRecvError};
use super::{Error, IndexedObjects};
use crate::fetch::verified::Described;
use crate::fetch::{FetchError, check_cancelled};
use crate::{Object, ObjectId};
const IN_FLIGHT_BYTES: usize = 64 * 1024 * 1024;
fn charge(object: &Object) -> usize {
let bytes = object.data().len();
if object.kind() == crate::ObjectKind::Tree {
bytes + bytes + bytes / 4
} else {
bytes
}
}
const MIN_PARALLEL_OBJECTS: usize = 1024;
struct Job {
position: usize,
object: Object,
}
struct Done {
position: usize,
id: ObjectId,
described: Described,
bytes: usize,
}
pub(super) struct Finisher {
jobs: Option<SyncSender<Job>>,
done: Option<Receiver<Done>>,
in_flight: usize,
pending: usize,
}
pub(super) fn with_finisher<R>(
objects: usize,
max_threads: usize,
body: impl FnOnce(&mut Finisher) -> R,
) -> R {
let threads = std::thread::available_parallelism()
.map_or(1, std::num::NonZeroUsize::get)
.min(max_threads);
if threads < 2 || objects < MIN_PARALLEL_OBJECTS {
return body(&mut Finisher {
jobs: None,
done: None,
in_flight: 0,
pending: 0,
});
}
let (jobs, queue) = std::sync::mpsc::sync_channel::<Job>(threads * 4);
let (results, done) = std::sync::mpsc::channel::<Done>();
let queue = Mutex::new(queue);
std::thread::scope(|scope| {
for _ in 0..threads {
let results = results.clone();
let queue = &queue;
scope.spawn(move || {
loop {
let job = queue.lock().unwrap_or_else(|e| e.into_inner()).recv();
let Ok(Job { position, object }) = job else {
return;
};
let id = object.id();
let described = Described::new(id, &object);
let bytes = charge(&object);
drop(object);
if results
.send(Done {
position,
id,
described,
bytes,
})
.is_err()
{
return;
}
}
});
}
drop(results);
let mut finisher = Finisher {
jobs: Some(jobs),
done: Some(done),
in_flight: 0,
pending: 0,
};
let result = body(&mut finisher);
finisher.jobs = None;
finisher.done = None;
result
})
}
impl Finisher {
pub fn submit(
&mut self,
objects: &mut IndexedObjects,
position: usize,
object: Object,
cancel: &AtomicBool,
) -> Result<(), FetchError> {
if self.jobs.is_none() {
let id = object.id();
let described = Described::new(id, &object);
return merge(objects, position, id, described);
}
let bytes = charge(&object);
while self.in_flight != 0 && self.in_flight + bytes > IN_FLIGHT_BYTES {
self.drain(objects, true)?;
check_cancelled(cancel)?;
}
self.in_flight += bytes;
if let Some(jobs) = &self.jobs {
jobs.send(Job { position, object })
.map_err(|_| FetchError::Pack(Error::Corrupt("import worker stopped")))?;
self.pending += 1;
}
self.drain(objects, false)
}
pub fn submit_now(
&mut self,
objects: &mut IndexedObjects,
position: usize,
object: Object,
) -> Result<(), FetchError> {
let id = object.id();
let described = Described::new(id, &object);
merge(objects, position, id, described)
}
fn drain(&mut self, objects: &mut IndexedObjects, mut block: bool) -> Result<(), FetchError> {
let Some(done) = &self.done else {
return Ok(());
};
while self.pending != 0 {
let next = if block {
block = false;
done.recv().map_err(|_| TryRecvError::Disconnected)
} else {
done.try_recv()
};
let finished = match next {
Ok(finished) => finished,
Err(TryRecvError::Empty) => break,
Err(TryRecvError::Disconnected) => {
return Err(FetchError::Pack(Error::Corrupt("import worker stopped")));
}
};
self.pending -= 1;
self.in_flight -= finished.bytes;
merge(objects, finished.position, finished.id, finished.described)?;
}
Ok(())
}
pub fn finish(&mut self, objects: &mut IndexedObjects) -> Result<(), FetchError> {
while self.pending != 0 {
self.drain(objects, true)?;
}
Ok(())
}
}
fn merge(
objects: &mut IndexedObjects,
position: usize,
id: ObjectId,
described: Described,
) -> Result<(), FetchError> {
if objects.identities.insert(id, position).is_some() {
return Err(Error::Corrupt("duplicate received identity").into());
}
if let Some(resolved) = &mut objects.entries[position].resolved {
resolved.0 = id;
}
objects.verified.record(id, described);
Ok(())
}
#[cfg(test)]
mod tests;