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 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 pub fn key(mut self, key: String) -> Self {
242 self.key = non_empty_string(Some(key));
243 self
244 }
245
246 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 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 pub fn connect_timeout(mut self, connect_timeout: Duration) -> Self {
264 self.connect_timeout = connect_timeout;
265 self
266 }
267
268 pub fn request_timeout(mut self, request_timeout: Duration) -> Self {
271 self.request_timeout = request_timeout;
272 self
273 }
274
275 pub fn retry_policy(mut self, retry_policy: ExponentialBackoff) -> Self {
278 self.retry_policy = retry_policy;
279 self
280 }
281
282 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 pub fn cache_size(mut self, cache_size: u64) -> Self {
292 self.cache_size = cache_size;
293 self
294 }
295
296 pub fn negative_cache_ttl(mut self, ttl: Duration) -> Self {
301 self.negative_cache_ttl = Some(ttl);
302 self
303 }
304
305 pub fn circuit_breaker(mut self, circuit_breaker: CircuitBreakerBuilder) -> Self {
307 self.circuit_breaker = circuit_breaker;
308 self
309 }
310
311 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}