Skip to main content

alpaca_data/cache/
state.rs

1use std::collections::{HashMap, HashSet};
2use std::time::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 StockBarsRequest {
12    pub key: String,
13    pub symbols: Vec<String>,
14    pub timeframe: TimeFrame,
15    pub start: Option<String>,
16    pub end: Option<String>,
17    pub limit: u32,
18    pub adjustment: Option<Adjustment>,
19    pub feed: Option<DataFeed>,
20    pub currency: Option<Currency>,
21    pub chunk_size: usize,
22}
23
24impl StockBarsRequest {
25    pub(crate) fn normalized(mut self) -> Self {
26        self.symbols = normalize_stock_list(&self.symbols);
27        self
28    }
29
30    pub(crate) fn merge_from(&mut self, next: &Self) {
31        self.symbols = merge_values(&self.symbols, &next.symbols);
32        self.timeframe = next.timeframe.clone();
33        self.start = next.start.clone();
34        self.end = next.end.clone();
35        self.limit = next.limit;
36        self.adjustment = next.adjustment.clone();
37        self.feed = next.feed;
38        self.currency = next.currency.clone();
39        self.chunk_size = next.chunk_size;
40    }
41}
42
43#[derive(Debug, Default)]
44pub(crate) struct SnapshotCache<T> {
45    pub subscribed: HashSet<String>,
46    pub values: HashMap<String, T>,
47    pub empty: HashSet<String>,
48    pub updated_at: Option<SystemTime>,
49}
50
51impl<T: Clone> SnapshotCache<T> {
52    pub(crate) fn reconcile(
53        &mut self,
54        requested: &[String],
55        fetched: &HashMap<String, T>,
56        updated_at: SystemTime,
57    ) -> usize {
58        let mut count = 0;
59        let mut seen = HashSet::new();
60
61        for key in requested {
62            if !seen.insert(key) {
63                continue;
64            }
65
66            self.subscribed.insert(key.clone());
67            if let Some(value) = fetched.get(key) {
68                self.values.insert(key.clone(), value.clone());
69                self.empty.remove(key);
70                count += 1;
71            } else {
72                self.values.remove(key);
73                self.empty.insert(key.clone());
74            }
75        }
76
77        if !seen.is_empty() {
78            self.updated_at = Some(updated_at);
79        }
80        count
81    }
82}
83
84#[derive(Debug, Default)]
85pub(crate) struct StockBarsCache {
86    pub requests: HashMap<String, StockBarsRequest>,
87    pub values: HashMap<String, BarsMap>,
88    pub empty: HashMap<String, HashSet<String>>,
89    pub updated_at: HashMap<String, SystemTime>,
90}
91
92#[derive(Debug, Default)]
93pub(crate) struct CacheState {
94    pub stocks: SnapshotCache<stocks::Snapshot>,
95    pub options: SnapshotCache<options::Snapshot>,
96    pub bars: StockBarsCache,
97}
98
99pub(crate) fn normalize_values<S: AsRef<str>>(values: &[S]) -> Vec<String> {
100    let mut normalized = Vec::new();
101    let mut seen = HashSet::new();
102    for value in values {
103        let value = value.as_ref().trim();
104        if !value.is_empty() && seen.insert(value.to_string()) {
105            normalized.push(value.to_string());
106        }
107    }
108    normalized
109}
110
111pub(crate) fn normalize_stock_symbols<S: AsRef<str>>(symbols: &[S]) -> Vec<(String, String)> {
112    normalize_values(symbols)
113        .into_iter()
114        .map(|symbol| {
115            let resolved = stocks::display_stock_symbol(&symbol);
116            (symbol, resolved)
117        })
118        .collect()
119}
120
121pub(crate) fn normalize_option_symbols<S: AsRef<str>>(symbols: &[S]) -> Vec<String> {
122    let normalized = normalize_values(symbols);
123    let mut values = Vec::new();
124    let mut seen = HashSet::new();
125    for symbol in normalized {
126        let symbol = option_contract_symbol(&symbol);
127        if !symbol.is_empty() && seen.insert(symbol.clone()) {
128            values.push(symbol);
129        }
130    }
131    values
132}
133
134pub(crate) fn collect_cached_hits<T: Clone>(
135    requested: &[String],
136    cached: &HashMap<String, T>,
137    empty: &HashSet<String>,
138) -> (HashMap<String, T>, Vec<String>) {
139    let mut hits = HashMap::new();
140    let mut missing = Vec::new();
141    for key in requested {
142        if let Some(value) = cached.get(key) {
143            hits.insert(key.clone(), value.clone());
144        } else if !empty.contains(key) {
145            missing.push(key.clone());
146        }
147    }
148    (hits, missing)
149}
150
151fn normalize_stock_list<S: AsRef<str>>(symbols: &[S]) -> Vec<String> {
152    normalize_stock_symbols(symbols)
153        .into_iter()
154        .map(|(_, resolved)| resolved)
155        .collect()
156}
157
158fn merge_values(current: &[String], next: &[String]) -> Vec<String> {
159    let mut merged = Vec::new();
160    let mut seen = HashSet::new();
161    for value in current.iter().chain(next.iter()) {
162        if seen.insert(value.clone()) {
163            merged.push(value.clone());
164        }
165    }
166    merged
167}