diff --git a/include/async/cancellation.hpp b/include/async/cancellation.hpp index f2b1306..db6030c 100644 --- a/include/async/cancellation.hpp +++ b/include/async/cancellation.hpp @@ -293,29 +293,52 @@ struct cancellation_resolver final : private abstract_cancellation_callback { if (!(st & done_listen) || !(st & done_completion_path)) return; - if (!(st & done_cancellation_path)) { - // Try to unregister from the cancellation event once both listen() and complete() were called. - // Note that we enter this code path at most once since done_cancellation_path is the only missing - // bit in state_ and the next call to transition_() will necessarily set it. - assert(event_); - - frg::unique_lock guard{event_->_mutex}; - if (event_->_was_requested) - return; - auto it = event_->_cbs.iterator_to(this); - event_->_cbs.erase(it); + if (!(old_st & done_listen) || !(old_st & done_completion_path)) { + // If at least one of done_listen or done_completion_path was unset before, + // we perform the unregistration step. This happens exactly once. + unsigned int bits = done_unregister; + if (!(st & done_cancellation_path)) { + // Try to unregister from the cancellation event. + assert(event_); + + frg::unique_lock guard{event_->_mutex}; + if (!event_->_was_requested) { + auto it = event_->_cbs.iterator_to(this); + event_->_cbs.erase(it); + // call() can never run anymore. + bits |= done_cancellation_path; + } + } + // Set done_unregister (and potentially done_cancellation_path) last + // as accessing *this may be UB afterwards (see below). + st = state_.fetch_or(bits, std::memory_order_acq_rel) | bits; + } else { + // Both done_listen and done_completion_path were already set before. + assert(new_bits == done_cancellation_path); } - // Call resume() when all code paths are done. This can only happen once. + // Call resume() once all code paths are done and unregistration is finished (i.e., all four bits are set). + // If the cancellation path and unregistering race, both of them perform an RMW on state_ + // and exactly one of them observes the other's bit. resume() is hence called exactly once. + // If we are not the code path that calls resume_() here, + // accessing *this is UB as resume_() may concurrently destroy *this. + if (!(st & done_cancellation_path) || !(st & done_unregister)) + return; resume_(this); } + // Note: resume() is called once all four of the bits below are set. + // Hence, *this also only remains valid until this happens. + // All code paths that access *this must guarantee that at least one bit is still unset. + // Set in state_ when listen() is done. static constexpr unsigned int done_listen = 1u << 0; // Set in state_ when complete() is done or when we know that it will never be called. static constexpr unsigned int done_completion_path = 1u << 1; // Set in state_ when tryCancel() is done or when we know that it will never be called. static constexpr unsigned int done_cancellation_path = 1u << 2; + // Set in state_ when unregistration from the cancellation event is done or when we know that it is unnecessary. + static constexpr unsigned int done_unregister = 1u << 3; cancellation_event *event_{nullptr}; std::atomic state_{0}; diff --git a/tests/post-ack.cpp b/tests/post-ack.cpp index 97ddb1e..b9797f2 100644 --- a/tests/post-ack.cpp +++ b/tests/post-ack.cpp @@ -1,4 +1,6 @@ +#include #include +#include #include #include @@ -74,3 +76,55 @@ TEST(PostAck, ImmovableType) { ASSERT_EQ(ok1_ctr, 8); ASSERT_EQ(ok2_ctr, 8); } + +// Stress test for the race between completion and cancellation of a poll operation; +// resuming (and thereby destructing) the operation must not overlap with the racing +// thread's accesses inside cancellation_resolver (cf. managarm/managarm#1509). +// Best run under ASan/TSan to detect regressions. +TEST(PostAck, PollCancelRace) { + for (int round = 0; round < 10000; ++round) { + async::post_ack_mechanism mech; + async::post_ack_agent agent; + agent.attach(&mech); + async::cancellation_event ce; + + std::atomic consumerDone{false}; + std::atomic producerDone{false}; + + auto consumer = [&] () -> async::detached { + { + auto handle = co_await agent.poll(ce); + if (handle) + handle.ack(); + } + consumerDone.store(true, std::memory_order_release); + }; + + auto producer = [&] () -> async::detached { + co_await mech.post(round); + producerDone.store(true, std::memory_order_release); + }; + + std::atomic ready{0}; + auto spinUntilReady = [&] { + ready.fetch_add(1, std::memory_order_relaxed); + while (ready.load(std::memory_order_relaxed) != 3) + ; + }; + std::thread poller{[&] { spinUntilReady(); consumer(); }}; + std::thread poster{[&] { spinUntilReady(); producer(); }}; + std::thread canceller{[&] { spinUntilReady(); ce.cancel(); }}; + poller.join(); + poster.join(); + canceller.join(); + + while (!consumerDone.load(std::memory_order_acquire)) + std::this_thread::yield(); + + // If the poll was cancelled, detach() acks the outstanding post. + agent.detach(); + + while (!producerDone.load(std::memory_order_acquire)) + std::this_thread::yield(); + } +}