hypercall/read_cache/

tier.rs

use crate::rsm::MarginMode;
use anyhow::{anyhow, Result};
use hypercall_db::TierWriter;
use hypercall_db::UserTierRecord as UserTier;
use hypercall_db::UserTierUpdate;
use hypercall_types::api_models::{TradingLimits, UserTierData};
use hypercall_types::WalletAddress;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info};

fn default_user_tier(wallet: WalletAddress) -> UserTierData {
    UserTierData {
        wallet_address: wallet,
        tier: "tier2".to_string(),
    }
}

const DEFAULT_MAX_OPEN_ORDERS: i32 = 100;
const DEFAULT_MAX_OPEN_POSITIONS: i32 = 50;

#[derive(Debug, Clone, Copy)]
pub struct TierCacheConfig {
    pub max_open_orders_default: i32,
    pub max_open_positions_default: i32,
}

impl Default for TierCacheConfig {
    fn default() -> Self {
        Self {
            max_open_orders_default: DEFAULT_MAX_OPEN_ORDERS,
            max_open_positions_default: DEFAULT_MAX_OPEN_POSITIONS,
        }
    }
}

pub struct TierCache {
    tiers: Arc<RwLock<HashMap<WalletAddress, UserTier>>>,
    db: Arc<dyn TierWriter>,
    config: TierCacheConfig,
}

impl TierCache {
    pub fn new(db: Arc<dyn TierWriter>) -> Result<Self> {
        Self::new_with_config(db, TierCacheConfig::default())
    }

    pub fn new_with_config(db: Arc<dyn TierWriter>, config: TierCacheConfig) -> Result<Self> {
        Ok(Self {
            tiers: Arc::new(RwLock::new(HashMap::new())),
            db,
            config,
        })
    }

    /// Load all tier configurations from database
    pub async fn load_from_db(&self) -> Result<()> {
        let handler = self.db.clone();
        let tiers = tokio::task::spawn_blocking(move || handler.get_all_user_tiers_sync())
            .await
            .map_err(|e| anyhow::anyhow!("spawn_blocking join error: {}", e))??;
        let mut cache = self.tiers.write().await;
        for tier in tiers {
            cache.insert(tier.wallet_address, tier);
        }
        info!("Loaded {} user tiers from database", cache.len());
        Ok(())
    }

    /// Get tier for wallet (returns default if not found)
    pub async fn get_tier(&self, wallet: &WalletAddress) -> Option<UserTierData> {
        let cache = self.tiers.read().await;
        cache.get(wallet).map(|t| UserTierData {
            wallet_address: t.wallet_address,
            tier: t.tier.clone(),
        })
    }

    pub async fn get_tier_record(&self, wallet: &WalletAddress) -> Result<Option<UserTier>> {
        let handler = self.db.clone();
        let wallet_owned = *wallet;
        tokio::task::spawn_blocking(move || handler.get_user_tier_sync(&wallet_owned))
            .await
            .map_err(|e| anyhow::anyhow!("spawn_blocking join error: {}", e))?
    }

    pub async fn restore_tier_record(
        &self,
        wallet: &WalletAddress,
        previous_tier: Option<&UserTier>,
    ) -> Result<()> {
        match previous_tier {
            Some(previous) => self.set_tier(new_user_tier_from_existing(previous)).await,
            None => self.delete_tier(wallet).await,
        }
    }

    /// Get tier synchronously (for use in unified_engine)
    /// This is non-blocking and returns default if cache is locked or tier not found
    pub fn get_tier_sync(&self, wallet: &WalletAddress) -> UserTierData {
        if let Some(tier_data) = self.try_get_from_cache(wallet) {
            return tier_data;
        }

        default_user_tier(*wallet)
    }

    fn try_get_from_cache(&self, wallet: &WalletAddress) -> Option<UserTierData> {
        // Use try_read to avoid blocking
        if let Ok(cache) = self.tiers.try_read() {
            cache.get(wallet).map(|t| UserTierData {
                wallet_address: t.wallet_address,
                tier: t.tier.clone(),
            })
        } else {
            None
        }
    }

