Skip to main content

hive_console_sdk/
persisted_documents.rs

1use std::time::Duration;
2
3use crate::agent::usage_agent::non_empty_string;
4use crate::circuit_breaker::CircuitBreakerBuilder;
5use crate::circuit_breaker::CircuitBreakerError;
6use moka::future::Cache;
7use recloser::AsyncRecloser;
8use reqwest::header::HeaderMap;
9use reqwest::header::HeaderValue;
10use reqwest_middleware::ClientBuilder;
11use reqwest_middleware::ClientWithMiddleware;
12use reqwest_retry::RetryTransientMiddleware;
13use retry_policies::policies::ExponentialBackoff;
14use tracing::{debug, info, warn};
15
16#[derive(Debug)]
17pub struct PersistedDocumentsManager {
18    client: ClientWithMiddleware,
19    cache: Cache<String, String>,
20    negative_cache: Option<Cache<String, ()>>,
21    endpoints_with_circuit_breakers: Vec<(String, AsyncRecloser)>,
22}
23
24#[derive(Debug, thiserror::Error, Clone)]
25pub enum PersistedDocumentsError {
26    #[error("Failed to read body: {0}")]
27    FailedToReadBody(String),
28    #[error("Failed to parse body: {0}")]
29    FailedToParseBody(String),
30    #[error("Persisted document not found.")]
31    DocumentNotFound,
32    #[error("Failed to locate the persisted document key in request.")]
33    KeyNotFound,
34    #[error("Failed to validate persisted document")]
35    FailedToFetchFromCDN(String),
36    #[error("Failed to read CDN response body")]
37    FailedToReadCDNResponse(String),
38    #[error("No persisted document provided, or document id cannot be resolved.")]
39    PersistedDocumentRequired,
40    #[error("Missing required configuration option: {0}")]
41    MissingConfigurationOption(String),
42    #[error("Invalid CDN key {0}")]
43    InvalidCDNKey(String),
44    #[error("Failed to create HTTP client: {0}")]
45    HTTPClientCreationError(String),
46    #[error("unable to create circuit breaker: {0}")]
47    CircuitBreakerCreationError(String),
48    #[error("rejected by the circuit breaker")]
49    CircuitBreakerRejected,
50    #[error("unknown error")]
51    Unknown,
52}
53
54impl From<reqwest_middleware::Error> for PersistedDocumentsError {
55    fn from(err: reqwest_middleware::Error) -> Self {
56        PersistedDocumentsError::FailedToFetchFromCDN(err.to_string())
57    }
58}
59
60impl From<serde_json::Error> for PersistedDocumentsError {
61    fn from(err: serde_json::Error) -> Self {
62        PersistedDocumentsError::FailedToParseBody(err.to_string())
63    }
64}
65
66impl From<CircuitBreakerError> for PersistedDocumentsError {
67    fn from(err: CircuitBreakerError) -> Self {
68        PersistedDocumentsError::CircuitBreakerCreationError(err.to_string())
69    }
70}
71
72impl PersistedDocumentsError {
73    pub fn message(&self) -> String {
74        self.to_string()
75    }
76
77    pub fn code(&self) -> String {
78        match self {
79            PersistedDocumentsError::FailedToReadBody(_) => "FAILED_TO_READ_BODY".into(),
80            PersistedDocumentsError::FailedToParseBody(_) => "FAILED_TO_PARSE_BODY".into(),
81            PersistedDocumentsError::DocumentNotFound => "PERSISTED_DOCUMENT_NOT_FOUND".into(),
82            PersistedDocumentsError::KeyNotFound => "PERSISTED_DOCUMENT_KEY_NOT_FOUND".into(),
83            PersistedDocumentsError::FailedToFetchFromCDN(_) => "FAILED_TO_FETCH_FROM_CDN".into(),
84            PersistedDocumentsError::FailedToReadCDNResponse(_) => {
85                "FAILED_TO_READ_CDN_RESPONSE".into()
86            }
87            PersistedDocumentsError::PersistedDocumentRequired => {
88                "PERSISTED_DOCUMENT_REQUIRED".into()
89            }
90            PersistedDocumentsError::MissingConfigurationOption(_) => {
91                "MISSING_CONFIGURATION_OPTION".into()
92            }
93            PersistedDocumentsError::InvalidCDNKey(_) => "INVALID_CDN_KEY".into(),
94            PersistedDocumentsError::HTTPClientCreationError(_) => {
95                "HTTP_CLIENT_CREATION_ERROR".into()
96            }
97            PersistedDocumentsError::CircuitBreakerCreationError(_) => {
98                "CIRCUIT_BREAKER_CREATION_ERROR".into()
99            }
100            PersistedDocumentsError::CircuitBreakerRejected => "CIRCUIT_BREAKER_REJECTED".into(),
101            PersistedDocumentsError::Unknown => "UNKNOWN_ERROR".into(),
102        }
103    }
104}
105
106impl PersistedDocumentsManager {
107    pub fn builder() -> PersistedDocumentsManagerBuilder {
108        PersistedDocumentsManagerBuilder::default()
109    }
110    async fn resolve_from_endpoint(
111        &self,
112        endpoint: &str,
113        document_id: &str,
114        circuit_breaker: &AsyncRecloser,
115    ) -> Result<String, PersistedDocumentsError> {
116        let cdn_document_id = str::replace(document_id, "~", "/");
117        let cdn_artifact_url = format!("{}/apps/{}", endpoint, cdn_document_id);
118        info!(
119            "Fetching document {} from CDN: {}",
120            document_id, cdn_artifact_url
121        );
122        let response_fut = self.client.get(cdn_artifact_url).send();
123
124        let response = circuit_breaker
125            .call(response_fut)
126            .await
127            .map_err(|e| match e {
128                recloser::Error::Inner(e) => PersistedDocumentsError::from(e),
129                recloser::Error::Rejected => PersistedDocumentsError::CircuitBreakerRejected,
130            })?;
131
132        if response.status().is_success() {
133            let document = response
134                .text()
135                .await
136                .map_err(|e| PersistedDocumentsError::FailedToReadCDNResponse(e.to_string()))?;
137            debug!("Document fetched from CDN: {}", document);
138
139            return Ok(document);
140        }
141
142        let status = response.status();
143
144        warn!(
145            "Document fetch from CDN failed: HTTP {}, Body: {:?}",
146            status,
147            response
148                .text()
149                .await
150                .unwrap_or_else(|_| "Unavailable".to_string())
151        );
152
153        Err(PersistedDocumentsError::DocumentNotFound)
154    }
155
156    /// Resolves the document from the cache, or from the CDN
157    pub async fn resolve_document(
158        &self,
159        document_id: &str,
160    ) -> Result<String, PersistedDocumentsError> {
161        if let Some(negative_cache) = &self.negative_cache {
162            if negative_cache.get(document_id).await.is_some() {
163                debug!(
164                    "Document {} found in negative cache, skipping CDN fetch",
165                    document_id
166                );
167                return Err(PersistedDocumentsError::DocumentNotFound);
168            }
169        }
170
171        if let Some(cached_document) = self.cache.get(document_id).await {
172            return Ok(cached_document);
173        }
174
175        let result = self
176            .cache
177            .try_get_with_by_ref(document_id, async {
178                debug!(
179                    "Document {} not found in cache. Fetching from CDN",
180                    document_id
181                );
182
183                let mut last_error: Option<PersistedDocumentsError> = None;
184                for (endpoint, circuit_breaker) in self.endpoints_with_circuit_breakers.iter() {
185                    match self
186                        .resolve_from_endpoint(endpoint, document_id, circuit_breaker)
187                        .await
188                    {
189                        Ok(document) => return Ok(document),
190                        Err(error) => last_error = Some(error),
191                    }
192                }
193
194                Err(last_error.unwrap_or(PersistedDocumentsError::Unknown))
195            })
196            .await
197            .map_err(|error| error.as_ref().clone());
198
199        if matches!(&result, Err(PersistedDocumentsError::DocumentNotFound)) {
200            if let Some(negative_cache) = &self.negative_cache {
201                negative_cache.insert(document_id.to_string(), ()).await;
202            }
203        }
204
205        result
206    }
207}
208
209pub struct PersistedDocumentsManagerBuilder {
210    key: Option<String>,
211    endpoints: Vec<String>,
212    accept_invalid_certs: bool,
213    connect_timeout: Duration,
214    request_timeout: Duration,
215    retry_policy: ExponentialBackoff,
216    cache_size: u64,
217    negative_cache_ttl: Option<Duration>,
218    user_agent: Option<String>,
219    circuit_breaker: CircuitBreakerBuilder,
220}
221
222impl Default for PersistedDocumentsManagerBuilder {
223    fn default() -> Self {
224        Self {
225            key: None,
226            endpoints: vec![],
227            accept_invalid_certs: false,
228            connect_timeout: Duration::from_secs(5),
229            request_timeout: Duration::from_secs(15),
230            retry_policy: ExponentialBackoff::builder().build_with_max_retries(3),
231            cache_size: 10_000,
232            negative_cache_ttl: None,
233            user_agent: None,
234            circuit_breaker: CircuitBreakerBuilder::default(),
235        }
236    }
237}
238
239impl PersistedDocumentsManagerBuilder {
240    /// The CDN Access Token with from the Hive Console target.
241    pub fn key(mut self, key: String) -> Self {
242        self.key = non_empty_string(Some(key));
243        self
244    }
245
246    /// The CDN endpoint from Hive Console target.
247    pub fn add_endpoint(mut self, endpoint: String) -> Self {
248        if let Some(endpoint) = non_empty_string(Some(endpoint)) {
249            self.endpoints.push(endpoint);
250        }
251        self
252    }
253
254    /// Accept invalid SSL certificates
255    /// default: false
256    pub fn accept_invalid_certs(mut self, accept_invalid_certs: bool) -> Self {
257        self.accept_invalid_certs = accept_invalid_certs;
258        self
259    }
260
261    /// Connection timeout for the Hive Console CDN requests.
262    /// Default: 5 seconds
263    pub fn connect_timeout(mut self, connect_timeout: Duration) -> Self {
264        self.connect_timeout = connect_timeout;
265        self
266    }
267
268    /// Request timeout for the Hive Console CDN requests.
269    /// Default: 15 seconds
270    pub fn request_timeout(mut self, request_timeout: Duration) -> Self {
271        self.request_timeout = request_timeout;
272        self
273    }
274
275    /// Retry policy for fetching persisted documents
276    /// Default: ExponentialBackoff with max 3 retries
277    pub fn retry_policy(mut self, retry_policy: ExponentialBackoff) -> Self {
278        self.retry_policy = retry_policy;
279        self
280    }
281
282    /// Maximum number of retries for fetching persisted documents
283    /// Default: ExponentialBackoff with max 3 retries
284    pub fn max_retries(mut self, max_retries: u32) -> Self {
285        self.retry_policy = ExponentialBackoff::builder().build_with_max_retries(max_retries);
286        self
287    }
288
289    /// Size of the in-memory cache for persisted documents
290    /// Default: 10,000 entries
291    pub fn cache_size(mut self, cache_size: u64) -> Self {
292        self.cache_size = cache_size;
293        self
294    }
295
296    /// TTL for negative cache entries (failed lookups / not found responses).
297    ///
298    /// When set, repeated misses for the same document id are served from in-memory cache
299    /// until the TTL expires.
300    pub fn negative_cache_ttl(mut self, ttl: Duration) -> Self {
301        self.negative_cache_ttl = Some(ttl);
302        self
303    }
304
305    /// Circuit breaker configuration for persisted document CDN requests.
306    pub fn circuit_breaker(mut self, circuit_breaker: CircuitBreakerBuilder) -> Self {
307        self.circuit_breaker = circuit_breaker;
308        self
309    }
310
311    /// User-Agent header to be sent with each request
312    pub fn user_agent(mut self, user_agent: String) -> Self {
313        self.user_agent = non_empty_string(Some(user_agent));
314        self
315    }
316
317    pub fn build(self) -> Result<PersistedDocumentsManager, PersistedDocumentsError> {
318        let mut default_headers = HeaderMap::new();
319        let key = match self.key {
320            Some(key) => key,
321            None => {
322                return Err(PersistedDocumentsError::MissingConfigurationOption(
323                    "key".to_string(),
324                ));
325            }
326        };
327        default_headers.insert(
328            "X-Hive-CDN-Key",
329            HeaderValue::from_str(&key)
330                .map_err(|e| PersistedDocumentsError::InvalidCDNKey(e.to_string()))?,
331        );
332        let mut reqwest_agent = reqwest::Client::builder()
333            .danger_accept_invalid_certs(self.accept_invalid_certs)
334            .connect_timeout(self.connect_timeout)
335            .timeout(self.request_timeout)
336            .default_headers(default_headers);
337
338        if let Some(user_agent) = self.user_agent {
339            reqwest_agent = reqwest_agent.user_agent(user_agent);
340        }
341
342        let reqwest_agent = reqwest_agent
343            .build()
344            .map_err(|e| PersistedDocumentsError::HTTPClientCreationError(e.to_string()))?;
345        let client = ClientBuilder::new(reqwest_agent)
346            .with(RetryTransientMiddleware::new_with_policy(self.retry_policy))
347            .build();
348
349        let cache = Cache::<String, String>::new(self.cache_size);
350        let negative_cache = self.negative_cache_ttl.map(|ttl| {
351            Cache::builder()
352                .max_capacity(self.cache_size)
353                .time_to_live(ttl)
354                .build()
355        });
356
357        if self.endpoints.is_empty() {
358            return Err(PersistedDocumentsError::MissingConfigurationOption(
359                "endpoints".to_string(),
360            ));
361        }
362
363        Ok(PersistedDocumentsManager {
364            client,
365            cache,
366            negative_cache,
367            endpoints_with_circuit_breakers: self
368                .endpoints
369                .into_iter()
370                .map(move |endpoint| {
371                    let circuit_breaker = self.circuit_breaker.clone().build_async()?;
372                    Ok((endpoint, circuit_breaker))
373                })
374                .collect::<Result<Vec<(String, AsyncRecloser)>, PersistedDocumentsError>>()?,
375        })
376    }
377}