async_work.h

#pragma once
 
#include <tess/core/assert.h>
#include <tess/core/fail_fast.h>
#include <tess/diagnostics/diagnostics.h>
 
#include <cstddef>
#include <cstdint>
#include <limits>
#include <source_location>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace tess {

struct AsyncTicket {
  std::uint32_t index = 0;
  std::uint64_t generation = 0;
 
  friend constexpr bool operator==(AsyncTicket lhs,
                                   AsyncTicket rhs) noexcept = default;
};

struct AsyncVersion {
  std::uint64_t value = 0;
 
  constexpr AsyncVersion() noexcept = default;
  constexpr AsyncVersion(std::uint64_t version) noexcept : value(version) {}
 
  friend constexpr bool operator==(AsyncVersion lhs,
                                   AsyncVersion rhs) noexcept = default;
};

enum class AsyncResultState : std::uint8_t {
  Unbound,
  Immediate,
  Pending,
  Ready,
  Failed,
  Cancelled,
  Superseded,
  Stale,
};

struct AsyncWorkBudget {
  std::uint32_t max_items = 0;
};

enum class AsyncStepState : std::uint8_t {
  Pending,
  Ready,
  Failed,
  Stale,
};

struct AsyncWorkStep {
  AsyncStepState state = AsyncStepState::Pending;
  std::uint32_t items_done = 0;
  AsyncVersion result_version{};
};

struct AsyncAdvanceStats {
  std::uint32_t invoked = 0;
  std::uint32_t items_done = 0;
  std::uint32_t pending = 0;
  std::uint32_t ready = 0;
  std::uint32_t failed = 0;
  std::uint32_t stale = 0;
};

template <typename T>
class ResumableWorkQueue {
  static_assert(std::is_default_constructible_v<T>,
                "ResumableWorkQueue<T> requires T to be default "
                "constructible");
 
 public:
  using WorkFn = AsyncWorkStep (*)(void*, AsyncWorkBudget, T&);
 
  ResumableWorkQueue() = default;
  ResumableWorkQueue(const ResumableWorkQueue&) = delete;
  auto operator=(const ResumableWorkQueue&) -> ResumableWorkQueue& = delete;
  ResumableWorkQueue(ResumableWorkQueue&&) = delete;
  auto operator=(ResumableWorkQueue&&) -> ResumableWorkQueue& = delete;
  ~ResumableWorkQueue() = default;

  void set_flow_accounting(diagnostics::FlowAccounting* accounting) noexcept {
    if (!slots_.empty()) {
      detail::fail_fast(
          "ResumableWorkQueue::set_flow_accounting requires an empty queue");
    }
    accounting_ = accounting;
  }

  void observe_flow_tick(std::uint64_t tick) noexcept {
    if (reject_reentrant_mutation()) {
      return;
    }
    if (accounting_ == nullptr) {
      return;
    }
    accounting_->observe_tick(tick);
    const auto now = accounting_->last_observed_tick;
    auto oldest = now;
    auto any_pending = false;
    for (const auto& slot : slots_) {
      if (slot.state == AsyncResultState::Pending) {
        any_pending = true;
        oldest = slot.submitted_tick < oldest ? slot.submitted_tick : oldest;
      }
    }
    accounting_->counters.oldest_outstanding_age_ticks =
        any_pending ? now - oldest : 0;
  }
 
  void reserve_tickets(std::size_t count) {
    if (reject_reentrant_mutation()) {
      return;
    }
    slots_.reserve(count);
  }
 
  template <typename Work>
  [[nodiscard]] auto submit(
      Work& work, AsyncVersion required_version = {},
      std::source_location source = std::source_location::current())
      -> AsyncTicket {
    static_assert(
        std::is_invocable_r_v<AsyncWorkStep, Work&, AsyncWorkBudget, T&>,
        "ResumableWorkQueue::submit requires Work(AsyncWorkBudget, T&) to "
        "return AsyncWorkStep");
    return submit(
        static_cast<void*>(&work),
        [](void* context, AsyncWorkBudget budget, T& value) -> AsyncWorkStep {
          return (*static_cast<Work*>(context))(budget, value);
        },
        required_version, source);
  }
 
  [[nodiscard]] auto submit(
      void* context, WorkFn work, AsyncVersion required_version = {},
      std::source_location source = std::source_location::current())
      -> AsyncTicket {
    if (reject_reentrant_mutation()) {
      account_rejected_offer();
      return {};
    }
    if (work == nullptr) {
      detail::fail_fast("ResumableWorkQueue::submit received a null callback");
    }
    if (slots_.size() > std::numeric_limits<std::uint32_t>::max()) {
      detail::fail_fast(
          "ResumableWorkQueue::submit exhausted the AsyncTicket index space");
    }
    const auto index = static_cast<std::uint32_t>(slots_.size());
    slots_.push_back(Slot{});
    auto& slot = slots_.back();
    slot.context = context;
    slot.work = work;
    slot.required_version = required_version;
    slot.source = source;
    slot.state = AsyncResultState::Pending;
    if (accounting_ != nullptr) {
      ++accounting_->counters.offered;
      accounting_->record_admitted();
      slot.submitted_tick = accounting_->last_observed_tick;
    }
    return AsyncTicket{index, generation_};
  }
 