    /// Set or update tier
    pub async fn set_tier(&self, update: UserTierUpdate) -> Result<()> {
        debug!(
            "Setting tier for wallet {} to {}",
            update.wallet_address, update.tier
        );

        let wallet = update.wallet_address;
        let tier_name = update.tier.clone();
        let handler = self.db.clone();
        tokio::task::spawn_blocking(move || handler.save_user_tier_sync(&update))
            .await
            .map_err(|e| anyhow::anyhow!("spawn_blocking join error: {}", e))??;

        // Update cache
        let handler = self.db.clone();
        let tier = tokio::task::spawn_blocking(move || handler.get_user_tier_sync(&wallet))
            .await
            .map_err(|e| anyhow::anyhow!("spawn_blocking join error: {}", e))??
            .expect("Tier should exist after save");

        let mut cache = self.tiers.write().await;
        let wallet_address = tier.wallet_address;
        cache.insert(wallet_address, tier);

        info!(
            "Updated tier for wallet {} to {}",
            wallet_address, tier_name
        );

        Ok(())
    }

    /// Reset tier limits to defaults while preserving the explicit margin mode row.
    pub async fn delete_tier(&self, wallet: &WalletAddress) -> Result<()> {
        debug!("Resetting tier for wallet {} to defaults", wallet);

        let handler = self.db.clone();
        let wallet_owned = *wallet;
        tokio::task::spawn_blocking(move || handler.delete_user_tier_sync(&wallet_owned))
            .await
            .map_err(|e| anyhow::anyhow!("spawn_blocking join error: {}", e))??;

        let handler = self.db.clone();
        let wallet_owned = *wallet;
        let tier = tokio::task::spawn_blocking(move || handler.get_user_tier_sync(&wallet_owned))
            .await
            .map_err(|e| anyhow::anyhow!("spawn_blocking join error: {}", e))??
            .ok_or_else(|| anyhow!("missing user_tiers row after default tier reset"))?;

        let mut cache = self.tiers.write().await;
        cache.insert(*wallet, tier);

        info!("Reset tier for wallet {} to default tier2", wallet);

        Ok(())
    }

    /// Check if wallet is tier1 (restricted to long-only)
    pub fn is_tier1(&self, wallet: &WalletAddress) -> bool {
        let tier_data = self.get_tier_sync(wallet);
        tier_data.tier == "tier1"
    }

    /// Check if wallet is tier2 (unrestricted)
    pub fn is_tier2(&self, wallet: &WalletAddress) -> bool {
        let tier_data = self.get_tier_sync(wallet);
        tier_data.tier == "tier2"
    }

    // === Margin Mode Methods ===

    /// Get margin mode for wallet.
    ///
    /// This is the async version for use in handlers.
    pub async fn get_margin_mode(&self, wallet: &WalletAddress) -> Result<MarginMode> {
        let cache = self.tiers.read().await;
        // Missing user_tiers rows intentionally mean Standard margin. That is a
        // product default, not a fabricated financial input: Standard is the
        // most restrictive supported regime, while Portfolio margin is an
        // explicit opt-in stored in user_tiers and replayed into the engine.
        // Persisted rows must still parse successfully, because a corrupt
        // explicit margin mode would route risk checks through the wrong model.
        Ok(cache
            .get(wallet)
            .map(|tier| tier.margin_mode)
            .unwrap_or(MarginMode::Standard))
    }

    /// Get the existing margin mode for a wallet without fabricating a default.
    ///
    /// Returns `Ok(None)` only when no tier row exists. Invalid stored values
    /// still fail closed because the persisted margin mode is corrupted.
    pub async fn get_existing_margin_mode(
        &self,
        wallet: &WalletAddress,
    ) -> Result<Option<MarginMode>> {
        let handler = self.db.clone();
        let wallet_owned = *wallet;
        tokio::task::spawn_blocking(move || handler.get_existing_margin_mode_sync(&wallet_owned))
            .await
            .map_err(|e| anyhow::anyhow!("spawn_blocking join error: {}", e))?
    }

