Skip to content

perf(waitgroup): streamline completion for participants and observers - #395

Closed
tisonkun wants to merge 4 commits into
apache:mainfrom
tisonkun:codex/waitgroup-completion
Closed

tisonkun wants to merge 4 commits into
apache:mainfrom
tisonkun:codex/waitgroup-completion

Conversation

@tisonkun

@tisonkun tisonkun commented Nov 3, 2026 •

Copy link
Copy Markdown
Member

Summary

Speed up one-shot wait group completion. Participants share one allocation reference, so cloning and dropping a worker updates only the participant count. The first pending observer registers an inline waker. Additional observers use the existing cancellable waker set.

Extend the fixed-workload comparisons with late waits, nested workers, four scheduled tasks, and four independent observers of the same drain. The drain case requires each observer to be notified and reach Ready. It exercises the completion side of a graceful shutdown.

Design Notes

A participant can finish by dropping its handle or can finish and wait for its peers by consuming the handle. Cloning a Wait observes completion without adding work. This distinction lets ShutdownGuard keep the drain pending while cloned Shutdown handles independently await it. A supervisor interested in a single batch can simply give workers handles to drop and await its own.

Go also wakes multiple waiters at zero. Its Wait only observes, work is counted separately, and a new generation can start after earlier waits return. Asyncband follows the one-shot participant model also used by Crossbeam. Go WaitGroup · Crossbeam WaitGroup

Participants collectively hold one reference through final notification. Each observer owns a separate reference and the last release reclaims the allocation. The inline waker and overflow registrations share a single standard-library write lock. Poll and cancellation recheck the participant count under that lock. The last participant detaches all wakers under the same lock and wakes them after unlocking. This eliminates the unsafe waker cell and the atomic handoff protocol. Unsafe code remains in the allocation ownership implementation.

Benchmarks

Apple M4 Max, Darwin arm64, rustc 1.101.0-nightly (8006091eb), Cargo bench profile. Each median is elapsed milliseconds for ten equal-work repetitions with rotating adapter order. The revised head is 540858b. Base and first draft use the same benchmark code as the revision under review. Slashes denote separate launches. Compare peers within the same launch.

Scenario Rounds Base 36ba6f2 First draft 1912102 Revised PR waitgroup TaskTracker
complete-waiter 3,500,000 306 91 96 / 93 100 / 101 —
complete-4-workers 2,600,000 302 123 136 / 135 127 / 127 —
wait-after-work 5,500,000 330 120 125 / 114 130 / 126 —
nested-4-workers 3,000,000 316 107 93 / 93 116 / 114 —
tasks-4-workers-4w 25,000 214 212 210 / 210 212 / 213 —
drain-4-observers 2,500,000 294 372 230 / 221 1042 / 947 435 / 426

The four-observer row requires four independent wakeups. It compares native completion futures to waitgroup plus FutureExt::shared and to a closed Tokio TaskTracker. The scheduled row includes spawning and joining four tasks and should be interpreted separately from the hand-polled workloads.

Why the first draft regressed on four observers

The first observer avoided the mutex and arena, but the second observer initialized a boxed registry, a pthread mutex on this Darwin toolchain, and arena storage. A counting allocator observed one allocation at setup and three during overflow registration. The revised registry has one setup allocation and two during overflow registration. The original main implementation made one at setup and two at the first registration.

A separate 64-group batched diagnostic, in forward and reverse process order, placed the largest costs in registration and cleanup. In nanoseconds per group, the first draft spent 3.7–3.8 on the first observer, 55–56 on registering the next three and 34–35 on cleanup. The revised registry spent about 4, 32–36 and 14–16 respectively. The revised notification phase took 24–27 ns against 17–18 ns for the first draft. These phase timings are diagnostic and do not reproduce the allocation locality of the official one-group-per-round benchmark.

In a separate arm64 Linux VM control (rustc 1.95.0), a standalone driver repeated the same four-observer lifecycle and rotated all three Asyncband implementations. The ten-run medians were 159 ms for main, 182.5 ms for the first draft and 164 ms for the revised PR. This closed most of the gap on Linux; the revised implementation remained about 3% above main in that control.

Raw samples

Base

