finance-query 3.3.0

//! API keys scoped to one [`Providers`](crate::Providers) instance.
//!
//! The keyed adapters read their credentials from process-global
//! [`OnceLock`](std::sync::OnceLock) singletons, which allows exactly one key
//! per provider per process. A key configured on a builder is instead published
//! into a task-local for the duration of a dispatch call, so two `Providers`
//! can hold different keys for the same provider, and a key can be rotated
//! without restarting.
//!
//! `build_client` resolves in order: scoped key, singleton, environment
//! variable. Anything dispatched outside a scope behaves exactly as before.

// Everything below is reached from the keyed adapters' `build_client`, which
// only exists when one of their features is enabled.
#![allow(dead_code)]

use std::collections::HashMap;
use std::sync::{Arc, LazyLock, Mutex, OnceLock, RwLock, Weak};
use std::time::Duration;

use crate::rate_limiter::RateLimiter;

#[derive(Clone)]
pub(crate) struct ScopedKey {
    pub(crate) api_key: String,
    pub(crate) timeout: Duration,
    pub(crate) base_url: Option<String>,
    pub(crate) rpm: Option<u32>,
    limiter: Arc<OnceLock<Arc<RateLimiter>>>,
    http_clients: Arc<Mutex<RuntimeClients>>,
}

