rate_limit.hpp

#pragma once

#define ATL_TRANSPORT_LAYER_PRESENT 1
 
#include "../core/enum_values.hpp"
 
#include <array>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <stdexcept>
#include <string_view>
 
namespace algotradelib {
namespace network {

enum class RateLimitGroup {
  TRADING,
  AUTH,
  ACCOUNT,
  MARKET_DATA,
  REFERENCE,
};

inline constexpr std::array<RateLimitGroup, 5> all_rate_limit_groups{
  RateLimitGroup::TRADING,
  RateLimitGroup::AUTH,
  RateLimitGroup::ACCOUNT,
  RateLimitGroup::MARKET_DATA,
  RateLimitGroup::REFERENCE,
};

inline constexpr int rate_limit_priority(RateLimitGroup group) noexcept {
  switch (group) {
    case RateLimitGroup::AUTH:
      return 50;
    case RateLimitGroup::TRADING:
      return 40;
    case RateLimitGroup::ACCOUNT:
      return 30;
    case RateLimitGroup::MARKET_DATA:
      return 20;
    case RateLimitGroup::REFERENCE:
      return 10;
  }
  return 0;
}

inline constexpr std::string_view rate_limit_group_name(RateLimitGroup group) noexcept {
  switch (group) {
    case RateLimitGroup::AUTH:
      return "auth";
    case RateLimitGroup::TRADING:
      return "trading";
    case RateLimitGroup::ACCOUNT:
      return "account";
    case RateLimitGroup::MARKET_DATA:
      return "market_data";
    case RateLimitGroup::REFERENCE:
      return "reference";
  }
  return "";
}
 
static_assert(
  lists_every_value(all_rate_limit_groups, [](RateLimitGroup group) { return !rate_limit_group_name(group).empty(); }),
  "all_rate_limit_groups must list every RateLimitGroup"
);

inline constexpr std::size_t rate_limit_group_count = all_rate_limit_groups.size();

inline constexpr std::size_t rate_limit_group_index(RateLimitGroup group) noexcept {
  return static_cast<std::size_t>(group);
}

struct RateLimitBudget {
  int burst_{8};
  int per_second_{8};
  int trading_reserve_{2};
};

inline constexpr int usable_burst(const RateLimitBudget& budget, RateLimitGroup group) noexcept {
  switch (group) {
    // An order, and the token it needs, may use the reserve.
    case RateLimitGroup::TRADING:
    case RateLimitGroup::AUTH:
      return budget.burst_;
    case RateLimitGroup::ACCOUNT:
    case RateLimitGroup::MARKET_DATA:
    case RateLimitGroup::REFERENCE:
      break;
  }
  const int reserved = budget.burst_ - budget.trading_reserve_;
  return reserved > 0 ? reserved : 1;
}

template <typename Clock = std::chrono::steady_clock>
class BasicRateLimitBucket {
public:
  using ClockType = Clock;

  using TimePoint = typename Clock::time_point;

  using Duration = std::chrono::nanoseconds;

  BasicRateLimitBucket(int burst, int per_second) { resize(burst, per_second); }

  explicit BasicRateLimitBucket(const RateLimitBudget& budget)
  : BasicRateLimitBucket(budget.burst_, budget.per_second_) {}

  void resize(int burst, int per_second) {
    if (burst < 1) {
      throw std::invalid_argument("rate limit burst must be at least one request");
    }
    if (per_second < 1) {
      throw std::invalid_argument("rate limit rate must be at least one request per second");
    }
    burst_ = burst;
    per_second_ = per_second;
    interval_ = Duration(std::chrono::nanoseconds(std::chrono::seconds(1)).count() / per_second);
  }

  Duration delay_until(TimePoint now, int allowance) const {
    if (!started_) {
      return Duration::zero();
    }
    const auto earliest = schedule_ - tolerance(allowance);
    if (now >= earliest) {
      return Duration::zero();
    }
    return std::chrono::duration_cast<Duration>(earliest - now);
  }

  bool ready(TimePoint now, int allowance) const {
    return delay_until(now, allowance) == Duration::zero();
  }

  bool try_acquire(TimePoint now, int allowance) {
    if (!ready(now, allowance)) {
      return false;
    }
    if (!started_) {
      schedule_ = now + interval_;
      started_ = true;
      return true;
    }
    schedule_ = (schedule_ > now ? schedule_ : now) + interval_;
    return true;
  }