  [[nodiscard]] auto submit_immediate(
      T value, AsyncVersion result_version = {},
      std::source_location source = std::source_location::current())
      -> AsyncTicket {
    if (reject_reentrant_mutation()) {
      account_rejected_offer();
      return {};
    }
    if (slots_.size() > std::numeric_limits<std::uint32_t>::max()) {
      detail::fail_fast(
          "ResumableWorkQueue::submit_immediate exhausted the AsyncTicket "
          "index space");
    }
    const auto index = static_cast<std::uint32_t>(slots_.size());
    // Build the slot before storing it so a throwing value move admits
    // nothing and leaves no Unbound slot behind.
    auto slot = Slot{};
    slot.value = std::move(value);
    slot.result_version = result_version;
    slot.source = source;
    slot.state = AsyncResultState::Immediate;
    slots_.push_back(std::move(slot));
    if (accounting_ != nullptr) {
      auto& counters = accounting_->counters;
      ++counters.offered;
      ++counters.admitted;
      ++counters.completed;
    }
    return AsyncTicket{index, generation_};
  }
 
  [[nodiscard]] auto advance(AsyncWorkBudget budget) -> AsyncAdvanceStats {
    if (in_advance_) {
      detail::fail_fast(
          "ResumableWorkQueue::advance rejected mutation during advance");
    }
    struct AdvanceGuard {
      bool& active;
      // Exception propagation is intentional. This guard restores the
      // reentrancy sentinel while leaving the throwing slot Pending so the
      // caller may inspect state and choose whether to retry or fail it.
      ~AdvanceGuard() { active = false; }
    };
    in_advance_ = true;
    const auto guard = AdvanceGuard{in_advance_};
    auto stats = AsyncAdvanceStats{};
    if (accounting_ != nullptr) {
      accounting_->counters.offered_work_units += budget.max_items;
    }
    auto remaining = budget.max_items;
    auto invocations_remaining = budget.max_items;
    for (auto& slot : slots_) {
      if (remaining == 0 || invocations_remaining == 0) {
        break;
      }
      if (slot.state != AsyncResultState::Pending) {
        continue;
      }
      TESS_ASSERT(slot.work != nullptr);
      const auto step =
          slot.work(slot.context, AsyncWorkBudget{remaining}, slot.value);
      ++stats.invoked;
      --invocations_remaining;
      if (step.items_done > remaining) {
        slot.state = AsyncResultState::Failed;
        account_terminal(slot);
        continue;
      }
      remaining -= step.items_done;
      if (accounting_ != nullptr) {
        // Committed per step: a later callback's exception must not
        // discard work already consumed in this advance.
        accounting_->counters.consumed_work_units += step.items_done;
      }
      slot.result_version = step.result_version;
      switch (step.state) {
        case AsyncStepState::Pending:
          break;
        case AsyncStepState::Ready:
          slot.state = AsyncResultState::Ready;
          account_terminal(slot);
          break;
        case AsyncStepState::Failed:
          slot.state = AsyncResultState::Failed;
          account_terminal(slot);
          break;
        case AsyncStepState::Stale:
          slot.state = AsyncResultState::Stale;
          account_terminal(slot);
          break;
      }
      stats.items_done += step.items_done;
    }
    summarize_states(stats);
    return stats;
  }
 
  [[nodiscard]] auto state(AsyncTicket ticket) const noexcept
      -> AsyncResultState {
    const auto* slot = find(ticket);
    return slot == nullptr ? AsyncResultState::Unbound : slot->state;
  }

  [[nodiscard]] auto result(AsyncTicket ticket) const noexcept -> const T* {
    const auto* slot = find(ticket);
    if (slot == nullptr || (slot->state != AsyncResultState::Immediate &&
                            slot->state != AsyncResultState::Ready)) {
      return nullptr;
    }
    return &slot->value;
  }
 
  [[nodiscard]] auto required_version(AsyncTicket ticket) const noexcept
      -> AsyncVersion {
    const auto* slot = find(ticket);
    return slot == nullptr ? AsyncVersion{} : slot->required_version;
  }
 
  [[nodiscard]] auto result_version(AsyncTicket ticket) const noexcept
      -> AsyncVersion {
    const auto* slot = find(ticket);
    return slot == nullptr ? AsyncVersion{} : slot->result_version;
  }
 
  [[nodiscard]] auto source(AsyncTicket ticket) const noexcept
      -> std::source_location {
    const auto* slot = find(ticket);
    return slot == nullptr ? std::source_location{} : slot->source;
  }

  [[nodiscard]] bool cancel(AsyncTicket ticket) noexcept {
    return set_terminal(ticket, AsyncResultState::Cancelled);
  }

  [[nodiscard]] bool supersede(AsyncTicket ticket) noexcept {
    return set_terminal(ticket, AsyncResultState::Superseded);
  }

  [[nodiscard]] bool fail(AsyncTicket ticket) noexcept {
    return set_terminal(ticket, AsyncResultState::Failed);
  }