    /// Get an existing margin mode synchronously without fabricating a default.
    pub fn get_existing_margin_mode_sync(
        &self,
        wallet: &WalletAddress,
    ) -> Result<Option<MarginMode>> {
        let cache = self.tiers.try_read().map_err(|_| {
            anyhow!(
                "tier cache lock busy while reading margin mode for {}",
                wallet
            )
        })?;
        let Some(tier) = cache.get(wallet) else {
            return Ok(None);
        };
        Ok(Some(tier.margin_mode))
    }

    /// Get margin mode synchronously (for use in unified_engine).
    ///
    /// Non-blocking and fail-closed if the cache is locked. Missing rows
    /// intentionally resolve to Standard; see `get_margin_mode`.
    pub fn get_margin_mode_sync(&self, wallet: &WalletAddress) -> Result<MarginMode> {
        let cache = self.tiers.try_read().map_err(|_| {
            anyhow!(
                "tier cache lock busy while reading margin mode for {}",
                wallet
            )
        })?;
        Ok(cache
            .get(wallet)
            .map(|tier| tier.margin_mode)
            .unwrap_or(MarginMode::Standard))
    }

    /// Set margin mode for a wallet.
    ///
    /// Creates the tier entry if it doesn't exist.
    /// Returns the new version for use in event messages.
    pub async fn set_margin_mode(&self, wallet: &WalletAddress, mode: MarginMode) -> Result<i64> {
        debug!("Setting margin mode for wallet {} to {:?}", wallet, mode);

        // Atomically update margin mode and increment version in DB
        let handler = self.db.clone();
        let wallet_owned = *wallet;
        let new_version =
            tokio::task::spawn_blocking(move || handler.set_margin_mode_sync(&wallet_owned, mode))
                .await
                .map_err(|e| anyhow::anyhow!("spawn_blocking join error: {}", e))??;

        // Reload from DB (outside lock to avoid holding lock during I/O)
        let handler = self.db.clone();
        let wallet_owned = *wallet;
        let tier_opt =
            tokio::task::spawn_blocking(move || handler.get_user_tier_sync(&wallet_owned))
                .await
                .map_err(|e| anyhow::anyhow!("spawn_blocking join error: {}", e))??;
        if let Some(tier) = tier_opt {
            // Only insert if this is still the newest version
            // (a message from another process could have arrived with a higher version)
            let mut cache = self.tiers.write().await;
            let current_cached_version = cache.get(wallet).map(|t| t.version).unwrap_or(0);
            if tier.version > current_cached_version {
                cache.insert(*wallet, tier);
            }
        }

        info!(
            "Updated margin mode for wallet {} to {:?}, version={}",
            wallet, mode, new_version
        );

        Ok(new_version)
    }

    /// Create a margin mode row if one does not already exist.
    ///
    /// Returns `Some(version)` when a row was inserted. Existing rows are left
    /// unchanged so callers can safely use this for bootstrap paths.
    pub async fn insert_margin_mode_if_missing(
        &self,
        wallet: &WalletAddress,
        mode: MarginMode,
    ) -> Result<Option<i64>> {
        debug!(
            "Inserting margin mode for wallet {} to {:?} if missing",
            wallet, mode
        );

        let handler = self.db.clone();
        let wallet_owned = *wallet;
        let inserted_version = tokio::task::spawn_blocking(move || {
            handler.insert_margin_mode_if_missing_sync(&wallet_owned, mode)
        })
        .await
        .map_err(|e| anyhow::anyhow!("spawn_blocking join error: {}", e))??;

        if inserted_version.is_some() {
            let handler = self.db.clone();
            let wallet_owned = *wallet;
            let tier_opt =
                tokio::task::spawn_blocking(move || handler.get_user_tier_sync(&wallet_owned))
                    .await
                    .map_err(|e| anyhow::anyhow!("spawn_blocking join error: {}", e))??;

            if let Some(tier) = tier_opt {
                let mut cache = self.tiers.write().await;
                cache.insert(*wallet, tier);
            }
        }

        Ok(inserted_version)
    }