type RuntimeClients = HashMap<(tokio::runtime::Id, &'static str), reqwest::Client>;

impl ScopedKey {
    /// Connection pools are bound to the runtime that created them.
    pub(crate) fn http_client(
        &self,
        kind: &'static str,
        build: impl FnOnce() -> reqwest::Result<reqwest::Client>,
    ) -> reqwest::Result<reqwest::Client> {
        let mut clients = self.http_clients.lock().unwrap_or_else(|e| e.into_inner());
        let key = (tokio::runtime::Handle::current().id(), kind);
        if let Some(client) = clients.get(&key) {
            return Ok(client.clone());
        }
        let client = build()?;
        if clients.len() >= 8 {
            clients.clear();
        }
        clients.insert(key, client.clone());
        Ok(client)
    }

    pub(crate) fn new(api_key: String, timeout: Duration) -> Self {
        Self {
            api_key,
            timeout,
            base_url: None,
            rpm: None,
            limiter: Arc::new(OnceLock::new()),
            http_clients: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Rate limits are enforced per API key upstream, so each distinct key gets
    /// its own bucket rather than sharing the singleton's.
    pub(crate) fn limiter(
        &self,
        provider_key: &'static str,
        default_rate: f64,
    ) -> Arc<RateLimiter> {
        Arc::clone(self.limiter.get_or_init(|| {
            let explicit = self.rpm.map(|rpm| f64::from(rpm) / 60.0);
            shared_limiter(provider_key, &self.api_key, default_rate, explicit)
        }))
    }
}

pub(crate) type KeyMap = HashMap<&'static str, ScopedKey>;

tokio::task_local! {
    static SCOPED: Arc<KeyMap>;
}

type LimiterMap = HashMap<(&'static str, String), Weak<RateLimiter>>;

/// Weak so a rotated-out key's bucket is dropped with the `Providers` holding
/// it. The strong reference lives in that instance's [`ScopedKey`].
static LIMITERS: LazyLock<RwLock<LimiterMap>> = LazyLock::new(|| RwLock::new(HashMap::new()));

/// Instances sharing a key share one bucket. An explicit budget sets its rate,
/// the latest one winning; the provider default never changes an existing bucket.
fn shared_limiter(
    provider_key: &'static str,
    api_key: &str,
    default_rate: f64,
    explicit: Option<f64>,
) -> Arc<RateLimiter> {
    let id = (provider_key, api_key.to_string());
    let reuse = |limiter: Arc<RateLimiter>| {
        if let Some(rate) = explicit {
            limiter.set_rate(rate);
        }
        limiter
    };
    if let Some(limiter) = LIMITERS
        .read()
        .unwrap_or_else(|e| e.into_inner())
        .get(&id)
        .and_then(Weak::upgrade)
    {
        return reuse(limiter);
    }
    let mut limiters = LIMITERS.write().unwrap_or_else(|e| e.into_inner());
    if let Some(limiter) = limiters.get(&id).and_then(Weak::upgrade) {
        return reuse(limiter);
    }
    limiters.retain(|_, held| held.strong_count() > 0);
    let limiter = Arc::new(RateLimiter::new(explicit.unwrap_or(default_rate)));
    limiters.insert(id, Arc::downgrade(&limiter));
    limiter
}

pub(crate) fn scoped_key(provider_key: &str) -> Option<ScopedKey> {
    SCOPED
        .try_with(|keys| keys.get(provider_key).cloned())
        .ok()
        .flatten()
}

pub(crate) async fn scope<F>(keys: Arc<KeyMap>, f: F) -> F::Output
where
    F: std::future::Future,
{
    if keys.is_empty() {
        return f.await;
    }
    SCOPED.scope(keys, f).await
}

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

    #[tokio::test]
    async fn an_explicit_budget_sets_the_shared_rate_and_the_default_never_lowers_it() {
        let paid = shared_limiter("polygon", "explicit-budget", 0.1, Some(100.0));
        let unset = shared_limiter("polygon", "explicit-budget", 0.1, None);
        assert!(Arc::ptr_eq(&paid, &unset));
        assert!(unset.available_estimate().unwrap() > 50.0);

        let lowered = shared_limiter("polygon", "explicit-budget", 0.1, Some(1.0));
        assert!(Arc::ptr_eq(&paid, &lowered));
        assert!(paid.available_estimate().unwrap() <= 1.0);
        paid.acquire().await;
        assert!(
            tokio::time::timeout(Duration::from_millis(20), unset.acquire())
                .await
                .is_err()
        );

        shared_limiter("polygon", "explicit-budget", 0.1, Some(1000.0));
        tokio::time::timeout(Duration::from_millis(50), unset.acquire())
            .await
            .unwrap();
    }

    fn map(pairs: &[(&'static str, &str)]) -> Arc<KeyMap> {
        Arc::new(
            pairs
                .iter()
                .map(|(p, k)| {
                    (
                        *p,
                        ScopedKey::new((*k).to_string(), Duration::from_secs(30)),
                    )
                })
                .collect(),
        )
    }

    #[tokio::test]
    async fn no_scope_yields_no_key() {
        assert!(scoped_key("fmp").is_none());
    }

    #[tokio::test]
    async fn a_scope_publishes_its_own_keys() {
        scope(map(&[("fmp", "abc")]), async {
            assert_eq!(scoped_key("fmp").map(|k| k.api_key), Some("abc".into()));
            assert!(scoped_key("polygon").is_none());
        })
        .await;
    }

    #[tokio::test]
    async fn an_empty_map_leaves_the_scope_unset() {
        scope(Arc::new(KeyMap::new()), async {
            assert!(scoped_key("fmp").is_none());
        })
        .await;
    }

    #[tokio::test]
    async fn concurrent_scopes_do_not_see_each_other() {
        let one = tokio::spawn(scope(map(&[("fmp", "first")]), async {
            tokio::task::yield_now().await;
            scoped_key("fmp").map(|k| k.api_key)
        }));
        let two = tokio::spawn(scope(map(&[("fmp", "second")]), async {
            tokio::task::yield_now().await;
            scoped_key("fmp").map(|k| k.api_key)
        }));
        assert_eq!(one.await.unwrap(), Some("first".into()));
        assert_eq!(two.await.unwrap(), Some("second".into()));
    }

    #[test]
    fn each_key_gets_its_own_limiter() {
        let a = ScopedKey::new("key-a".into(), Duration::from_secs(30));
        let b = ScopedKey::new("key-b".into(), Duration::from_secs(30));
        assert!(Arc::ptr_eq(&a.limiter("fmp", 5.0), &a.limiter("fmp", 5.0)));
        assert!(!Arc::ptr_eq(&a.limiter("fmp", 5.0), &b.limiter("fmp", 5.0)));
    }

    #[test]
    fn two_instances_of_one_key_share_a_bucket() {
        let first = ScopedKey::new("shared".into(), Duration::from_secs(30));
        let second = ScopedKey::new("shared".into(), Duration::from_secs(30));
        assert!(Arc::ptr_eq(
            &first.limiter("polygon", 5.0),
            &second.limiter("polygon", 5.0)
        ));
    }

    #[test]
    fn a_dropped_key_stops_being_tracked() {
        {
            let key = ScopedKey::new("transient-xyz".into(), Duration::from_secs(30));
            let _live = key.limiter("fred", 5.0);
            // Assert on this entry, not the map size: parallel tests prune
            // their own dropped keys and shift the count.
            assert!(
                LIMITERS
                    .read()
                    .unwrap()
                    .contains_key(&("fred", "transient-xyz".to_string()))
            );
        }
        let _prune =
            ScopedKey::new("prune-trigger".into(), Duration::from_secs(30)).limiter("fred", 5.0);
        assert!(
            !LIMITERS
                .read()
                .unwrap()
                .contains_key(&("fred", "transient-xyz".to_string()))
        );
    }
}