use std::collections::BTreeMap;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use crate::interpreter::ExecResult;
pub struct JobTable {
jobs: BTreeMap<usize, JoinHandle<ExecResult>>,
next_id: usize,
last_job_id: Option<usize>,
}
impl Default for JobTable {
fn default() -> Self {
Self::new()
}
}
impl JobTable {
pub fn new() -> Self {
Self {
jobs: BTreeMap::new(),
next_id: 1,
last_job_id: None,
}
}
pub fn spawn(&mut self, handle: JoinHandle<ExecResult>) -> usize {
let id = self.next_id;
self.next_id += 1;
self.jobs.insert(id, handle);
self.last_job_id = Some(id);
id
}
pub fn last_job_id(&self) -> Option<usize> {
self.last_job_id
}
pub async fn wait_for(&mut self, job_id: usize) -> Option<ExecResult> {
if let Some(handle) = self.jobs.remove(&job_id) {
match handle.await {
Ok(result) => Some(result),
Err(_) => Some(ExecResult::err("job panicked".to_string(), 1)),
}
} else {
None
}
}
#[allow(dead_code)]
pub async fn wait_all(&mut self) -> i32 {
self.wait_all_results()
.await
.last()
.map(|r| r.exit_code)
.unwrap_or(0)
}
pub async fn wait_all_results(&mut self) -> Vec<ExecResult> {
let jobs: Vec<_> = std::mem::take(&mut self.jobs).into_iter().collect();
let mut results = Vec::new();
for (_, handle) in jobs {
match handle.await {
Ok(result) => results.push(result),
Err(_) => results.push(ExecResult::err("job panicked".to_string(), 1)),
}
}
results
}
#[allow(dead_code)]
pub fn has_jobs(&self) -> bool {
!self.jobs.is_empty()
}
#[allow(dead_code)]
pub fn job_count(&self) -> usize {
self.jobs.len()
}
}
pub type SharedJobTable = Arc<Mutex<JobTable>>;
pub fn new_shared_job_table() -> SharedJobTable {
Arc::new(Mutex::new(JobTable::new()))
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_spawn_and_wait() {
let mut table = JobTable::new();
let handle = tokio::spawn(async { ExecResult::ok("hello".to_string()) });
let job_id = table.spawn(handle);
assert_eq!(job_id, 1);
assert_eq!(table.last_job_id(), Some(1));
let result = table.wait_for(job_id).await;
assert!(result.is_some());
assert_eq!(result.unwrap().exit_code, 0);
}
#[tokio::test]
async fn test_wait_all() {
let mut table = JobTable::new();
for i in 0..3 {
let handle = tokio::spawn(async move { ExecResult::ok(format!("job {}", i)) });
table.spawn(handle);
}
assert_eq!(table.job_count(), 3);
let exit_code = table.wait_all().await;
assert_eq!(exit_code, 0);
assert!(!table.has_jobs());
}
#[tokio::test]
async fn test_wait_for_nonexistent() {
let mut table = JobTable::new();
let result = table.wait_for(999).await;
assert!(result.is_none());
}
}