nautilus_lighter/common/

rate_limit.rs

// -------------------------------------------------------------------------------------------------
//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
//  https://nautechsystems.io
//
//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
//  You may not use this file except in compliance with the License.
//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
// -------------------------------------------------------------------------------------------------

//! Rate-limit keys, quotas, and limiters for the Lighter adapter.
//!
//! Lighter meters requests against both the caller IP and the account L1 address.
//!
//! REST reads draw on a per-client read quota. Transactions (`sendTx` /
//! `sendTxBatch`) are metered in one venue bucket per account regardless of
//! transport; single orders go over the WebSocket and batches over HTTP, so both
//! share one [`LighterTxRateLimiter`] to keep their combined rate under the single
//! venue transaction limit.
//!
//! # WebSocket client messages
//!
//! Non-transaction WS frames (subscribe, unsubscribe, resubscribe) face two
//! independent per-IP caps: 200 messages per minute and 50 unacknowledged
//! (inflight) messages. The adapter honours each with a separate mechanism:
//!
//! - Rate: one [`ws_message_rate_limiter`] per venue URL, shared by the data and
//!   execution clients so their combined send rate counts against a single
//!   bucket. It paces at the documented 200/min with a matching 50-message burst.
//! - Inflight: a rate limiter cannot bound inflight, because the unacknowledged
//!   count tracks venue acknowledgement latency (multi-second and fat-tailed),
//!   not emission rate. The feed handler instead gates subscribe dispatch on a
//!   closed-loop count of unacknowledged subscribes
//!   ([`crate::common::consts::SUBSCRIBE_INFLIGHT_MAX`]), releasing a slot on each
//!   ack. Without it, a subscribe storm at startup or reconnect drives inflight
//!   past 50 and the venue returns `30009` / `30010`.
//!
//! `sendTx` is metered in the transaction bucket, not the WS message bucket.

use std::{
    num::NonZeroU32,
    sync::{Arc, LazyLock},
};

use ahash::AHashMap;
use nautilus_network::ratelimiter::{RateLimiter, clock::MonotonicClock, quota::Quota};
use parking_lot::Mutex;
use ustr::Ustr;

/// Conservative Lighter REST rate limit for standard accounts.
///
/// Lighter documents 60 REST requests per rolling minute for standard accounts. Builder and
/// premium accounts can authenticate requests to get higher weighted limits.
pub static LIGHTER_REST_QUOTA: LazyLock<Quota> =
    LazyLock::new(|| Quota::per_minute(NonZeroU32::new(60).expect("non-zero")));

/// Rate-limit bucket key shared by all REST read endpoints.
pub const LIGHTER_REST_BUCKET: &str = "lighter:rest";

/// Rate-limit bucket key for the venue transaction bucket.
///
/// Lighter meters `sendTx` and `sendTxBatch` (HTTP) and the WebSocket `sendTx`
/// path in one per-account bucket. Both transports share a
/// [`LighterTxRateLimiter`] keyed on this so their combined rate stays under the
/// single venue limit.
pub const LIGHTER_TX_BUCKET: &str = "lighter:tx";

/// Rate-limit bucket key for non-transaction WebSocket client messages.
pub const LIGHTER_WS_MESSAGE_BUCKET: &str = "lighter:ws:messages";

/// Lighter's documented WebSocket message rate: 200 per IP per minute.
pub const LIGHTER_WS_MESSAGE_RATE_PER_MIN: u32 = 200;

/// Rate-limiter burst, at Lighter's documented 50-message inflight cap. The
/// closed-loop subscribe gate ([`crate::common::consts::SUBSCRIBE_INFLIGHT_MAX`])
/// is the real inflight bound; this burst only shapes send rate.
pub const LIGHTER_WS_MESSAGE_BURST: u32 = 50;

/// Lighter WebSocket client-message quota, excluding `sendTx` / `sendTxBatch`.
///
/// Paces at the documented 200/min with a 50-message burst; the inflight cap is
/// enforced separately by the subscribe gate (see the module docs).
pub static LIGHTER_WS_MESSAGE_QUOTA: LazyLock<Quota> = LazyLock::new(|| {
    Quota::per_minute(NonZeroU32::new(LIGHTER_WS_MESSAGE_RATE_PER_MIN).expect("non-zero"))
        .allow_burst(NonZeroU32::new(LIGHTER_WS_MESSAGE_BURST).expect("non-zero"))
});

/// Pre-interned rate-limit key for non-transaction WebSocket client messages.
pub static LIGHTER_WS_MESSAGE_RATE_LIMIT_KEY: LazyLock<[Ustr; 1]> =
    LazyLock::new(|| [Ustr::from(LIGHTER_WS_MESSAGE_BUCKET)]);

/// Per-account transaction rate limiter, shared across the HTTP and WebSocket
/// `sendTx` paths so their combined rate honours the single venue bucket.
pub type LighterTxRateLimiter = RateLimiter<Ustr, MonotonicClock>;

/// Shared WebSocket message limiter, keyed by venue WS URL. Both data and
/// execution clients (and the backend balance poller) draw from one bucket
/// per URL so their combined send rate honours the venue's per-IP cap.
pub type LighterWsMessageRateLimiter = Arc<RateLimiter<Ustr, MonotonicClock>>;

// Process-global registry of Lighter WS message limiters, keyed by resolved
// WS URL. Clients on the same URL (a network's data, execution, and backend
// poller) share one bucket honouring the venue per-IP cap; distinct URLs
// (testnet, or a custom endpoint in tests) stay isolated so unrelated traffic
// never contends for the same tokens.
static LIGHTER_WS_MESSAGE_LIMITERS: LazyLock<Mutex<AHashMap<String, LighterWsMessageRateLimiter>>> =
    LazyLock::new(|| Mutex::new(AHashMap::new()));

