finance_query_core/models/
logo.rs1use crate::client::FetchClient;
2use std::collections::HashMap;
3use std::env;
4use std::sync::Arc;
5use std::time::{Duration, Instant};
6use tokio::sync::{Mutex, RwLock};
7use tokio::time::timeout;
8use tracing::debug;
9
10#[derive(Debug)]
12struct CircuitBreaker {
13 failure_threshold: u32,
14 timeout_duration: Duration,
15 failure_count: u32,
16 last_failure_time: Option<Instant>,
17 state: CircuitState,
18}
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21enum CircuitState {
22 Closed,
23 Open,
24 HalfOpen,
25}
26
27impl CircuitBreaker {
28 fn new(failure_threshold: u32, timeout_duration: Duration) -> Self {
29 Self {
30 failure_threshold,
31 timeout_duration,
32 failure_count: 0,
33 last_failure_time: None,
34 state: CircuitState::Closed,
35 }
36 }
37
38 fn allow(&mut self) -> bool {
39 match self.state {
40 CircuitState::Closed => true,
41 CircuitState::Open => {
42 if let Some(last) = self.last_failure_time {
43 if last.elapsed() > self.timeout_duration {
44 self.state = CircuitState::HalfOpen;
45 return true;
46 }
47 }
48 false
49 }
50 CircuitState::HalfOpen => true,
51 }
52 }
53
54 fn record_success(&mut self) {
55 self.failure_count = 0;
56 self.state = CircuitState::Closed;
57 self.last_failure_time = None;
58 }
59
60 fn record_failure(&mut self) {
61 self.failure_count = self.failure_count.saturating_add(1);
62 self.last_failure_time = Some(Instant::now());
63 if self.failure_count >= self.failure_threshold {
64 self.state = CircuitState::Open;
65 }
66 }
67}
68
69#[derive(Debug, Clone)]
70struct CacheEntry {
71 value: String,
72 expires_at: Instant,
73}
74
75#[derive(Debug, Clone)]
86pub struct LogoFetcher {
87 fetch_client: Arc<FetchClient>,
88 cache: Arc<RwLock<HashMap<String, CacheEntry>>>,
89 circuit_breaker: Arc<Mutex<CircuitBreaker>>,
90 timeout: Duration,
91 cache_ttl: Duration,
92 enabled: bool,
93}
94
95impl LogoFetcher {
96 pub fn new(fetch_client: Arc<FetchClient>) -> Self {
104 let enabled = env::var("DISABLE_LOGO_FETCHING")
105 .map(|v| v.to_lowercase() != "true")
106 .unwrap_or(true);
107
108 let timeout_secs = env::var("LOGO_TIMEOUT_SECONDS")
109 .ok()
110 .and_then(|v| v.parse::<f64>().ok())
111 .unwrap_or(2.0)
112 .max(0.1);
113
114 let breaker_threshold = env::var("LOGO_CIRCUIT_BREAKER_THRESHOLD")
115 .ok()
116 .and_then(|v| v.parse::<u32>().ok())
117 .unwrap_or(5)
118 .max(1);
119
120 let breaker_timeout = env::var("LOGO_CIRCUIT_BREAKER_TIMEOUT")
121 .ok()
122 .and_then(|v| v.parse::<u64>().ok())
123 .unwrap_or(300);
124
125 Self {
126 fetch_client,
127 cache: Arc::new(RwLock::new(HashMap::new())),
128 circuit_breaker: Arc::new(Mutex::new(CircuitBreaker::new(
129 breaker_threshold,
130 Duration::from_secs(breaker_timeout),
131 ))),
132 timeout: Duration::from_secs_f64(timeout_secs),
133 cache_ttl: Duration::from_secs(60 * 60 * 24),
134 enabled,
135 }
136 }
137
138 pub fn is_enabled(&self) -> bool {
140 self.enabled
141 }
142
143 pub async fn fetch_logo(&self, symbol: &str, _website: Option<&str>) -> Option<String> {
148 if !self.enabled {
149 return None;
150 }
151
152 if symbol.is_empty() {
153 return None;
154 }
155
156 if !self.allow_request().await {
157 debug!("Logo circuit breaker is open; skipping fetch");
158 return None;
159 }
160
161 let cache_key = format!("logo:{}", symbol.to_uppercase());
162
163 if let Some(cached) = self.get_cached(&cache_key).await {
164 return Some(cached);
165 }
166
167 let logo_url = format!(
169 "https://img.logo.dev/ticker/{}?token=pk_NNp9abu9TMm9II6Z0666YA&format=png&fallback=404&size=50&theme=dark",
170 symbol.to_uppercase()
171 );
172
173 match self.try_fetch_logo(&logo_url).await {
174 Ok(url) => {
175 self.cache_value(&cache_key, url.clone()).await;
176 self.record_success().await;
177 Some(url)
178 }
179 Err(err) => {
180 debug!("Logo fetch failed for {}: {}", symbol, err);
181 self.record_failure().await;
182 None
183 }
184 }
185 }
186
187 async fn get_cached(&self, key: &str) -> Option<String> {
188 let now = Instant::now();
189 let mut cache = self.cache.write().await;
190 if let Some(entry) = cache.get(key) {
191 if entry.expires_at > now {
192 return Some(entry.value.clone());
193 }
194 }
195 cache.remove(key);
196 None
197 }
198
199 async fn cache_value(&self, key: &str, value: String) {
200 let expires_at = Instant::now() + self.cache_ttl;
201 let mut cache = self.cache.write().await;
202 cache.insert(key.to_string(), CacheEntry { value, expires_at });
203 }
204
205 async fn allow_request(&self) -> bool {
206 let mut breaker = self.circuit_breaker.lock().await;
207 breaker.allow()
208 }
209
210 async fn record_success(&self) {
211 let mut breaker = self.circuit_breaker.lock().await;
212 breaker.record_success();
213 }
214
215 async fn record_failure(&self) {
216 let mut breaker = self.circuit_breaker.lock().await;
217 breaker.record_failure();
218 }
219
220 async fn try_fetch_logo(&self, url: &str) -> Result<String, String> {
221 let fetch = self.fetch_client.fetch_response(url);
222 match timeout(self.timeout, fetch).await {
223 Ok(Ok(response)) => {
224 let status = response.status();
225 if status.is_success() {
226 Ok(response.url().to_string())
227 } else {
228 Err(format!("HTTP {}", status))
229 }
230 }
231 Ok(Err(err)) => Err(err.to_string()),
232 Err(_) => Err(format!("Timed out after {:?}", self.timeout)),
233 }
234 }
235}