如何在线程池中实现内存屏障,确保任务按依赖关系执行?
为线程池实现任务依赖同步的高效方案
假设我们有一个非常简单的线程池模型,仅支持添加任务和执行任务的操作(任务还可以添加新任务)。现在想实现任务同步:添加任务时可以指定该任务必须等待另一个指定任务完全执行完毕后,才能被线程执行。
我自己想到了一个可行方案,但担心性能极差:
- 为初始任务集合中的每个任务分配字符串标识符,同时为每个任务设置阻塞标识符(即它依赖的任务ID)
- 将任务存储在以标识符为键的map中
- 启动工作线程
- 线程遍历map中的元素,检查当前任务的阻塞标识符对应的任务是否仍在map中;若存在则跳过,否则将该任务从map中移除并执行
这个方案肯定能跑,但会让程序几乎变成单线程:每个线程都必须锁定整个map,然后遍历所有元素才能找到可执行的任务。我相信有更优雅高效的解决方案。
以下是我目前实现的无依赖版本代码:
pub struct TaskMessage<'a> { pub more_tasks: Vec<Task<'a>>, pub should_repeat: bool, } pub struct TaskList<'a> { tasks: VecDeque<Task<'a>>, } impl<'a> TaskList<'a> { pub fn push_task<F>(&mut self, task: F) where F: for<'b> FnMut(&'b mut TaskMessage<'a>) + Send + 'a, { self.tasks.push_back(Task(Box::new(task))); } fn pop_task(&mut self) -> Option<Task<'a>> { self.tasks.pop_front() } } pub struct ThreadPool<'a> { task_list: Arc<Mutex<TaskList<'a>>>, max_threads: usize, } fn ceil_div(a: usize, b: usize) -> usize { (a / b) + ((a % b) != 0) as usize } impl<'a> ThreadPool<'a> { pub fn new<'b>() -> Self { let thread_count = thread::available_parallelism().unwrap().get(); ThreadPool { task_list: Arc::new(Mutex::new(TaskList { tasks: VecDeque::new(), })), max_threads: thread_count, } } pub fn add_task<T>(&mut self, task: T) where T: for<'b> FnMut(&'b mut TaskMessage<'a>) -> () + Send + 'a, { self.task_list.lock().unwrap().push_task(task); } pub fn add_tasks<F: 'a, D: 'a>(&mut self, task: F, data: &'a mut [D]) where F: Clone + for<'b> FnMut(&'b mut D, &'b mut TaskMessage<'a>) + Send + 'a, D: Send + Sync + std::fmt::Debug, { let num_chunks = ceil_div(data.len(), self.max_threads); let chunk_size = ceil_div(data.len(), num_chunks); for chunk in data.chunks_mut(chunk_size) { let mut task = task.clone(); self.task_list.lock().unwrap().push_task( move |proxy: &mut TaskMessage<'a>| { for item in chunk.iter_mut() { task(item, proxy) } }, ); } } pub fn run(&mut self) { thread::scope(|s| { for _ in 0..self.max_threads { let task_clone = Arc::clone(&self.task_list); s.spawn(move || { loop { let top_task = task_clone.lock().unwrap().pop_task(); if top_task.is_none() { break; } let mut proxy = TaskMessage { more_tasks: Vec::new(), should_repeat: false, }; let mut top_task = top_task.unwrap(); top_task.0(&mut proxy); for task in proxy.more_tasks { task_clone.lock().unwrap().push_task(task.0); } if proxy.should_repeat { task_clone.lock().unwrap().push_task(top_task.0); } } }); } }); } }
内容的提问来源于stack exchange,提问作者Makogan
相关产品推荐
相关产品推荐

