timeout.hpp

Go to the documentation of this file.

#pragma once
 
// timeout / with_timeout — bound a coroutine's execution by a deadline.
//
// Usage:
//
//   auto v = co_await with_timeout(timer, 3s, [](CancellationToken tok) {
//       return http::get(tok, url);
//   });
//
//   // No-token form: timeout still works, but inner work cannot
//   // cooperatively cancel — it just runs to completion in the background.
//   auto v = co_await with_timeout(timer, 3s, []{ return work(); });
//
//   // With a parent token; parent-cancel takes priority over timeout.
//   auto v = co_await with_timeout(parent_tok, timer, 3s, factory);
//
//   // Fail-fast mode: do not wait for inner to unwind; resume parent
//   // immediately with TimeoutError when the deadline expires.
//   auto v = co_await with_timeout(timer, 3s, factory, OnTimeout::Fail);
//
// Behaviour matrix:
//
//   OnTimeout::Cancel  (default; cooperative)
//     When `duration` expires, the inner CancellationToken is flipped.
//     The awaiting coroutine resumes only when the inner factory unwinds
//     via OperationCancelled (or completes naturally just before the
//     timeout). If the timer was the cause, the caller observes
//     `TimeoutError`; if the inner work threw something else, that
//     exception is propagated.
//
//     ⚠️ Footgun — Cancel + no-token factory:
//       The no-token overload (`with_timeout(timer, dur, []{ return work(); })`)
//       has NO way to signal cancellation into the inner work. Combined
//       with OnTimeout::Cancel, this means the timer can only OBSERVE
//       the inner outcome and re-label it as `TimeoutError` AFTER the
//       inner work has finished naturally — the caller does NOT get a
//       prompt timeout. If you need fail-fast behaviour, either pass a
//       factory that takes `CancellationToken` (so cooperative unwind
//       is possible), or use `OnTimeout::Fail`. Mixing no-token + Cancel
//       is a deliberate "best-effort observe" mode and is documented
//       here so reviewers can flag it; for "real" deadlines, prefer
//       one of the other two combinations.
//
//   OnTimeout::Fail    (non-cooperative; fail-fast)
//     When `duration` expires the parent is resumed IMMEDIATELY with
//     TimeoutError. The inner Task is NOT awaited any further — it
//     continues running detached in the background until it naturally
//     completes (its result and any exception are discarded). The
//     inner CancellationToken is still flipped, so cooperative inner
//     work CAN unwind early; non-cooperative inner work simply runs
//     to completion off-stage.
//
//     ⚠️ Caveat — only safe when the inner work has no observable
//     side effects after the timeout, OR side effects are idempotent.
//     If inner mutates shared state the caller can no longer track,
//     prefer OnTimeout::Cancel.
//
// Composition with retry:
//   When wrapped inside retry/retry_with_backoff, each attempt gets its
//   OWN deadline (factory() is re-invoked per attempt). A "global"
//   timeout across all attempts is NOT provided — wrap the outer call
//   in another with_timeout if you need it.
 
#include "aria/async/task.hpp"
#include "aria/async/executor.hpp"
#include "aria/async/cancellation.hpp"
#include "aria/async/virtual_time_executor.hpp"  // for IDelayedScheduler
#include "aria/async/detail/race_slot.hpp"
#include "aria/async/detail/race_trace.hpp"
#include "aria/error.hpp"
 
#include <atomic>
#include <chrono>
#include <coroutine>
#include <exception>
#include <memory>
#include <mutex>
#include <optional>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <variant>
 