/// Returns the shared WS message limiter for `url`, creating it on first
/// access. Subsequent calls with the same `url` return the same `Arc`.
#[must_use]
pub fn ws_message_rate_limiter(url: &str) -> LighterWsMessageRateLimiter {
    LIGHTER_WS_MESSAGE_LIMITERS
        .lock()
        .entry(url.to_string())
        .or_insert_with(|| {
            Arc::new(RateLimiter::new_with_quota(
                None,
                vec![(
                    Ustr::from(LIGHTER_WS_MESSAGE_BUCKET),
                    *LIGHTER_WS_MESSAGE_QUOTA,
                )],
            ))
        })
        .clone()
}

/// Resolves a per-minute override to a quota, falling back to the conservative
/// standard-account quota when unset or zero.
#[must_use]
pub fn resolve_quota(per_min: Option<u32>) -> Quota {
    per_min
        .and_then(NonZeroU32::new)
        .map_or(*LIGHTER_REST_QUOTA, Quota::per_minute)
}

/// Builds the shared transaction limiter from a `sendtx_quota_per_min` override,
/// keyed on [`LIGHTER_TX_BUCKET`]. Unset or zero falls back to the standard
/// 60 req/min.
#[must_use]
pub fn build_tx_rate_limiter(sendtx_per_min: Option<u32>) -> Arc<LighterTxRateLimiter> {
    Arc::new(RateLimiter::new_with_quota(
        None,
        vec![(Ustr::from(LIGHTER_TX_BUCKET), resolve_quota(sendtx_per_min))],
    ))
}

/// Awaits transaction-bucket capacity before a `sendTx` on either transport.
///
/// Paces in the caller's task before the frame is enqueued, so neither the HTTP
/// client nor the WebSocket feed-handler task sleeps mid-loop.
pub async fn await_tx_quota(limiter: &LighterTxRateLimiter) {
    limiter
        .await_keys_ready(Some(&[Ustr::from(LIGHTER_TX_BUCKET)]))
        .await;
}

#[cfg(test)]
mod tests {
    use rstest::rstest;

    use super::*;

    #[rstest]
    fn test_resolve_quota_defaults_when_unset_or_zero() {
        assert_eq!(resolve_quota(None), *LIGHTER_REST_QUOTA);
        assert_eq!(resolve_quota(Some(0)), *LIGHTER_REST_QUOTA);
    }

    #[rstest]
    fn test_resolve_quota_uses_override() {
        let expected = Quota::per_minute(NonZeroU32::new(24_000).unwrap());
        assert_eq!(resolve_quota(Some(24_000)), expected);
    }

    #[rstest]
    fn test_build_tx_rate_limiter_handles_unset_zero_and_override() {
        // Builds without panicking for unset/zero (NonZeroU32 guard) and keys on
        // the tx bucket: a fresh limiter admits the first transaction.
        let key = Ustr::from(LIGHTER_TX_BUCKET);

        for sendtx in [None, Some(0), Some(4_000)] {
            let limiter = build_tx_rate_limiter(sendtx);
            assert!(limiter.check_key(&key).is_ok());
        }
    }

    #[rstest]
    fn test_ws_message_quota_matches_venue_caps() {
        // Documented caps, not under-paced: the subscribe gate owns the inflight cap
        assert_eq!(LIGHTER_WS_MESSAGE_RATE_PER_MIN, 200);
        assert_eq!(LIGHTER_WS_MESSAGE_BURST, 50);
        let expected = Quota::per_minute(NonZeroU32::new(LIGHTER_WS_MESSAGE_RATE_PER_MIN).unwrap())
            .allow_burst(NonZeroU32::new(LIGHTER_WS_MESSAGE_BURST).unwrap());
        assert_eq!(*LIGHTER_WS_MESSAGE_QUOTA, expected);
    }

    #[rstest]
    fn test_ws_message_rate_limit_key_matches_bucket() {
        assert_eq!(
            LIGHTER_WS_MESSAGE_RATE_LIMIT_KEY.as_slice(),
            [Ustr::from(LIGHTER_WS_MESSAGE_BUCKET)].as_slice(),
        );
    }

    #[rstest]
    fn test_ws_message_rate_limiter_shared_per_url() {
        let shared_a = ws_message_rate_limiter("wss://example.invalid/share");
        let shared_b = ws_message_rate_limiter("wss://example.invalid/share");
        let isolated = ws_message_rate_limiter("wss://example.invalid/isolated");

        // Same URL: data, execution, and backend poller share one bucket.
        assert!(Arc::ptr_eq(&shared_a, &shared_b));
        // Distinct URL: testnet and custom endpoints stay isolated.
        assert!(!Arc::ptr_eq(&shared_a, &isolated));
    }

    #[rstest]
    fn test_ws_message_rate_limiter_enforces_inflight_burst() {
        let limiter = RateLimiter::new_with_quota(
            None,
            vec![(
                Ustr::from(LIGHTER_WS_MESSAGE_BUCKET),
                *LIGHTER_WS_MESSAGE_QUOTA,
            )],
        );
        let key = LIGHTER_WS_MESSAGE_RATE_LIMIT_KEY[0];

        for _ in 0..LIGHTER_WS_MESSAGE_BURST {
            assert!(limiter.check_key(&key).is_ok());
        }
        assert!(limiter.check_key(&key).is_err());
    }
}