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 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 pub fn key(mut self, key: String) -> Self {
251 self.key = non_empty_string(Some(key));
252 self
253 }
254
255 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 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 pub fn connect_timeout(mut self, connect_timeout: Duration) -> Self {
273 self.connect_timeout = connect_timeout;
274 self
275 }
276
277 pub fn request_timeout(mut self, request_timeout: Duration) -> Self {
280 self.request_timeout = request_timeout;
281 self
282 }
283
284 pub fn retry_policy(mut self, retry_policy: ExponentialBackoff) -> Self {
287 self.retry_policy = retry_policy;
288 self
289 }
290
291 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 pub fn cache_size(mut self, cache_size: u64) -> Self {
301 self.cache_size = cache_size;
302 self
303 }
304
305 pub fn negative_cache_ttl(mut self, ttl: Duration) -> Self {
310 self.negative_cache_ttl = Some(ttl);
311 self
312 }
313
314 pub fn circuit_breaker(mut self, circuit_breaker: CircuitBreakerBuilder) -> Self {
316 self.circuit_breaker = circuit_breaker;
317 self
318 }
319
320 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}