vaea_flash_sdk/
warm_cache.rs1use crate::types::{CapacityResponse, TokenCapacity, VAEA_API_URL};
8use crate::local_builder::update_registry_from_capacity;
9use std::sync::{Arc, Mutex};
10use tokio::task::JoinHandle;
11
12pub struct WarmCache {
29 api_url: String,
30 refresh_ms: u64,
31 inner: Arc<Mutex<WarmCacheInner>>,
32 task_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
33}
34
35struct WarmCacheInner {
36 capacity: Option<CapacityResponse>,
37 listeners: Vec<Box<dyn Fn(&CapacityResponse) + Send + 'static>>,
38}
39
40impl WarmCache {
41 pub fn new(api_url: Option<&str>, refresh_ms: Option<u64>) -> Self {
46 Self {
47 api_url: api_url.unwrap_or(VAEA_API_URL).to_string(),
48 refresh_ms: refresh_ms.unwrap_or(2000),
49 inner: Arc::new(Mutex::new(WarmCacheInner {
50 capacity: None,
51 listeners: Vec::new(),
52 })),
53 task_handle: Arc::new(Mutex::new(None)),
54 }
55 }
56
57 pub async fn start(&self) {
59 self.refresh().await;
61
62 let api_url = self.api_url.clone();
64 let refresh_ms = self.refresh_ms;
65 let inner = Arc::clone(&self.inner);
66
67 let handle = tokio::spawn(async move {
68 let client = reqwest::Client::new();
69 let mut interval = tokio::time::interval(std::time::Duration::from_millis(refresh_ms));
70 loop {
71 interval.tick().await;
72 if let Ok(res) = client.get(&format!("{}/v1/capacity", api_url)).send().await {
73 if let Ok(capacity) = res.json::<CapacityResponse>().await {
74 update_registry_from_capacity(&capacity.tokens);
76 let mut guard = inner.lock().unwrap();
77 guard.capacity = Some(capacity.clone());
78 for listener in &guard.listeners {
79 listener(&capacity);
80 }
81 }
82 }
83 }
84 });
85
86 *self.task_handle.lock().unwrap() = Some(handle);
87 }
88
89 pub fn stop(&self) {
91 if let Some(handle) = self.task_handle.lock().unwrap().take() {
92 handle.abort();
93 }
94 }
95
96 pub fn on_update<F: Fn(&CapacityResponse) + Send + 'static>(&self, handler: F) {
98 self.inner.lock().unwrap().listeners.push(Box::new(handler));
99 }
100
101 pub fn get_capacity(&self) -> Option<CapacityResponse> {
103 self.inner.lock().unwrap().capacity.clone()
104 }
105
106 pub fn get_token_capacity(&self, symbol: &str) -> Option<TokenCapacity> {
108 let guard = self.inner.lock().unwrap();
109 guard.capacity.as_ref()?.tokens.iter()
110 .find(|t| t.symbol.eq_ignore_ascii_case(symbol))
111 .cloned()
112 }
113
114 pub fn is_warm(&self) -> bool {
116 self.inner.lock().unwrap().capacity.is_some()
117 }
118
119 async fn refresh(&self) {
120 let client = reqwest::Client::new();
121 if let Ok(res) = client.get(&format!("{}/v1/capacity", self.api_url)).send().await {
122 if let Ok(capacity) = res.json::<CapacityResponse>().await {
123 update_registry_from_capacity(&capacity.tokens);
125 let mut guard = self.inner.lock().unwrap();
126 guard.capacity = Some(capacity.clone());
127 for listener in &guard.listeners {
128 listener(&capacity);
129 }
130 }
131 }
132 }
133}