Skip to content

perf: push cancellation from Rust and register handler token sources - #82

Draft
jmbryan4 wants to merge 8 commits into
mainfrom
perf/push-cancel-cts-pool
Draft

perf: push cancellation from Rust and register handler token sources#82
jmbryan4 wants to merge 8 commits into
mainfrom
perf/push-cancel-cts-pool

Conversation

@jmbryan4

@jmbryan4 jmbryan4 commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

Experimentation on cancellation changes

Replace the pull cancellation bridge with a push path. Rust spawns a
watcher task per handler call. The watcher races on_cancel() against a
stop signal and pushes cancel() to a C# callback. This removes the
per-message TaskCompletionSource, Task.WhenAny, monitor task, and the
CA2025 pragma.

Pool the CancellationTokenSource per handler call. Each slot carries a
rental epoch. A stale push checks the epoch first, so it cannot cancel
a later rental. Rent() retires a slot that a stale push cancelled while
it sat in the pool, so a handler never starts with a cancelled token.

Deleted tests and their surviving stronger tests:

- BridgeCancellationCancelsCtsWhenOnCancelCompletes ->
  HandleMessageCancelsTokenWhileHandlerIsRunning
- BridgeCancellationDoesNotCancelCtsWhenHandlerCompletesFirst ->
  CancelPushAfterHandlerCompletedIsANoOp and
  StaleEpochCannotCancelReRentedSlot
- BridgeCancellationSwallowsSynchronousOnCancelFault ->
  HandleMessageCompletesWhenWatchRegistrationFaults
- BridgeCancellationSwallowsLateFaultWhenHandlerCompletesFirst is
  deleted without replacement. The hazard is gone: cancel() is a
  synchronous call from a Rust task, so no C# Task can fault late.
@jmbryan4
jmbryan4 marked this pull request as draft August 5, 2026 05:22
Give every handler a unique ID. Hold the slot gate across the ID check, cancellation, and reset. This makes stale callbacks no-ops and keeps TryReset away from concurrent cancellation. Run cancellation in the Rust handler task. Remove the watcher task, callback trait, and fixed primitive pool capacity. Size each pool from client concurrency. Track regenerated bindings.

StaleEpochCannotCancelReRentedSlot and StaleCancelWhilePooledDoesNotLeakACancelledToken now use StaleHandlerIdCannotCancelReRentedSlot. CancelledSourceIsDisposedNotPooled now uses CancelledSourceIsReplacedBeforeNextRental. PoolRetainsAtMostCapacity now uses PoolRetainsConfiguredCapacity. RentedSourceIsNeverCancelledUnderAnyOpSequence now uses StaleHandlerIdCannotCancelReRentedSlot and ReturnWaitsForCancellationBeforeReset. CancelPushAfterHandlerCompletedIsANoOp now uses CancelAfterHandlerCompletedIsANoOp. HandleMessageCompletesWhenWatchRegistrationFaults and HandleTimerCompletesWhenWatchStopFaults now use HandleMessageCancelsTokenWhileHandlerIsRunning and HandleTimerCancelsTokenWhileHandlerIsRunning.
Every build path runs bindgen, so the committed copy is never consumed.
Restore the .gitignore entry and untrack ProsodyFfi.cs. Update the
Definition of Done to match.
The generated shim starts handlers with Task.Run, so Rust can call
cancel(handler_id) before the handler rents its cancellation slot. The
bridge now probes should_cancel once after the rent and pre-cancels its
own token, so an early cancel is never lost.

CancelIfCurrent no longer cancels under the slot gate. It marks the
source cancel-pending and queues the cancel to the thread pool, so user
token callbacks never run on the Rust runtime thread or under the lock.
Return retires a cancel-pending source instead of resetting it, so a
late cancel cannot reach the next renter. Faults from token callbacks
are logged, not thrown across the FFI.

The FFI client now resolves max_concurrency once (options, then
PROSODY_MAX_CONCURRENCY, then the scheduler default) and exposes it as
max_concurrency(). The C# client sizes the cancellation pool from that
value. This deletes the duplicated resolution in ClientOptions. Timer
and message handlers share the same scheduler permits upstream, so the
pool capacity equals the true concurrency bound.

Also extract the duplicated select dance into CsHandler::invoke, remove
the unused CsCheck package, and correct the false ordering comment.

Deleted tests and their surviving stronger tests:
- ReturnWaitsForCancellationBeforeReset and
  CancelledSourceIsReplacedBeforeNextRental are covered by
  CancelPendingSourceIsRetiredOnReturn and
  SynchronouslyCancelledSourceIsReplacedOnReturn.