waitgroup/complete-waiter  (3500000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               306        293..318     1.00x  [318, 306, 314, 302, 299, 300, 293, 305, 312, 314]
  WaitgroupRs              98         97..102     0.32x  [99, 102, 98, 98, 97, 97, 97, 97, 98, 99]

waitgroup/complete-4-workers  (2600000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               302        294..310     1.00x  [310, 309, 304, 299, 304, 298, 294, 299, 307, 295]
  WaitgroupRs             123        118..126     0.41x  [122, 118, 123, 120, 124, 121, 122, 125, 123, 126]

waitgroup/wait-after-work  (5500000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               330        321..337     1.00x  [328, 333, 332, 332, 326, 337, 321, 325, 334, 326]
  WaitgroupRs             120        118..127     0.37x  [119, 120, 120, 121, 123, 118, 124, 127, 121, 120]

waitgroup/nested-4-workers  (3000000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               316        306..322     1.00x  [306, 321, 318, 319, 315, 322, 314, 319, 313, 308]
  WaitgroupRs              93          93..97     0.30x  [93, 94, 93, 93, 97, 93, 94, 94, 95, 93]

waitgroup/tasks-4-workers-4w  (25000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               214        208..228     1.00x  [222, 209, 211, 222, 208, 212, 221, 214, 214, 228]
  WaitgroupRs             207        200..230     0.97x  [200, 200, 208, 207, 206, 222, 208, 225, 230, 207]

waitgroup/drain-4-observers  (2500000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               294        286..306     1.00x  [290, 306, 286, 292, 294, 295, 295, 295, 286, 303]
  WaitgroupRsShared        796        772..802     2.70x  [791, 790, 802, 796, 778, 796, 800, 798, 772, 796]
  TokioTaskTracker        434        424..438     1.47x  [434, 438, 429, 436, 436, 428, 434, 435, 424, 434]

First draft

waitgroup/complete-waiter  (3500000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband                91         87..105     1.00x  [105, 89, 87, 94, 89, 97, 92, 88, 94, 88]
  WaitgroupRs             100         98..107     1.10x  [100, 99, 99, 104, 98, 101, 100, 99, 107, 101]

waitgroup/complete-4-workers  (2600000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               123        120..128     1.00x  [123, 120, 121, 122, 123, 124, 124, 120, 125, 128]
  WaitgroupRs             129        127..135     1.05x  [132, 127, 131, 129, 135, 127, 129, 130, 129, 129]

waitgroup/wait-after-work  (5500000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               120        117..126     1.00x  [120, 120, 121, 120, 119, 124, 126, 117, 117, 120]
  WaitgroupRs             129        126..132     1.07x  [132, 129, 129, 132, 126, 127, 131, 127, 126, 130]

waitgroup/nested-4-workers  (3000000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               107        105..115     1.00x  [107, 111, 108, 107, 105, 107, 105, 111, 106, 115]
  WaitgroupRs             118        113..123     1.10x  [118, 120, 118, 115, 121, 121, 118, 123, 113, 115]

waitgroup/tasks-4-workers-4w  (25000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               212        200..222     1.00x  [216, 202, 222, 213, 211, 206, 207, 218, 212, 200]
  WaitgroupRs             211        206..221     1.00x  [206, 221, 207, 209, 207, 215, 211, 214, 213, 210]

waitgroup/drain-4-observers  (2500000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               372        365..389     1.00x  [365, 377, 371, 371, 372, 374, 367, 389, 372, 367]
  WaitgroupRsShared        831        822..854     2.23x  [822, 832, 827, 827, 830, 827, 850, 837, 854, 844]
  TokioTaskTracker        430        419..459     1.16x  [429, 419, 420, 425, 422, 459, 431, 447, 435, 441]

Revised PR, first launch

waitgroup/complete-waiter  (3500000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband                96          92..99     1.00x  [96, 99, 93, 96, 99, 95, 99, 94, 96, 92]
  WaitgroupRs             100         98..105     1.04x  [99, 105, 98, 98, 100, 100, 101, 100, 104, 100]

waitgroup/complete-4-workers  (2600000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               136        130..140     1.00x  [131, 130, 133, 137, 140, 139, 135, 140, 131, 137]
  WaitgroupRs             127        125..139     0.94x  [139, 126, 133, 126, 128, 127, 125, 130, 126, 129]

waitgroup/wait-after-work  (5500000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               125        122..127     1.00x  [126, 123, 122, 126, 124, 127, 123, 125, 125, 125]
  WaitgroupRs             130        125..136     1.04x  [126, 127, 125, 131, 128, 128, 133, 136, 134, 134]

waitgroup/nested-4-workers  (3000000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband                93         91..100     1.00x  [94, 91, 93, 92, 93, 93, 92, 100, 93, 99]
  WaitgroupRs             116        113..124     1.25x  [124, 116, 121, 117, 117, 115, 113, 118, 113, 115]

waitgroup/tasks-4-workers-4w  (25000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               210        205..222     1.00x  [210, 211, 205, 209, 221, 210, 218, 209, 222, 212]
  WaitgroupRs             212        206..220     1.01x  [208, 216, 217, 209, 207, 220, 218, 206, 208, 216]

waitgroup/drain-4-observers  (2500000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               230        225..238     1.00x  [227, 225, 238, 228, 232, 230, 226, 234, 231, 234]
  WaitgroupRsShared       1042       969..1133     4.52x  [985, 992, 986, 993, 969, 1127, 1133, 1090, 1130, 1125]
  TokioTaskTracker        435        423..458     1.88x  [423, 432, 435, 434, 430, 434, 458, 457, 453, 453]

Revised PR, second launch

waitgroup/complete-waiter  (3500000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband                93          91..99     1.00x  [98, 91, 94, 91, 92, 96, 92, 94, 92, 99]
  WaitgroupRs             101         98..107     1.09x  [99, 98, 104, 99, 104, 107, 101, 104, 100, 101]

waitgroup/complete-4-workers  (2600000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               135        131..139     1.00x  [134, 138, 131, 132, 134, 131, 137, 137, 136, 139]
  WaitgroupRs             127        125..135     0.94x  [126, 127, 131, 135, 128, 127, 126, 125, 127, 128]

waitgroup/wait-after-work  (5500000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               114        111..117     1.00x  [117, 113, 113, 115, 117, 115, 114, 113, 112, 111]
  WaitgroupRs             126        125..135     1.11x  [126, 130, 125, 135, 131, 126, 125, 127, 126, 126]

waitgroup/nested-4-workers  (3000000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband                93          91..99     1.00x  [92, 91, 93, 99, 94, 97, 99, 93, 91, 99]
  WaitgroupRs             114        113..118     1.22x  [114, 115, 113, 118, 114, 113, 117, 115, 113, 115]

waitgroup/tasks-4-workers-4w  (25000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               210        204..221     1.00x  [215, 206, 210, 204, 204, 206, 221, 214, 212, 210]
  WaitgroupRs             213        206..222     1.02x  [210, 217, 209, 217, 206, 219, 209, 215, 222, 211]

waitgroup/drain-4-observers  (2500000 rounds, 10 runs; reference: Asyncband)
  implementation    median ms     min..max ms    vs ref
  Asyncband               221        216..228     1.00x  [217, 222, 220, 216, 220, 228, 223, 222, 220, 226]
  WaitgroupRsShared        947        932..959     4.28x  [959, 956, 949, 958, 948, 944, 943, 944, 946, 932]
  TokioTaskTracker        426        417..444     1.93x  [432, 427, 430, 425, 433, 425, 417, 423, 444, 418]

Validation and review

Workspace lint, tests and the benchmark smoke suite passed. On 540858b the WaitGroup Miri target passed on all eight seeds with strict provenance, symbolic alignment checks and a preemption rate of 0.25. Each seed ran 14 tests. All 11 targets selected by cargo x miri passed locally with explicit --target x86_64-unknown-linux-gnu, strict provenance, symbolic alignment checks and the four CI seeds. The OS-backed Tokio integration case runs in the native suite. The full native macOS Miri workflow passed WaitGroup before encountering unsupported kqueue in the existing broadcast tests. The new test drops the last observer from inside its wake callback and exercises the participant anchor through notification.

MIRIFLAGS='-Zmiri-strict-provenance -Zmiri-symbolic-alignment-check -Zmiri-preemption-rate=0.25 -Zmiri-many-seeds=0..8' cargo +nightly miri test -p tests-integration --test waitgroup_test

Kimi Code independently reviewed the waiter registry, reference ownership, memory ordering, reentrant cancellation/wake and benchmark fairness. It reported no blocking finding on e01ccff. A registration already detached by completion can still receive a wake after a Ready poll. That legal spurious wake does not justify another lock or a generation counter. A second Kimi review of the final 540858b diff found no blocking issue. It verified the explicit count acquisition on every pointer duplication and the reentrant last-observer test.

The full GitHub Actions Linux Miri job passed on 540858b in 5m17s. Miri CI on this revision

Reproduction

cargo x bench --bench ecosystem -- waitgroup/ --repeat 10 --raw

@tisonkun tisonkun closed this Nov 3, 2026
@tisonkun

tisonkun commented Nov 3, 2026

Copy link
Copy Markdown
Member Author

Release the lock - anyone can pick up this performance improvement.

@tisonkun
tisonkun deleted the codex/waitgroup-completion branch November 3, 2026 11:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant