rate_limit.hpp

Go to the documentation of this file.

#pragma once
 
#include "kalshi/error.hpp"
#include "kalshi/http_client.hpp"
 
#include <chrono>
#include <memory>
#include <mutex>
#include <optional>
#include <string>
#include <string_view>
#include <vector>
 
namespace kalshi {
 
class TokenBucket {
public:
    struct Config {
        double capacity{0.0};
        double refill_per_second{0.0};
        std::optional<double> initial_tokens;
    };
 
    explicit TokenBucket(Config config);
 
    [[nodiscard]] bool try_acquire(double cost = 1.0) noexcept;
 
    [[nodiscard]] bool acquire_for(double cost, std::chrono::nanoseconds max_wait);
 
    [[nodiscard]] std::chrono::nanoseconds wait_time(double cost) const noexcept;
 
    [[nodiscard]] double available() const noexcept;
    void reset() noexcept;
    [[nodiscard]] const Config& config() const noexcept { return config_; }
 
private:
    using Clock = std::chrono::steady_clock;
 
    void refill(Clock::time_point now) const noexcept;
    [[nodiscard]] std::chrono::nanoseconds wait_time_locked(double cost) const noexcept;
 
    Config config_;
    mutable std::mutex mutex_;
    mutable double tokens_{0.0};
    mutable Clock::time_point last_refill_;
};
 
struct EndpointCostRule {
    HttpMethod method{HttpMethod::GET};
    std::string path;
    double cost{0.0};
};
 
struct RateLimitConfig {
    TokenBucket::Config read{.capacity = 600.0, .refill_per_second = 200.0, .initial_tokens = {}};
    TokenBucket::Config write{.capacity = 100.0, .refill_per_second = 100.0, .initial_tokens = {}};
    double default_cost{10.0};
    std::vector<EndpointCostRule> costs;
    std::chrono::milliseconds max_wait{std::chrono::seconds{5}};
};
 
class RateLimitedTransport final : public HttpTransport {
public:
    RateLimitedTransport(std::shared_ptr<const HttpTransport> inner, RateLimitConfig config);
 
    [[nodiscard]] Result<HttpResponse> request(HttpMethod method, std::string_view path,
                                               std::string_view body = {}) const override;
 
    [[nodiscard]] static bool is_write(HttpMethod method, std::string_view path) noexcept;
 
    [[nodiscard]] double cost(HttpMethod method, std::string_view path,
                              std::string_view body) const;
 
private:
    std::shared_ptr<const HttpTransport> inner_;
    RateLimitConfig config_;
    mutable TokenBucket read_;
    mutable TokenBucket write_;
};
 
} // namespace kalshi