- ResolveMaxConcurrencyUsesConfiguredValue and
  ResolveMaxConcurrencyRejectsZero die with the deleted code; the
  native scheduler validates the bound.
The node reported its internal Docker IP to clients, so host-run tests
failed after driver peer discovery. Broadcast 127.0.0.1 instead. The
prosody repo compose file carries the same change.
@jmbryan4

jmbryan4 commented Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

Benchmark findings

Three mechanisms were benchmarked :

  1. current existing design on main (per-invocation CancellationTokenSource + TaskCompletionSource + Task.WhenAny monitor)
  2. push cancellation with a fresh CancellationTokenSource registered in a ConcurrentDictionary
  3. push cancellation with pooled sources

BenchmarkDotNet — one handler invocation, mechanism only (.NET 10, Arm64)

Mechanism Mean Allocated vs current
Current (main): CTS + TCS + WhenAny monitor 1,449 ns 832 B baseline
Push + fresh CTS in dictionary 31 ns 96 B 47x faster, -89% alloc
Push + pooled CTS 40 ns 0 B 37x faster, -100% alloc

The benchmark simulates OnCancel() with a plain TaskCompletionSource. The real call crosses the FFI boundary and allocates uniffi PollFuture machinery per poll.

dotnet-counters — 32 concurrent workers, 25 s sustained load

Mechanism Throughput Alloc rate GC pause Lock contentions/s
Current (main) 3.09 M ops/s 2,689 MB/s 178 ms/s ~1
Push + dictionary 3.95 M ops/s 1,211 MB/s 99 ms/s 3,055
Push + pool 2.05 M ops/s 449 MB/s 32 ms/s 7,021

Conclusions

  1. Push cancellation delivers almost the entire win: 47x faster and -89% allocation against main, before the removed FFI machinery is counted.
  2. The pool saves only the last 96 B per operation. In exchange it is slower single-threaded, and its single gate cuts throughput 48% under heavy contention with ~9x the per-operation lock contention.
  3. The pool's GC advantage per operation (16 ms pause per million ops against 25 ms for the dictionary) is invisible at real Kafka handler rates. The load harness runs about 1000x hotter than any real workload.
  4. The pool is also where the race complexity lives: slot reuse, cancel-pending state, TryReset, abandon-to-GC, and the cross-boundary max_concurrency invariant. The dictionary variant has no slot reuse, so the stale-cancel-to-next-renter race class cannot exist.

…registry

Benchmarks on PR #82 show the push-cancellation path carries the whole win. The pool saved only 96 B per invocation, ran slower than a fresh source per invocation, and serialized handlers on one gate under load. It also carried the slot-reuse race machinery and a cross-boundary concurrency invariant.

Each invocation now registers a fresh CancellationTokenSource in a ConcurrentDictionary keyed by handler ID and removes it on completion. Sources are never reused, so a stale cancel can only reach an abandoned source. The native max_concurrency report and the pool-sizing plumbing are gone.

Deleted tests and their surviving coverage:
- CancellationTokenSourcePoolTests.CancelReachesTheActiveRenter -> CancellationRegistryTests.CancelReachesTheRegisteredHandler.
- StaleHandlerIdCannotCancelReRentedSlot, CancelPendingSourceIsRetiredOnReturn, SynchronouslyCancelledSourceIsReplacedOnReturn -> LateCancelCannotReachALaterHandler and RegisterReturnsAFreshUncancelledSource; without reuse the stale-cancel class is structural, not defended.
- ReturnedHandlerIdCanRentAgain -> CompletedHandlerIdCanRegisterAgain.
- PoolRetainsConfiguredCapacity, PoolRejectsCapacityAboveTheAllocationLimit -> removed with the capacity invariant; SecondRegistrationForAnActiveHandlerThrows keeps the identity guard.
- ClientOptionsValidatorTests.MaxConcurrencyOutsidePoolCapacityFails -> MaxConcurrencyOfZeroFails; the upper bound was pool-only.
@jmbryan4 jmbryan4 changed the title perf: push cancellation from Rust and pool handler token sources perf: push cancellation from Rust and register handler token sources Aug 5, 2026
@jmbryan4 jmbryan4 changed the title perf: push cancellation from Rust and register handler token sources push cancellation from Rust and register handler token sources Aug 6, 2026
@jmbryan4 jmbryan4 changed the title push cancellation from Rust and register handler token sources perf: push cancellation from Rust and register handler token sources Aug 6, 2026
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