use crate::settings::LiveComputeTarget;
use std::{
collections::VecDeque,
sync::{Mutex, PoisonError},
};
#[must_use]
pub fn local_lane_count() -> usize {
std::thread::available_parallelism()
.map_or(1, std::num::NonZeroUsize::get)
.saturating_sub(1)
.max(1)
}
pub struct WorkQueue<T> {
shared: Mutex<VecDeque<T>>,
local_retry: Mutex<VecDeque<T>>,
}
impl<T> WorkQueue<T> {
pub fn new(items: impl IntoIterator<Item = T>) -> Self {
Self {
shared: Mutex::new(items.into_iter().collect()),
local_retry: Mutex::new(VecDeque::new()),
}
}
pub fn claim_shared(&self) -> Option<T> {
self.shared
.lock()
.unwrap_or_else(PoisonError::into_inner)
.pop_front()
}
pub fn claim_local(&self) -> Option<T> {
let retried = self
.local_retry
.lock()
.unwrap_or_else(PoisonError::into_inner)
.pop_front();
if let Some(item) = retried {
return Some(item);
}
self.claim_shared()
}
pub fn return_to_local(&self, item: T) {
self.local_retry
.lock()
.unwrap_or_else(PoisonError::into_inner)
.push_back(item);
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct LanePlan {
pub run_local: bool,
pub run_remote: bool,
pub fallback_to_local: bool,
}
impl LanePlan {
#[must_use]
pub const fn for_target(target: LiveComputeTarget, worker_configured: bool) -> Self {
let run_remote = matches!(target, LiveComputeTarget::RemoteOnly)
|| (matches!(target, LiveComputeTarget::Both) && worker_configured);
Self {
run_local: !matches!(target, LiveComputeTarget::RemoteOnly),
run_remote,
fallback_to_local: matches!(target, LiveComputeTarget::Both),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn claim_shared_drains_fifo() {
let queue = WorkQueue::new([1, 2, 3]);
assert_eq!(queue.claim_shared(), Some(1));
assert_eq!(queue.claim_shared(), Some(2));
assert_eq!(queue.claim_shared(), Some(3));
assert_eq!(queue.claim_shared(), None);
}
#[test]
fn claim_local_prefers_the_retry_pile_over_fresh_work() {
let queue = WorkQueue::new(["fresh-1", "fresh-2"]);
queue.return_to_local("retried");
assert_eq!(queue.claim_local(), Some("retried"));
assert_eq!(queue.claim_local(), Some("fresh-1"));
assert_eq!(queue.claim_local(), Some("fresh-2"));
assert_eq!(queue.claim_local(), None);
}
#[test]
fn returned_items_are_invisible_to_claim_shared() {
let queue: WorkQueue<i32> = WorkQueue::new(std::iter::empty());
queue.return_to_local(42);
assert_eq!(queue.claim_shared(), None);
assert_eq!(queue.claim_local(), Some(42));
}
#[test]
fn local_only_never_runs_remote() {
for worker_configured in [false, true] {
let plan = LanePlan::for_target(LiveComputeTarget::LocalOnly, worker_configured);
assert!(plan.run_local);
assert!(!plan.run_remote);
assert!(!plan.fallback_to_local);
}
}
#[test]
fn remote_only_always_runs_remote_and_never_falls_back() {
for worker_configured in [false, true] {
let plan = LanePlan::for_target(LiveComputeTarget::RemoteOnly, worker_configured);
assert!(!plan.run_local);
assert!(plan.run_remote);
assert!(!plan.fallback_to_local);
}
}
#[test]
fn both_runs_remote_only_when_a_worker_is_configured() {
let with_worker = LanePlan::for_target(LiveComputeTarget::Both, true);
assert!(with_worker.run_local);
assert!(with_worker.run_remote);
assert!(with_worker.fallback_to_local);
let without_worker = LanePlan::for_target(LiveComputeTarget::Both, false);
assert!(without_worker.run_local);
assert!(!without_worker.run_remote);
assert!(without_worker.fallback_to_local);
}
#[test]
fn empty_queue_claims_are_none_on_both_paths() {
let queue: WorkQueue<i32> = WorkQueue::new(std::iter::empty());
assert_eq!(queue.claim_shared(), None);
assert_eq!(queue.claim_local(), None);
}
#[test]
fn n_concurrent_claim_shared_callers_each_get_a_disjoint_subset() {
const ITEM_COUNT: i32 = 2_000;
const LANE_COUNT: usize = 16;
let queue = WorkQueue::new(0..ITEM_COUNT);
let claimed: Mutex<Vec<i32>> = Mutex::new(Vec::with_capacity(ITEM_COUNT as usize));
std::thread::scope(|scope| {
for _ in 0..LANE_COUNT {
scope.spawn(|| {
let mut mine = Vec::new();
while let Some(item) = queue.claim_shared() {
mine.push(item);
}
claimed.lock().unwrap().extend(mine);
});
}
});
let mut claimed = claimed.into_inner().unwrap();
assert_eq!(
claimed.len(),
ITEM_COUNT as usize,
"no item lost or duplicated"
);
claimed.sort_unstable();
claimed.dedup();
assert_eq!(
claimed.len(),
ITEM_COUNT as usize,
"every item claimed exactly once, by exactly one lane"
);
}
#[test]
fn n_concurrent_claim_local_callers_never_duplicate_across_retry_and_shared_piles() {
const FRESH_COUNT: i32 = 1_000;
const RETRIED_COUNT: i32 = 500;
const LANE_COUNT: usize = 16;
let queue = WorkQueue::new(0..FRESH_COUNT);
for retried in FRESH_COUNT..(FRESH_COUNT + RETRIED_COUNT) {
queue.return_to_local(retried);
}
let claimed: Mutex<Vec<i32>> =
Mutex::new(Vec::with_capacity((FRESH_COUNT + RETRIED_COUNT) as usize));
std::thread::scope(|scope| {
for _ in 0..LANE_COUNT {
scope.spawn(|| {
let mut mine = Vec::new();
while let Some(item) = queue.claim_local() {
mine.push(item);
}
claimed.lock().unwrap().extend(mine);
});
}
});
let mut claimed = claimed.into_inner().unwrap();
let expected = (FRESH_COUNT + RETRIED_COUNT) as usize;
assert_eq!(claimed.len(), expected, "no item lost or duplicated");
claimed.sort_unstable();
claimed.dedup();
assert_eq!(
claimed.len(),
expected,
"every fresh and retried item claimed exactly once"
);
}
#[test]
fn local_lane_count_is_at_least_one_and_leaves_headroom() {
let n = local_lane_count();
assert!(n >= 1);
let available = std::thread::available_parallelism().map_or(1, std::num::NonZeroUsize::get);
assert!(n <= available.saturating_sub(1).max(1));
}
}