  void cool_down(TimePoint now, Duration duration) {
    if (duration < Duration::zero()) {
      duration = Duration::zero();
    }
    const auto held = now + duration + tolerance(burst_);
    if (!started_ || held > schedule_) {
      schedule_ = held;
      started_ = true;
    }
  }

  void observe_limit(TimePoint now, int burst, int per_second, Duration refill) {
    resize(burst, per_second);
    cool_down(now, refill > Duration::zero() ? refill : interval_);
  }

  int available(TimePoint now, int allowance) const {
    if (!started_) {
      return allowance;
    }
    const auto earliest = schedule_ - tolerance(allowance);
    if (now < earliest) {
      return 0;
    }
    const auto spare = std::chrono::duration_cast<Duration>(now - earliest);
    const auto slots = spare / interval_ + 1;
    return slots > allowance ? allowance : static_cast<int>(slots);
  }

  int burst() const noexcept { return burst_; }

  int per_second() const noexcept { return per_second_; }

  Duration interval() const noexcept { return interval_; }
 
private:
  // Returns how far ahead of now the schedule may run for `allowance`.
  Duration tolerance(int allowance) const noexcept {
    const int slots = allowance < 1 ? 1 : (allowance > burst_ ? burst_ : allowance);
    return interval_ * (slots - 1);
  }
 
  int burst_{1};
  int per_second_{1};
  Duration interval_{std::chrono::seconds(1)};
  TimePoint schedule_{};
  bool started_{false};
};

using RateLimitBucket = BasicRateLimitBucket<>;

struct RateLimitGroupStats {
  std::uint64_t admitted_{0};
  std::uint64_t delayed_{0};
  std::uint64_t shed_{0};
  std::chrono::nanoseconds waited_{0};
};

struct RateLimitStats {
  RateLimitGroupStats groups_[rate_limit_group_count];
  int burst_{0};
  int per_second_{0};
  int available_{0};
  std::uint64_t broker_rejections_{0};

  RateLimitGroupStats& group(RateLimitGroup value) {
    return groups_[rate_limit_group_index(value)];
  }

  const RateLimitGroupStats& group(RateLimitGroup value) const {
    return groups_[rate_limit_group_index(value)];
  }
};

class RateLimitController {
public:
  using Duration = RateLimitBucket::Duration;

  explicit RateLimitController(RateLimitBudget budget = {})
  : budget_(budget), bucket_(budget) {
    stats_.burst_ = budget.burst_;
    stats_.per_second_ = budget.per_second_;
  }

  Duration delay_until(RateLimitGroup group) const {
    return bucket_.delay_until(RateLimitBucket::ClockType::now(), usable_burst(budget_, group));
  }

  bool acquire(RateLimitGroup group) {
    const auto now = RateLimitBucket::ClockType::now();
    if (!bucket_.try_acquire(now, usable_burst(budget_, group))) {
      return false;
    }
    ++stats_.group(group).admitted_;
    stats_.available_ = bucket_.available(now, usable_burst(budget_, group));
    return true;
  }

  void record_delay(RateLimitGroup group, Duration waited) {
    auto& counters = stats_.group(group);
    ++counters.delayed_;
    counters.waited_ += waited;
  }

  void record_shed(RateLimitGroup group) { ++stats_.group(group).shed_; }

  void hold(Duration duration) {
    ++stats_.broker_rejections_;
    bucket_.cool_down(RateLimitBucket::ClockType::now(), duration);
  }

  void shrink(int burst, int per_second, Duration retry) {
    // The bucket validates first, so a refused figure leaves the reported budget the one in force.
    bucket_.observe_limit(RateLimitBucket::ClockType::now(), burst, per_second, retry);
    ++stats_.broker_rejections_;
    apply_budget(burst, per_second);
  }

  const RateLimitBudget& budget() const noexcept { return budget_; }

  const RateLimitStats& stats() const noexcept { return stats_; }
 
private:
  // Records a new allowance, keeping the reserve smaller than the burst.
  void apply_budget(int burst, int per_second) {
    budget_.burst_ = burst;
    budget_.per_second_ = per_second;
    if (budget_.trading_reserve_ >= burst) {
      budget_.trading_reserve_ = burst > 1 ? burst - 1 : 0;
    }
    stats_.burst_ = burst;
    stats_.per_second_ = per_second;
  }
 
  RateLimitBudget budget_;
  RateLimitBucket bucket_;
  RateLimitStats stats_{};
};
 
} // namespace network
} // namespace algotradelib