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