namespace aria::async {
 
class TimeoutError : public std::runtime_error {
public:
    TimeoutError() : std::runtime_error("operation timed out") {}
};

enum class OnTimeout {
    Cancel,
    Fail,
};
 
namespace detail {
 
template<typename Factory>
concept TimeoutTokenAcceptingFactory =
    std::invocable<Factory, CancellationToken>;
 
template<typename Factory>
concept TimeoutPlainFactory =
    std::invocable<Factory>;
 
template<typename Factory>
struct timeout_factory_value;
 
template<typename Factory>
    requires TimeoutTokenAcceptingFactory<Factory>
struct timeout_factory_value<Factory> {
    using type = typename std::invoke_result_t<Factory, CancellationToken>::promise_type::value_type;
};
 
template<typename Factory>
    requires (!TimeoutTokenAcceptingFactory<Factory>) && TimeoutPlainFactory<Factory>
struct timeout_factory_value<Factory> {
    using type = typename std::invoke_result_t<Factory>::promise_type::value_type;
};
 
template<typename Factory>
using timeout_factory_value_t = typename timeout_factory_value<Factory>::type;

template<typename Factory>
auto invoke_factory(Factory& f, CancellationToken tok) {
    if constexpr (TimeoutTokenAcceptingFactory<Factory>) {
        return f(tok);
    } else {
        return f();
    }
}
 
// ── Cancel mode (cooperative) ────────────────────────────────────────────

template<typename Factory>
auto with_timeout_impl_(std::optional<CancellationToken> parent,
                        IDelayedScheduler& timer,
                        std::chrono::milliseconds duration,
                        Factory factory)
    -> Task<timeout_factory_value_t<Factory>>
{
    using R = timeout_factory_value_t<Factory>;
 
    // Per-invocation cancellation source for the inner work; flipped by
    // either the deadline or (if engaged) the parent token.
    auto inner_src = std::make_shared<CancellationSource>();
    auto inner_tok = inner_src->token();
 
    // Race flag: 0 = pending, 1 = inner-won, 2 = timer-won.
    auto winner = std::make_shared<std::atomic<int>>(0);
 
    publish_race_trace(race_source::kWithTimeout, race_op::kStart);
    struct EndTrace {
        ~EndTrace() {
            publish_race_trace(race_source::kWithTimeout, race_op::kEnd);
        }
    } end_trace;
 
    // Arm the deadline. The lambda below takes care of cancelling the
    // inner work iff it wins the race.
    timer.post_after(duration, [winner, inner_src]() {
        int expected = 0;
        if (winner->compare_exchange_strong(expected, 2)) {
            // D-31.1: the deadline claimed the race. Cancelling the sole
            // participant IS the timeout here, so no separate
            // race_loser_cancel is published.
            publish_race_trace(race_source::kWithTimeout, race_op::kTimeout, 0,
                               ::aria::Error::timeout("with_timeout"));
            inner_src->cancel();
        }
    });
 
    // If a parent token is engaged, propagate its cancellation to the
    // inner work synchronously. Parent-cancel always wins over the
    // deadline (we re-throw the original OperationCancelled below
    // instead of converting to TimeoutError).
    if (parent) {
        parent->on_cancel([inner_src]{ inner_src->cancel(); });
    }
 
    auto finish_success = [&] {
        int expected = 0;
        const bool inner_won = winner->compare_exchange_strong(expected, 1);
        if (parent && parent->is_cancelled()) throw OperationCancelled{};
        // Cancel mode waits for non-cooperative work to settle, but the
        // elapsed deadline still determines the caller's result.
        if (!inner_won) throw TimeoutError{};
        publish_race_trace(race_source::kWithTimeout, race_op::kWon);
    };
 
    try {
        if constexpr (std::is_void_v<R>) {
            co_await invoke_factory(factory, inner_tok);
            finish_success();
            co_return;
        } else {
            R value = co_await invoke_factory(factory, inner_tok);
            finish_success();
            co_return value;
        }
    } catch (const OperationCancelled&) {
        int expected = 0;
        winner->compare_exchange_strong(expected, 1);
        if (parent && parent->is_cancelled()) {
            publish_race_trace(race_source::kWithTimeout,
                               race_op::kParentCancel, 0,
                               ::aria::Error::cancellation("with_timeout"));
            throw;
        }
        if (winner->load() == 2) throw TimeoutError{};
        throw;
    } catch (...) {
        // A failed inner task has also finished the race. Prevent the
        // still-queued timer from cancelling it or reporting a later timeout.
        int expected = 0;
        winner->compare_exchange_strong(expected, 1);
        throw;
    }
}
 
// ── Fail mode (non-cooperative, fail-fast) ───────────────────────────────
 
// Race state and awaiter live in detail/race_slot.hpp (shared with
// when_any's race-hardening path).
//
// Winner codes for with_timeout::Fail:
//   1 = Inner completed first
//   2 = Timer fired first
//   3 = Parent token cancelled (engaged parent only)

template<typename Factory, typename R>
Task<void> drive_inner_for_fail_(
    Factory factory,
    CancellationToken inner_tok,
    std::shared_ptr<RaceSlot<R>> slot)
{
    try {
        if constexpr (std::is_void_v<R>) {
            co_await invoke_factory(factory, inner_tok);
            if (slot->try_claim(/*Inner=*/1)) {
                slot->store_value_or_exception();   // void success
                slot->publish(/*Inner=*/1);
                publish_race_trace(race_source::kWithTimeout, race_op::kWon);
                slot->notify_winner_resume();
            }
            // else: timer or parent already won; result discarded.
        } else {
            R value = co_await invoke_factory(factory, inner_tok);
            if (slot->try_claim(/*Inner=*/1)) {
                slot->store_value_or_exception(std::move(value));
                slot->publish(/*Inner=*/1);
                publish_race_trace(race_source::kWithTimeout, race_op::kWon);
                slot->notify_winner_resume();
            }
        }
    } catch (...) {
        if (slot->try_claim(/*Inner=*/1)) {
            slot->result.template emplace<2>(std::current_exception());
            slot->publish(/*Inner=*/1);
            // The inner task still won the race — it just won with a
            // failure. The exception itself surfaces through
            // await_resume, not through an arbitration event.
            publish_race_trace(race_source::kWithTimeout, race_op::kWon);
            slot->notify_winner_resume();
        }
        // else: timer/parent won; exception discarded silently.
    }
}

template<typename Factory>
auto with_timeout_fail_impl_(std::optional<CancellationToken> parent,
                             IDelayedScheduler& timer,
                             std::chrono::milliseconds duration,
                             Factory factory)
    -> Task<timeout_factory_value_t<Factory>>
{
    using R = timeout_factory_value_t<Factory>;
 
    auto slot      = std::make_shared<RaceSlot<R>>();
    auto inner_src = std::make_shared<CancellationSource>();
    auto inner_tok = inner_src->token();
 
    publish_race_trace(race_source::kWithTimeout, race_op::kStart);
    struct EndTrace {
        ~EndTrace() {
            publish_race_trace(race_source::kWithTimeout, race_op::kEnd);
        }
    } end_trace;
 
    // Arm the deadline.
    timer.post_after(duration, [slot, inner_src]() {
        if (slot->try_claim(/*Timer=*/2)) {
            // Mark inner as cancelled so cooperative work CAN unwind
            // (best-effort; non-cooperative inner just runs in the bg).
            inner_src->cancel();
            slot->result.template emplace<2>(
                std::make_exception_ptr(TimeoutError{}));
            slot->publish(/*Timer=*/2);
            publish_race_trace(race_source::kWithTimeout, race_op::kTimeout, 0,
                               ::aria::Error::timeout("with_timeout"));
            slot->notify_winner_resume();
        }
    });
 
    // Wire parent token (if engaged): parent-cancel beats timeout, same
    // as Cancel mode. Manifest as OperationCancelled in await_resume.
    if (parent) {
        parent->on_cancel([slot, inner_src]{
            if (slot->try_claim(/*ParentCancel=*/3)) {
                inner_src->cancel();
                slot->result.template emplace<2>(
                    std::make_exception_ptr(OperationCancelled{}));
                slot->publish(/*ParentCancel=*/3);
                publish_race_trace(race_source::kWithTimeout,
                                   race_op::kParentCancel, 0,
                                   ::aria::Error::cancellation("with_timeout"));
                slot->notify_winner_resume();
            }
        });
    }
 
    // Start the inner driver detached. It owns its own coroutine frame
    // via Task::start_detached and may outlive this wrapper if the
    // timer wins the race.
    drive_inner_for_fail_<Factory, R>(std::move(factory), inner_tok, slot)
        .start_detached();
 
    // Park until a winner publishes a result. The guard closes the trace
    // on both successful decoding and exception unwinding.
    if constexpr (std::is_void_v<R>) {
        co_await RaceSlotAwaiter<R>{slot};
        co_return;
    } else {
        auto decoded = co_await RaceSlotAwaiter<R>{slot};
        co_return decoded;
    }
}
 
}  // namespace detail
 
// ── Public API: 4 thin wrappers, dispatched on OnTimeout ─────────────────
 
template<typename Factory>
auto with_timeout(IDelayedScheduler& timer,
                  std::chrono::milliseconds duration,
                  Factory factory,
                  OnTimeout on_timeout = OnTimeout::Cancel)
    -> Task<detail::timeout_factory_value_t<Factory>>
{
    if (on_timeout == OnTimeout::Fail) {
        return detail::with_timeout_fail_impl_(
            std::optional<CancellationToken>{}, timer, duration, std::move(factory));
    }
    return detail::with_timeout_impl_(
        std::optional<CancellationToken>{}, timer, duration, std::move(factory));
}
 
template<typename Factory>
auto with_timeout(CancellationToken parent,
                  IDelayedScheduler& timer,
                  std::chrono::milliseconds duration,
                  Factory factory,
                  OnTimeout on_timeout = OnTimeout::Cancel)
    -> Task<detail::timeout_factory_value_t<Factory>>
{
    if (on_timeout == OnTimeout::Fail) {
        return detail::with_timeout_fail_impl_(
            std::optional<CancellationToken>{std::move(parent)},
            timer, duration, std::move(factory));
    }
    return detail::with_timeout_impl_(
        std::optional<CancellationToken>{std::move(parent)},
        timer, duration, std::move(factory));
}
 
}  // namespace aria::async