    /// Extract all wallet margin modes from the cache.
    ///
    /// Used at startup to seed engine-owned wallet_margin_modes so the engine
    /// never needs the fallback try_read() path.
    pub fn get_all_margin_modes_sync(&self) -> Result<HashMap<WalletAddress, MarginMode>> {
        let cache = self
            .tiers
            .try_read()
            .map_err(|_| anyhow!("tier cache lock busy during margin-mode seed"))?;
        Ok(cache
            .iter()
            .map(|(wallet, tier)| (*wallet, tier.margin_mode))
            .collect())
    }

    /// Extract all wallet trading limits from the cache.
    ///
    /// Returns an error if the lock cannot be acquired (e.g. contention during
    /// engine seed). Callers must not silently default to empty.
    pub fn get_all_trading_limits_sync(
        &self,
    ) -> Result<HashMap<WalletAddress, TradingLimits>, String> {
        let cache = self
            .tiers
            .try_read()
            .map_err(|_| "tier cache lock busy during trading-limits seed".to_string())?;
        Ok(cache
            .iter()
            .map(|(wallet, tier)| (*wallet, self.trading_limits_from_tier(tier)))
            .collect())
    }

    /// Extract all wallet tier strings from the cache.
    ///
    /// Returns an error if the lock cannot be acquired.
    pub fn get_all_tiers_sync(&self) -> Result<HashMap<WalletAddress, String>, String> {
        let cache = self
            .tiers
            .try_read()
            .map_err(|_| "tier cache lock busy during tier seed".to_string())?;
        Ok(cache
            .iter()
            .map(|(wallet, tier)| (*wallet, tier.tier.clone()))
            .collect())
    }

    /// Check if wallet is in Standard margin mode.
    pub fn is_standard_margin(&self, wallet: &WalletAddress) -> Result<bool> {
        Ok(self.get_margin_mode_sync(wallet)?.is_standard())
    }

    /// Check if wallet is in Portfolio margin mode.
    pub fn is_portfolio_margin(&self, wallet: &WalletAddress) -> Result<bool> {
        Ok(self.get_margin_mode_sync(wallet)?.is_portfolio())
    }

    /// Apply a margin mode update from the event bus (for cross-process cache sync).
    ///
    /// This is called when we receive a TierUpdate message.
    /// Only applies the update if the version is newer than the cached version.
    pub async fn apply_margin_mode_update(
        &self,
        wallet: WalletAddress,
        mode: MarginMode,
        version: i64,
    ) {
        // Early check with read lock to avoid unnecessary DB reads for stale messages
        let cached_version = {
            let cache = self.tiers.read().await;
            cache.get(&wallet).map(|t| t.version).unwrap_or(0)
        };

        if version <= cached_version {
            debug!(
                "Ignoring stale tier update: wallet={}, msg_version={}, cached_version={}",
                wallet, version, cached_version
            );
            return;
        }

        // Fetch from DB (outside lock to avoid holding lock during I/O)
        let handler = self.db.clone();
        let tier =
            match tokio::task::spawn_blocking(move || handler.get_user_tier_sync(&wallet)).await {
                Ok(Ok(Some(t))) => t,
                Ok(Ok(None)) => {
                    debug!("No tier found in DB for wallet {}", wallet);
                    return;
                }
                Ok(Err(e)) => {
                    tracing::error!("Failed to load tier from DB for {}: {}", wallet, e);
                    return;
                }
                Err(e) => {
                    tracing::error!("spawn_blocking failed for tier load {}: {}", wallet, e);
                    return;
                }
            };

        // Re-check version under write lock to prevent TOCTOU race
        let mut cache = self.tiers.write().await;
        let current_cached_version = cache.get(&wallet).map(|t| t.version).unwrap_or(0);

        // Only insert if still newer (another thread may have updated in the meantime)
        if tier.version > current_cached_version {
            debug!(
                "Applied margin mode update from event bus: wallet={}, mode={:?}, version={}",
                wallet, mode, tier.version
            );
            cache.insert(wallet, tier);
        } else {
            debug!(
                "Skipping insert after re-check: db_version={}, cached_version={}",
                tier.version, current_cached_version
            );
        }
    }