  [[nodiscard]] bool mark_stale(AsyncTicket ticket) noexcept {
    return set_terminal(ticket, AsyncResultState::Stale);
  }

  [[nodiscard]] bool mark_stale_if_version(AsyncTicket ticket,
                                           AsyncVersion current) noexcept {
    if (reject_reentrant_mutation()) {
      return false;
    }
    auto* slot = find(ticket);
    if (slot == nullptr ||
        (slot->state != AsyncResultState::Immediate &&
         slot->state != AsyncResultState::Ready) ||
        slot->result_version == current) {
      return false;
    }
    slot->state = AsyncResultState::Stale;
    if (accounting_ != nullptr) {
      // Reclassification, not a second terminal outcome: the result
      // completed earlier and its residence was already recorded, so
      // only the buckets swap. `completed` is non-monotonic here by
      // documented design.
      auto& counters = accounting_->counters;
      if (counters.completed > 0) {
        --counters.completed;
      }
      ++counters.stale;
    }
    return true;
  }
 
  [[nodiscard]] auto size() const noexcept -> std::size_t {
    return slots_.size();
  }
 
  [[nodiscard]] auto generation() const noexcept -> std::uint64_t {
    return generation_;
  }
 
  void clear() noexcept {
    if (reject_reentrant_mutation()) {
      return;
    }
    if (accounting_ != nullptr) {
      for (auto& slot : slots_) {
        if (slot.state == AsyncResultState::Pending) {
          ++accounting_->counters.dropped_after_admission;
          close_outstanding(slot);
        }
      }
    }
    slots_.clear();
    ++generation_;
    if (generation_ == 0) {
      ++generation_;
    }
  }
 
 private:
  struct Slot {
    T value{};
    void* context = nullptr;
    WorkFn work = nullptr;
    AsyncVersion required_version{};
    AsyncVersion result_version{};
    std::source_location source = std::source_location::current();
    AsyncResultState state = AsyncResultState::Unbound;
    std::uint64_t submitted_tick = 0;
  };
 
  [[nodiscard]] auto reject_reentrant_mutation() const noexcept -> bool {
    if (!in_advance_) {
      return false;
    }
    detail::fail_fast("ResumableWorkQueue rejected mutation during advance");
  }
 
  void account_rejected_offer() noexcept {
    if (accounting_ != nullptr) {
      ++accounting_->counters.offered;
      ++accounting_->counters.rejected;
    }
  }
 
  [[nodiscard]] auto find(AsyncTicket ticket) noexcept -> Slot* {
    if (ticket.generation != generation_ || ticket.index >= slots_.size()) {
      return nullptr;
    }
    return &slots_[ticket.index];
  }
 
  [[nodiscard]] auto find(AsyncTicket ticket) const noexcept -> const Slot* {
    if (ticket.generation != generation_ || ticket.index >= slots_.size()) {
      return nullptr;
    }
    return &slots_[ticket.index];
  }
 
  [[nodiscard]] bool set_terminal(AsyncTicket ticket,
                                  AsyncResultState state) noexcept {
    if (reject_reentrant_mutation()) {
      return false;
    }
    auto* slot = find(ticket);
    if (slot == nullptr || slot->state != AsyncResultState::Pending) {
      return false;
    }
    slot->state = state;
    account_terminal(*slot);
    return true;
  }

  void account_terminal(Slot& slot) noexcept {
    if (accounting_ == nullptr) {
      return;
    }
    auto& counters = accounting_->counters;
    switch (slot.state) {
      case AsyncResultState::Ready:
        ++counters.completed;
        break;
      case AsyncResultState::Failed:
        ++counters.failed;
        break;
      case AsyncResultState::Cancelled:
        ++counters.cancelled;
        break;
      case AsyncResultState::Superseded:
        ++counters.superseded;
        break;
      case AsyncResultState::Stale:
        ++counters.stale;
        break;
      case AsyncResultState::Unbound:
      case AsyncResultState::Immediate:
      case AsyncResultState::Pending:
        return;
    }
    close_outstanding(slot);
  }
 
  void close_outstanding(const Slot& slot) noexcept {
    accounting_->record_left_outstanding();
    accounting_->counters.residence_ticks_accumulated +=
        accounting_->last_observed_tick - slot.submitted_tick;
  }
 
  void summarize_states(AsyncAdvanceStats& stats) const noexcept {
    for (const auto& slot : slots_) {
      switch (slot.state) {
        case AsyncResultState::Pending:
          ++stats.pending;
          break;
        case AsyncResultState::Immediate:
        case AsyncResultState::Ready:
          ++stats.ready;
          break;
        case AsyncResultState::Failed:
        case AsyncResultState::Cancelled:
        case AsyncResultState::Superseded:
          ++stats.failed;
          break;
        case AsyncResultState::Stale:
          ++stats.stale;
          break;
        case AsyncResultState::Unbound:
          break;
      }
    }
  }
 
  std::vector<Slot> slots_;
  std::uint64_t generation_ = 1;
  bool in_advance_ = false;
  diagnostics::FlowAccounting* accounting_ = nullptr;
};
 
}  // namespace tess