Skip to main content

alpaca_data/cache/
state.rs

1use std::collections::{HashMap, HashSet};
2use std::time::{Duration, Instant, SystemTime};
3
4use crate::options;
5use crate::stocks::{self, Adjustment, BarPoint, Currency, DataFeed, TimeFrame};
6use crate::symbols::option_contract_symbol;
7
8pub(crate) type BarsMap = HashMap<String, Vec<BarPoint>>;
9
10#[derive(Debug, Clone)]
11pub struct CachedEntry<T> {
12    pub value: T,
13    pub stored_at: Instant,
14}
15
16impl<T> CachedEntry<T> {
17    pub fn new(value: T) -> Self {
18        Self {
19            value,
20            stored_at: Instant::now(),
21        }
22    }
23
24    pub fn is_fresh(&self, ttl: Duration, now: Instant) -> bool {
25        is_timestamp_fresh(self.stored_at, ttl, now)
26    }
27}
28
29pub fn is_timestamp_fresh(stored_at: Instant, ttl: Duration, now: Instant) -> bool {
30    now.saturating_duration_since(stored_at) < ttl
31}
32
33#[derive(Debug, Clone)]
34pub struct StockBarsRequest {
35    pub key: String,
36    pub symbols: Vec<String>,
37    pub timeframe: TimeFrame,
38    pub start: Option<String>,
39    pub end: Option<String>,
40    pub limit: u32,
41    pub adjustment: Option<Adjustment>,
42    pub feed: Option<DataFeed>,
43    pub currency: Option<Currency>,
44    pub chunk_size: usize,
45}
46
47impl StockBarsRequest {
48    pub(crate) fn normalized(mut self) -> Self {
49        self.symbols = normalize_stock_list(&self.symbols);
50        self
51    }
52
53    pub(crate) fn merge_from(&mut self, next: &Self) {
54        self.symbols = merge_values(&self.symbols, &next.symbols);
55        self.timeframe = next.timeframe.clone();
56        self.start = next.start.clone();
57        self.end = next.end.clone();
58        self.limit = next.limit;
59        self.adjustment = next.adjustment.clone();
60        self.feed = next.feed;
61        self.currency = next.currency.clone();
62        self.chunk_size = next.chunk_size;
63    }
64}
65
66#[derive(Debug, Default)]
67pub(crate) struct SnapshotCache<T> {
68    pub subscribed: HashSet<String>,
69    pub values: HashMap<String, CachedEntry<T>>,
70    pub empty: HashMap<String, Instant>,
71    pub updated_at: Option<SystemTime>,
72}
73
74impl<T: Clone> SnapshotCache<T> {
75    pub(crate) fn reconcile(
76        &mut self,
77        requested: &[String],
78        fetched: &HashMap<String, T>,
79        updated_at: SystemTime,
80    ) -> usize {
81        let mut count = 0;
82        let mut seen = HashSet::new();
83        let stored_at = Instant::now();
84
85        for key in requested {
86            if !seen.insert(key) {
87                continue;
88            }
89
90            self.subscribed.insert(key.clone());
91            if let Some(value) = fetched.get(key) {
92                self.values.insert(
93                    key.clone(),
94                    CachedEntry {
95                        value: value.clone(),
96                        stored_at,
97                    },
98                );
99                self.empty.remove(key);
100                count += 1;
101            } else {
102                self.values.remove(key);
103                self.empty.insert(key.clone(), stored_at);
104            }
105        }
106
107        if !seen.is_empty() {
108            self.updated_at = Some(updated_at);
109        }
110        count
111    }
112}
113
114#[derive(Debug, Default)]
115pub(crate) struct StockBarsCache {
116    pub requests: HashMap<String, StockBarsRequest>,
117    pub values: HashMap<String, HashMap<String, CachedEntry<Vec<BarPoint>>>>,
118    pub empty: HashMap<String, HashMap<String, Instant>>,
119    pub updated_at: HashMap<String, SystemTime>,
120}
121
122#[derive(Debug, Default)]
123pub(crate) struct CacheState {
124    pub stocks: SnapshotCache<stocks::Snapshot>,
125    pub options: SnapshotCache<options::Snapshot>,
126    pub bars: StockBarsCache,
127}
128
129pub(crate) fn normalize_values<S: AsRef<str>>(values: &[S]) -> Vec<String> {
130    let mut normalized = Vec::new();
131    let mut seen = HashSet::new();
132    for value in values {
133        let value = value.as_ref().trim();
134        if !value.is_empty() && seen.insert(value.to_string()) {
135            normalized.push(value.to_string());
136        }
137    }
138    normalized
139}
140
141pub(crate) fn normalize_stock_symbols<S: AsRef<str>>(symbols: &[S]) -> Vec<(String, String)> {
142    normalize_values(symbols)
143        .into_iter()
144        .map(|symbol| {
145            let resolved = stocks::display_stock_symbol(&symbol);
146            (symbol, resolved)
147        })
148        .collect()
149}
150
151pub(crate) fn normalize_option_symbols<S: AsRef<str>>(symbols: &[S]) -> Vec<String> {
152    let normalized = normalize_values(symbols);
153    let mut values = Vec::new();
154    let mut seen = HashSet::new();
155    for symbol in normalized {
156        let symbol = option_contract_symbol(&symbol);
157        if !symbol.is_empty() && seen.insert(symbol.clone()) {
158            values.push(symbol);
159        }
160    }
161    values
162}
163
164pub fn collect_cached_hits<T: Clone>(
165    requested: &[String],
166    cached: &HashMap<String, CachedEntry<T>>,
167    empty: &HashMap<String, Instant>,
168    ttl: Duration,
169    now: Instant,
170) -> (HashMap<String, T>, Vec<String>) {
171    let mut hits = HashMap::new();
172    let mut missing = Vec::new();
173    for key in requested {
174        if let Some(entry) = cached.get(key) {
175            if entry.is_fresh(ttl, now) {
176                hits.insert(key.clone(), entry.value.clone());
177            } else {
178                missing.push(key.clone());
179            }
180        } else if empty
181            .get(key)
182            .is_some_and(|stored_at| is_timestamp_fresh(*stored_at, ttl, now))
183        {
184            continue;
185        } else {
186            missing.push(key.clone());
187        }
188    }
189    (hits, missing)
190}
191
192pub(crate) fn unwrap_bars_map(
193    cached: &HashMap<String, CachedEntry<Vec<BarPoint>>>,
194) -> BarsMap {
195    cached
196        .iter()
197        .map(|(symbol, entry)| (symbol.clone(), entry.value.clone()))
198        .collect()
199}
200
201pub(crate) fn missing_bar_symbols(
202    requested: &[String],
203    cached: Option<&HashMap<String, CachedEntry<Vec<BarPoint>>>>,
204    empty: Option<&HashMap<String, Instant>>,
205    ttl: Duration,
206    now: Instant,
207) -> Vec<String> {
208    requested
209        .iter()
210        .filter(|symbol| {
211            if cached
212                .and_then(|values| values.get(*symbol))
213                .is_some_and(|entry| entry.is_fresh(ttl, now))
214            {
215                return false;
216            }
217            if empty
218                .and_then(|values| values.get(*symbol))
219                .is_some_and(|stored_at| is_timestamp_fresh(*stored_at, ttl, now))
220            {
221                return false;
222            }
223            true
224        })
225        .cloned()
226        .collect()
227}
228
229fn normalize_stock_list<S: AsRef<str>>(symbols: &[S]) -> Vec<String> {
230    normalize_stock_symbols(symbols)
231        .into_iter()
232        .map(|(_, resolved)| resolved)
233        .collect()
234}
235
236fn merge_values(current: &[String], next: &[String]) -> Vec<String> {
237    let mut merged = Vec::new();
238    let mut seen = HashSet::new();
239    for value in current.iter().chain(next.iter()) {
240        if seen.insert(value.clone()) {
241            merged.push(value.clone());
242        }
243    }
244    merged
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250
251    #[test]
252    fn expired_price_entries_are_treated_as_missing() {
253        let now = Instant::now();
254        let stale = now
255            .checked_sub(Duration::from_secs(20))
256            .expect("ttl test clock");
257        let mut cached = HashMap::new();
258        cached.insert(
259            "AAPL".to_string(),
260            CachedEntry {
261                value: 1,
262                stored_at: stale,
263            },
264        );
265        cached.insert(
266            "MSFT".to_string(),
267            CachedEntry {
268                value: 2,
269                stored_at: now,
270            },
271        );
272
273        let (hits, missing) = collect_cached_hits(
274            &[
275                "AAPL".to_string(),
276                "MSFT".to_string(),
277                "GOOG".to_string(),
278            ],
279            &cached,
280            &HashMap::new(),
281            Duration::from_secs(15),
282            now,
283        );
284
285        assert_eq!(hits.get("MSFT"), Some(&2));
286        assert!(!hits.contains_key("AAPL"));
287        assert_eq!(missing, vec!["AAPL".to_string(), "GOOG".to_string()]);
288    }
289
290    #[test]
291    fn expired_empty_entries_are_retried() {
292        let now = Instant::now();
293        let stale = now
294            .checked_sub(Duration::from_secs(20))
295            .expect("ttl test clock");
296        let mut empty = HashMap::new();
297        empty.insert("AAPL".to_string(), stale);
298        empty.insert("MSFT".to_string(), now);
299
300        let (hits, missing) = collect_cached_hits(
301            &["AAPL".to_string(), "MSFT".to_string()],
302            &HashMap::<String, CachedEntry<i32>>::new(),
303            &empty,
304            Duration::from_secs(15),
305            now,
306        );
307
308        assert!(hits.is_empty());
309        assert_eq!(missing, vec!["AAPL".to_string()]);
310    }
311}