    /// Get the current version for a wallet's tier configuration.
    pub async fn get_version(&self, wallet: &WalletAddress) -> i64 {
        let cache = self.tiers.read().await;
        cache.get(wallet).map(|t| t.version).unwrap_or(0)
    }

    // === Trading Limits Methods ===

    /// Get trading limits for a wallet.
    ///
    /// Returns the wallet's configured limits, or defaults if not found.
    pub fn get_trading_limits(&self, wallet: &WalletAddress) -> TradingLimits {
        if let Ok(cache) = self.tiers.try_read() {
            if let Some(tier) = cache.get(wallet) {
                return self.trading_limits_from_tier(tier);
            }
        }
        self.default_trading_limits()
    }

    /// Get trading limits asynchronously.
    pub async fn get_trading_limits_async(&self, wallet: &WalletAddress) -> TradingLimits {
        let cache = self.tiers.read().await;
        if let Some(tier) = cache.get(wallet) {
            self.trading_limits_from_tier(tier)
        } else {
            self.default_trading_limits()
        }
    }

    /// Check if the number of open orders exceeds the limit.
    /// Returns true if the limit is exceeded.
    /// Returns false if limit is negative (unlimited).
    pub fn exceeds_open_orders_limit(&self, wallet: &WalletAddress, current_count: usize) -> bool {
        let limits = self.get_trading_limits(wallet);
        if limits.max_open_orders < 0 {
            return false;
        }
        current_count >= limits.max_open_orders as usize
    }

    /// Check if the number of positions exceeds the limit.
    /// Returns true if the limit is exceeded.
    /// Returns false if limit is -1 (unlimited).
    pub fn exceeds_positions_limit(&self, wallet: &WalletAddress, current_count: usize) -> bool {
        let limits = self.get_trading_limits(wallet);
        if limits.max_open_positions < 0 {
            return false; // -1 means unlimited
        }
        current_count >= limits.max_open_positions as usize
    }

    pub fn default_trading_limits(&self) -> TradingLimits {
        let mut limits = TradingLimits::default();
        limits.max_open_orders = self.config.max_open_orders_default;
        limits.max_open_positions = self.config.max_open_positions_default;
        limits
    }

    fn trading_limits_from_tier(&self, tier: &UserTier) -> TradingLimits {
        TradingLimits {
            max_open_orders: tier
                .max_open_orders
                .unwrap_or(self.config.max_open_orders_default),
            max_open_positions: tier.max_open_positions,
            orders_per_minute: tier.orders_per_minute,
            cancels_per_minute: tier.cancels_per_minute,
            api_requests_per_minute: tier.api_requests_per_minute,
        }
    }

    /// Set the margin mode for a wallet in the in-memory cache only (no DB write).
    /// Use this in tests that need to simulate PM vs SM accounts.
    #[cfg(any(test, feature = "test-utils"))]
    pub async fn set_margin_mode_in_memory(
        &self,
        wallet: &WalletAddress,
        mode: hypercall_types::MarginMode,
    ) {
        let mut cache = self.tiers.write().await;
        if let Some(tier) = cache.get_mut(wallet) {
            tier.margin_mode = mode;
        } else {
            cache.insert(
                *wallet,
                UserTier {
                    wallet_address: *wallet,
                    tier: "tier2".to_string(),
                    margin_mode: mode,
                    version: 1,
                    max_open_orders: None,
                    max_open_positions: 50,
                    orders_per_minute: 60,
                    cancels_per_minute: 120,
                    api_requests_per_minute: 600,
                    created_at: None,
                    updated_at: None,
                },
            );
        }
    }
}

fn new_user_tier_from_existing(tier: &UserTier) -> UserTierUpdate {
    UserTierUpdate {
        wallet_address: tier.wallet_address,
        tier: tier.tier.clone(),
        margin_mode: Some(tier.margin_mode),
        version: Some(tier.version),
        max_open_orders: tier.max_open_orders,
        max_open_positions: Some(tier.max_open_positions),
        orders_per_minute: Some(tier.orders_per_minute),
        cancels_per_minute: Some(tier.cancels_per_minute),
        api_requests_per_minute: Some(tier.api_requests_per_minute),
    }
}