首页 / 资讯中心 / 文章详情

Work-Stealing调度算法:从rayon到tokio的Rust实践

Work-Stealing调度算法:从rayon到tokio的Rust实践 ★ FEATURED ARTICLE
开门见山说个事如果你在 Rust 里写过需要并行的代码你大概率已经被 rayon 的par_iter()惯坏了——一行代码把迭代器扔到多线程里跑剩下的什么都不用管。又或者你在用 tokio 写异步服务时会发现spawn出去的任务好像能自动保持负载均衡不管你的 future 怎么嵌套、怎么await系统总能把它们安排得明明白白。这背后就是 work-stealing 调度算法在干活。但有意思的是很多写了几年 Rust 的开发者对 work-stealing 的理解停留在线程闲着就去偷任务这个模糊概念上。真被问到它怎么偷偷谁的偷的时候会不会冲突为什么 Rust 生态里 rayon 和 tokio 都选了这条路就卡壳了。这篇我就把这个话题掰开揉碎从算法本身的原理讲到 Rust 生态里的工程实现再给出我自己在实际项目中踩过的坑和调优建议。1. 先从调度这件事说起为什么静态分配靠不住1.1 一个反直觉的结论任务均分不一定均衡想象一个非常典型的场景你有一个VecVecf32要对里面每一个子数组做某种耗时的数值计算。最直觉的并行方案是什么把数组切成 N 份N 等于你的 CPU 核数每个线程处理一份。这就是所谓的静态分配static partitioning。看起来公平得不得了每个线程拿同样多的数据谁也不吃亏。但实际跑起来你会发现负载不均衡才是常态。举个我踩过的例子我用这种静态切分方式去处理一组稀疏矩阵的分解结果四个核上跑出来的时间差异接近 3 倍。为什么因为每个矩阵的稀疏结构不同有的矩阵分解起来几乎不费劲有的则因为填充因子fill-in严重导致计算量暴涨。均分数据量不等于均分计算量。这还没算上系统层面的干扰——CPU 频率波动、其他进程抢占、cache 命中率差异这些都会让看起来公平的任务分配变得实际不公平。更麻烦的是你没法预先知道每个任务到底要跑多久。有些异构负载你要么提前做 profiling要么靠经验调切分权重但这两者都很脆弱。数据变了、场景变了之前调好的权重就废了。静态分配的脆弱性根源在于它把任务分发和任务执行硬绑在一起缺少动态调整能力。1.2 为什么动态任务队列是必然方向既然静态分配不行自然想到动态分配搞一个全局共享的任务队列所有线程从里面取任务。谁闲了谁来拿拿完就跑。这样计算量不均衡的任务在全局层面被平均掉了——重的任务多占点时间轻的任务少占点反正每个线程都在不停地从队列里取活干。这个方案能解决负载不均衡问题但引入了新问题竞争contention。所有线程都从同一个全局队列取任务每次入队出队都要抢一把锁。线程数少的时候还好线程数一多——尤其是现在服务器 CPU 动不动几十核——锁竞争会直接把并行的收益吃掉大半。你可能想说那用无锁队列不就行了确实可行但全局无锁队列也有瓶颈。一个队列再快它的头部只有一个所有线程都从同一个地方取元素这本身就制造了热点。更严重的是 cache 一致性问题队列头部的 cache line 会被所有核心来回争抢表现为严重的 cache line bouncing。实际测试中你甚至能看到线程数翻倍性能却掉头往下的负扩展现象。到这里核心矛盾就很清晰了静态分配无法应对负载不均衡动态分配又受制于同步开销。Work-stealing 的出现正是在这两者之间找到了一个工程上非常漂亮的平衡点。2. Work-Stealing 的核心机制人人都有一亩三分地闲了就去邻居家借活干2.1 双端队列算法的心脏Work-stealing 的基本结构其实不复杂。每个线程维护一个私有的双端队列deque这个队列只属于自己平时的入队和出队操作都发生在自己身上全程无锁。线程干完了自己队列里的活才会尝试去别人的队列里偷任务。这就像每家农户各自经营自己的田地平时各干各的不用跟邻居打招呼。只有自家的活干完了才会扛着锄头去邻居家看看有没有能帮忙的活。关键问题来了任务怎么往自己队列里放和取放是放哪一头取又是取哪一头这涉及到 deque 的三种操作push线程往自己的队列底部放入一个新任务。pop线程从自己的队列底部取出一个任务来执行。注意是底部不是顶部。steal空闲线程从别的线程队列的顶部取出一个任务。你发现没有push 和 pop 都发生在底部也就是同一端而 steal 发生在顶部另一端。这意味着什么正常干活时线程不会碰队列的顶端只有小偷才会从顶端下手。这个设计的精妙之处在于并发安全。因为一个队列只有两种并发访问模式要么是所有者线程操作底部要么是多个来袭线程操作顶部多个空闲线程可能会同时盯上同一个队列。底部的 push/pop 因为只有所有者自己在做天然不用锁顶部的 steal 多个线程可能同时来所以顶部必须用原子操作比如 compare-and-swap保证安全。底部无锁 顶部原子操作的组合拳就把锁竞争压到了最低。用生活类比来解释就是我是小店老板只从我自己的柜台最里面的货架取货pop 底部补货也放在最里面push 底部。外面的顾客要从店门口货架拿货steal 顶部。我和顾客永远不会碰到同一格货架所以大部分时候我们各干各的不会有冲突。只有我在最里面拿货、又有顾客刚好要从门口货架最前面拿货时才需要小心一点用原子操作保证两边不会乱套。2.2 偷取策略LIFO 优先从最新分裂下手说完了怎么偷再说偷谁的和偷哪个。这是 work-stealing 里最有讲究的部分之一。首先空闲线程挑一个 victim受害者队列去偷。最简单的策略是随机选一个线程的队列去尝试。为什么随机因为随机性可以保证在大量线程间负载会逐渐趋向均衡而且实现代价极低。也有一些更精细的策略考虑比如优先偷取最后被访问过的队列利用 cache 局部性但工程上随机轮询已经足够。其次偷哪个任务上面说过作者线程 pop 底部偷取者 steal 顶部。这背后其实是两个经典的调度原则在打架作者线程的行为是LIFO后进先出跟栈一样。为什么因为后产生的任务通常是被刚执行的任务分裂出来的子任务它的数据很可能还热乎着——相关的变量、中间结果还留在 CPU cache 里。先执行子任务可以提高 cache 命中率。偷取者的行为是FIFO先进先出从顶部拿最老的任务。为什么因为最老的任务往往是最顶层的大块任务代表一整棵任务树的根。偷走一个大的顶层任务相当于分走了一整棵子树降低了与 victim 递归分裂产生更多任务时的交互频率。这里有一个很重要的洞察**victim 线程留下底部任务继续执行实际上是在沿着任务树的某条路径深入而 thief 偷走顶部任务是拿走了另一整块待开发的树枝。**两者各取所需工作负载就这样在树形结构上被优雅地切分了。2.3 什么时候偷从自适应等待到触发阈值另一个容易被忽略的细节是线程何时变成 thief。线程干完自己的活立刻就去偷吗不是。反复偷不到会浪费不少 CPU 周期。工程实现里通常会有一个自适应策略先自旋等待一小段时间比如几十微秒。目的是保留这个线程在 CPU 上避免频繁进入睡眠又唤醒的开销。因为如果别的线程正忙着分裂任务很快就会有新任务产生此时偷取会很快成功。如果自旋一段时间还是偷不到就调用操作系统提供的 yield/sleep 原语真正让出 CPU降低能耗和系统资源占用。rayon 里做的稍微复杂一点有空闲线程会逐步降级从自旋到短 sleep 再到长 sleep最终进入类似于lock-in的状态。这个细节对性能影响很大——在小任务上自旋时间占比过高CPU 空转严重在大任务上睡眠时间过多又会导致负载均衡响应变慢。所以如果你看到 rayon 或者 tokio 里面有大量关于spin count的常量调优请不要觉得这是过度工程化。这正是 work-stealing 调度在微观层面精细调校的体现。3. Rust 生态里的两大代表rayon 和 tokio 是怎么落地的3.1 rayon数据并行的分治大师rayon 是 Rust 生态里最经典的 work-stealing 实现主要面向数据并行和分治divide-and-conquer场景。核心 API 是两个par_iter()处理集合的并行迭代。join(a, b)将两个闭包并行执行等待两者都完成。join(a, b)的底层机制最能体现 work-stealing 的设计。算法大致是当前线程先执行其中一个闭包a。执行完后检查另一个闭包b是否已经完成。如果b没完成那么不是直接阻塞等待而是先从当前线程的队列里 pop 出其他任务继续执行也就是帮忙来做别的事。如果队列也空了才开始去其他线程偷任务。如果偷也偷不到最后才进入阻塞等待任务完成。这段逻辑透露出的核心思想是在没有真正需要阻塞之前线程永远要让自己 有事可做。一个线程在等待子任务完成时不是傻等而是去执行其他可以推进的工作。这其实是 work-stealing 调度器免于死锁的关键——每个线程都是工作饥渴的只要有任何可执行的任务它就会去做而不是空等。rayon 还有一个设计细节值得提任务分裂是递归发生的。在par_iter里一个大集合会被递归地一分为二直到切分成足够小的子任务。每个子任务在被执行时如果发现还可以继续分裂就会把另一半 push 到自己队列的底部自己先执行其中一个半区。这种递归分裂 LIFO 优先执行的组合让任务树的生长非常自然。3.2 tokio异步任务调度的偷跑者tokio 的 work-stealing 跟 rayon 不太一样因为它的任务不是 CPU 密集型的计算任务而是一个个异步状态机future。tokio 的多线程运行时multi-threaded runtime本质上是一个 work-stealing 任务调度器每个工作线程维护一个任务队列。但 tokio 做了一个有意思的优化每个 worker 持有一小部分额外的固定插槽比如一个 LIFO slot当前 worker 最新 spawn 的任务优先往这个 slot 里放。这样做的好处是新 spawn 的任务往往和当前正在执行的任务有数据关联比如共享 Future 的状态把它放在 worker 自己的手边可以最大化缓存局部性并减少跨线程唤醒。只有当 LIFO slot 满了或者说超过某个阈值新的任务才会被 push 到这个 worker 的全局队列里其他 worker steal 的时候专门从那棵队列的尾部来取比如从全局队列的其他位置。这算是一种尽量让任务留在源 worker 本地的思路只有当本地装不下了才交出去给邻居。另外 tokio 里还有一个非常重要的配套机制yield_now()和调度器内部的公平性守卫。因为 tokio 的任务都是协作式调度的如果一个 future 里有死循环不await这个 worker 就会被卡死永远不还给调度器。这一点跟 rayon 完全不同后者是按 OS 线程抢占的不存在长任务卡死调度器的问题。3.3 它们的取舍为什么不同rayon 和 tokio 虽然底层都是 work-stealing但设计取向完全不同。rayon 是计算并行任务本身不阻塞执行完一个就下一个调度器要尽可能压榨 CPU。所以 rayon 倾向于给任务更多嵌套分裂的机会并用频繁的自旋保持 worker 热启动。tokio 是IO 异步任务会频繁地await挂起和恢复的代价极大。所以 tokio 的任务数可能极多任务生命周期很短调度器必须轻量且避免过多线程间的任务迁移这也是 LIFO slot 和本地优先策略存在的意义。表格对比如下维度rayontokio目标场景CPU 密集型计算IO 密集型异步任务粒度中等偏大递归分裂极细future 状态机本地优化主要靠 LIFO 队列LIFO slot 全局队列混合等待策略自旋到睡眠的降级任务挂起原线程可继续执行调度公平性不太敏感任务基本同权敏感防止单个任务饿死或卡死4. 工程实践从原理到调优你要注意的几件大事4.1 任务粒度是头号性能开关我自己最开始用 rayon 时踩过最大的坑就是任务割得太碎。你如果把一个几千个元素的par_iter()拆成每个元素一个任务rayon 确实也能跑但调度开销会非常惊人。每个任务都要被 push/pop/steal 一次这些操作虽然是原子操作和内存栅栏但不是免费的。你节省的计算时间全被调度成本吃掉了。经验法则每个任务的执行时间应该在微秒到几十微秒之间不要低于几微秒。对于数值计算你可以通过手动 chunk 来控制粒度。rayon 提供了一些机制比如par_chunks()或par_iter().with_min_len()在 rayon 的某些版本/特性中来保证每个工作块至少包含若干元素。这里给出一个我常用的基准测试思路用一个固定数据集逐步改变任务切分的粒度比如从每个元素一个任务到每 16 个元素一个任务再到每 128 个元素一个任务对比总耗时。你几乎总能看到一条 U 型曲线右边是任务太少导致负载不均左边是任务太多导致调度开销过大中间是甜点区。4.2 线程数不是越多越好很多新手拿到 work-stealing 调度器第一反应是把线程数设成 等于核心数 甚至 等于核心数 1理由是 IO 等待要额外线程。但在 work-stealing 里线程数和并行度强相关线程数太多多出来的线程没事干只能天天去 steal反而增加了全局原子的竞争和内存带宽负担。更好用的思路是rayon 的线程池大小默认等于逻辑 CPU 核数这个默认值在生产中已经够用了。如果你不确定可以先用默认值跑一遍再用--features rayon/num_cpus参考 num_cpus 的值做基准。并行度的收益只在 N 小于等于核心数时是线性的超过核心数之后增加的往往只是开销和调度延迟。如果某个阶段你希望给特定任务插队或者提高优先级不要试图增加 worker 数量而是考虑拆分任务或者调整嵌套层级让更重要的任务更早被 pop 出来执行。4.3 注意任务之间的数据依赖和锁Work-stealing 调度最怕的不是任务本身慢而是任务之间产生隐式依赖。比如多个任务共享一个Mutex保护的全局计数器每个任务执行时都要抢锁。你可能觉得我就加个互斥锁保护一个 usize不碍事吧实际跑起来可能发现并行度越调越低——多线程抢锁导致 cache line bouncing锁竞争让线程频繁阻塞调度器以为这个线程很忙阻塞在锁上但实际它一个任务都没推进。如果必须共享数据优先使用无锁结构比如crossbeam_epoch或chashmap或者把共享数据彻底变成每个线程私有一份最后再归并。Work-stealing 的本质是任务隔离你的数据也应该跟着任务走尽量保持任务间的独立性。一旦出现大家都去抢同一个资源的情况你的并行代码就变成了隐式串行work-stealing 调度器再怎么优化也无济于事。4.4 饥饿问题和任务分裂的失衡理论上work-stealing 不会饿死任务——因为每个空闲线程最终都会去偷任务只要有任务存在总会被执行。但有一种场景容易让某些任务延迟很高任务树极度不平衡一棵树上有一个巨大的右子树其他线程把别的小任务偷完之后大量线程都去偷这个右子树的子任务而左子树的某个后代任务反而一直没被偷到。这在一些分支定界搜索里非常常见。解决思路有两个方向在任务分裂时保持一定平衡不要把一个任务分裂成一个巨块加一个碎块而是尽量拆成大小相近的两半。引入 力度控制 的偷取策略比如在 rayon 里你可以通过scope配合手动spawn来调整任务结构在 tokio 里可以通过yield_now()让出当前执行权给其他任务机会。实际上偷到哪个任务这个问题直接影响最坏情况延迟。如果你写的是一个实时性要求较高的系统要意识到 work-stealing 提供的只是期望延迟低而不是最坏延迟有界。这是它的数学性质决定的你在设计时就要心里有数。5. 自己实现一个最小 work-stealing从零到可运行读到这里你可能觉得算法原理懂了库也会用了但有些东西不自己写一遍总觉得不踏实。我个人建议你至少用 Rust 亲手实现一个最小版的 work-stealing 调度器。不是为了在生产环境用它而是为了彻底理解它的并发控制细节。下面给出一个非常精简的核心结构示意覆盖 4 个关键部分任务表示、双端队列、worker 循环、偷取逻辑。use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::Arc; use std::thread; // 用一个 Boxdyn FnOnce() Send 表示可运行任务 type Task Boxdyn FnOnce() Send; /// 每个 worker 私有的双端队列基于 crossbeam 的 ArrayDeque /// 这里为了演示直接用了 crossbeam_deque::Worker 和 Stealer use crossbeam_deque::{Worker, Stealer, Injector}; use crossbeam_utils::CachePadded; struct WorkerHandle { worker: WorkerTask, stealer: StealerTask, // 记录偷取次数用于观察行为 steal_count: CachePaddedAtomicUsize, } /// 全局注入队列可选用于放置初始任务 struct Shared { injector: InjectorTask, } fn spawn_worker( id: usize, handles: VecArcWorkerHandle, shared: ArcShared, running: ArcAtomicBool, ) - thread::JoinHandle() { thread::spawn(move || { // 每个 worker 循环执行任务 while running.load(Ordering::Relaxed) { // 1. 先尝试从自己的队列里 pop 一个任务 if let Some(task) handles[id].worker.pop() { task(); continue; } // 2. 本地队列空了尝试从全局注入队列取任务 if let Some(task) shared.injector.steal() { task(); continue; } // 3. 本地和全局都没有开始随机偷取别人的队列 let mut stolen false; let n handles.len(); for offset in 1..n { let victim_id (id offset) % n; if let Some(task) handles[victim_id].stealer.steal() { handles[victim_id].steal_count.fetch_add(1, Ordering::Relaxed); task(); stolen true; break; } } // 4. 实在偷不到短暂自旋后再次尝试 if !stolen { thread::yield_now(); } } }) }这段代码虽然短但你可以通过它跑很多小实验。比如统计每个 worker 的steal_count对比不同任务粒度下的偷取频率。观察某些 worker 是不是饿死了从没偷到过任务。修改自旋策略对比整体耗时。我在做这个最小实现时发现一个有趣的现象**当任务粒度较小时偷取次数会很高甚至出现任务还没进队列就被偷走的情况。**这是 work-stealing 调度行为的正常一瞥——调度开销随任务数线性增长这也是为什么你始终应该控制任务切分的下限。6. 实战问题排查与性能分析记录6.1 现象一并行代码跑得比串行还慢这是我被问得最多的问题。排查思路按顺序走先确认任务粒度如果任务是几纳秒就结束的轻量操作work-stealing 调度开销可能超过并行收益。可以先手动把任务合并成较大的 chunk 再测。检查是不是内存带宽受限如果你的任务是单纯遍历一个大数组做算术运算那瓶颈可能在内存带宽不在 CPU。多线程不仅没有收益还会因为多核竞争内存带宽导致更慢。这时候用 work-stealing 没有意义不如优化访存模式比如改成 read-only 遍历或者分块处理。看是否有锁竞争用perf top或者heaptrack辅助分析如果看到大量的原子操作和 futex 调用说明你的任务间共享状态过多破坏了 work-stealing 的任务独立性。6.2 现象二线程数增加但吞吐不再提升这类情况通常是因为任务树不够深。work-stealing 的并行度依赖于任务分裂空间的充足性。如果你的任务在很浅的层级就分裂完了后面线程再多也无事可做。检查方法是在任务里加一个一次性的计数日志统计总共产生了多少个任务。如果任务数只比线程数多一点说明分裂太浅。解决方案是改变数据划分策略让任务树更深比如递归地把任务一分为二而不是一次性全部分配完毕。6.3 现象三偶尔出现任务延迟高、或者 involution 深度过大如果你发现某个特定任务延迟高到一个令人不满意的程度大概率是某个空闲线程刚好随机选到了错误的 victim或者 victim 队列的局部性不好。work-stealing 是随机化的它不保证某个任务一定在某个时间内启动。缓解方式有两类如果任务之间实际有依赖关系手动把有依赖的任务放到同一个 scope 或同一个 worker 内执行比如 rayon 的scopespawn避免依赖方等待过久。如果系统中本来就有大量短任务和一条长任务链考虑用优先级队列替代 work-stealing——但这不是本文的范畴只是提醒你 work-stealing 不是万灵丹。工作一年多、写过不少并行和异步场景之后我的最大体会是work-stealing 不是一个加了就变快的黑魔法它是一个让你在正确粒度的任务树下保持线程忙碌的调度框架。它的优势在于无锁本地操作和低竞争偷取但前提是你得把任务切分好、数据隔离好。rayon 和 tokio 已经帮你把最难的调度逻辑藏好了普通开发者真正要花心思的地方反而是那些看起来很不起眼的任务怎么切、共享什么数据、预期什么负载模型。最后分享一个小技巧当你对某个并行逻辑的调度行为没有把握时先不要优化性能把每个任务的关键路径上打一个Instant::now()的时间戳记录在共享的内存数组里。跑完后打印出每个 worker 执行了多少任务、空转了多少时间、偷取了几次。这个数据往往比任何 profiling 工具的火焰图都直观能让你一眼看出负载到底平不平衡、任务粒度到底合不合理。多来几次这样的实验你对 work-stealing 的直觉就真正建立起来了。
阅读完成 · 觉得有帮助?
咨询建站