Skip to main content

arete_interpreter/
runtime_resolvers.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::OnceLock;
5use std::time::Instant;
6
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9
10use crate::ast::{HttpMethod, ResolverType};
11use crate::compiler::MultiEntityBytecode;
12use crate::resolvers::{TokenMetadataResolverClient, UrlResolverClient};
13use crate::vm::{
14    CachedResolverValue, ResolverRequest, TokenResolverMissAction, UpdateContext, VmContext,
15};
16use crate::Mutation;
17
18pub type ResolverBatchResult =
19    Result<HashMap<String, Value>, Box<dyn std::error::Error + Send + Sync>>;
20pub type ResolverBatchFuture<'a> = Pin<Box<dyn Future<Output = ResolverBatchResult> + Send + 'a>>;
21pub type ResolverApplyFuture<'a> = Pin<Box<dyn Future<Output = Vec<Mutation>> + Send + 'a>>;
22pub type SharedRuntimeResolver = std::sync::Arc<dyn RuntimeResolver>;
23
24#[derive(Clone, Copy)]
25struct RuntimeResolverProfileConfig {
26    enabled: bool,
27    slow_ms: f64,
28}
29
30fn runtime_resolver_profile_config() -> RuntimeResolverProfileConfig {
31    static CONFIG: OnceLock<RuntimeResolverProfileConfig> = OnceLock::new();
32    *CONFIG.get_or_init(|| RuntimeResolverProfileConfig {
33        enabled: std::env::var("ARETE_PROFILE_CODEPATHS")
34            .ok()
35            .or_else(|| std::env::var("ARETE_PROFILE_RUNTIME_RESOLVER").ok())
36            .map(|value| {
37                matches!(
38                    value.to_ascii_lowercase().as_str(),
39                    "1" | "true" | "yes" | "on"
40                )
41            })
42            .unwrap_or(false),
43        slow_ms: std::env::var("ARETE_PROFILE_CODEPATHS_SLOW_MS")
44            .ok()
45            .or_else(|| std::env::var("ARETE_PROFILE_RUNTIME_RESOLVER_SLOW_MS").ok())
46            .and_then(|value| value.parse::<f64>().ok())
47            .unwrap_or(5.0),
48    })
49}
50
51fn elapsed_ms(start: Instant) -> f64 {
52    start.elapsed().as_secs_f64() * 1000.0
53}
54
55fn should_log_runtime_resolver_profile(total_ms: f64) -> bool {
56    let config = runtime_resolver_profile_config();
57    config.enabled && total_ms >= config.slow_ms
58}
59
60#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
61#[serde(tag = "type", rename_all = "snake_case")]
62pub enum RuntimeResolverRequest {
63    TokenMetadata {
64        key: String,
65        mint: String,
66    },
67    UrlJson {
68        key: String,
69        url: String,
70        method: HttpMethod,
71    },
72}
73
74impl RuntimeResolverRequest {
75    pub fn key(&self) -> &str {
76        match self {
77            Self::TokenMetadata { key, .. } | Self::UrlJson { key, .. } => key,
78        }
79    }
80}
81
82#[derive(Debug, Clone, Serialize, Deserialize)]
83pub struct RuntimeResolverBatchRequest {
84    pub requests: Vec<RuntimeResolverRequest>,
85}
86
87#[derive(Debug, Clone, Serialize, Deserialize)]
88pub struct RuntimeResolverResponse {
89    pub key: String,
90    pub value: Value,
91}
92
93#[derive(Debug, Clone, Serialize, Deserialize)]
94pub struct RuntimeResolverBatchResponse {
95    #[serde(default)]
96    pub resolved: Vec<RuntimeResolverResponse>,
97}
98
99#[derive(Debug)]
100struct PendingRuntimeResolverRequest {
101    request: ResolverRequest,
102    backend_request: RuntimeResolverRequest,
103}
104
105pub trait RuntimeResolver: Send + Sync {
106    fn resolve_batch<'a>(
107        &'a self,
108        requests: &'a [RuntimeResolverRequest],
109    ) -> ResolverBatchFuture<'a>;
110
111    fn resolve_and_apply<'a>(
112        &'a self,
113        vm: &'a std::sync::Mutex<VmContext>,
114        bytecode: &'a MultiEntityBytecode,
115        requests: Vec<ResolverRequest>,
116        apply_context: Option<UpdateContext>,
117    ) -> ResolverApplyFuture<'a> {
118        Box::pin(async move {
119            let total_started_at = Instant::now();
120            if requests.is_empty() {
121                return Vec::new();
122            }
123
124            let request_count = requests.len();
125            let mut cached = Vec::new();
126            let mut pending = Vec::new();
127            let mut invalid = Vec::new();
128            let mut negative_cached = Vec::new();
129
130            let classify_started_at = Instant::now();
131            let (
132                classify_ms,
133                cached_count,
134                pending_count,
135                invalid_count,
136                negative_cached_count,
137                unique_request_count,
138            ) = {
139                let mut vm_guard = vm.lock().unwrap_or_else(|e| e.into_inner());
140
141                for request in requests {
142                    let canonical_key =
143                        runtime_resolver_cache_key(&request.resolver, &request.input);
144
145                    if let Some(cached_value) = vm_guard.get_cached_resolver_value(&canonical_key) {
146                        match cached_value {
147                            CachedResolverValue::Resolved(resolved_value) => {
148                                cached.push((request, resolved_value));
149                            }
150                            CachedResolverValue::Negative => {
151                                negative_cached.push(request);
152                            }
153                        }
154                        continue;
155                    }
156
157                    match runtime_request_from_vm_request(&request) {
158                        Some(backend_request) => pending.push(PendingRuntimeResolverRequest {
159                            request,
160                            backend_request,
161                        }),
162                        None => invalid.push(request),
163                    }
164                }
165
166                let invalid_count = invalid.len();
167                if !invalid.is_empty() {
168                    vm_guard.restore_resolver_requests(invalid);
169                }
170
171                let classify_ms = elapsed_ms(classify_started_at);
172                let cached_count = cached.len();
173                let pending_count = pending.len();
174                let negative_cached_count = negative_cached.len();
175                let unique_request_count = if pending.is_empty() {
176                    0
177                } else {
178                    let mut unique = HashMap::new();
179                    for entry in &pending {
180                        unique
181                            .entry(entry.backend_request.key().to_string())
182                            .or_insert_with(|| entry.backend_request.clone());
183                    }
184                    unique.len()
185                };
186
187                (
188                    classify_ms,
189                    cached_count,
190                    pending_count,
191                    invalid_count,
192                    negative_cached_count,
193                    unique_request_count,
194                )
195            };
196
197            let resolve_batch_started_at = Instant::now();
198            let resolved_map = if pending.is_empty() {
199                Ok(HashMap::new())
200            } else {
201                let mut unique = HashMap::new();
202                for entry in &pending {
203                    unique
204                        .entry(entry.backend_request.key().to_string())
205                        .or_insert_with(|| entry.backend_request.clone());
206                }
207
208                let unique_requests: Vec<RuntimeResolverRequest> = unique.into_values().collect();
209                self.resolve_batch(&unique_requests).await
210            };
211            let resolve_batch_ms = elapsed_ms(resolve_batch_started_at);
212
213            let mut mutations = Vec::new();
214            let mut failed = Vec::new();
215            let mut retry_scheduled_count = 0usize;
216            let mut negative_cached_miss_count = 0usize;
217
218            let apply_cached_started_at = Instant::now();
219            let mut vm_guard = vm.lock().unwrap_or_else(|e| e.into_inner());
220            let previous_context = if apply_context.is_some() {
221                let previous = vm_guard.current_context().cloned();
222                vm_guard.set_current_context(apply_context.clone());
223                previous
224            } else {
225                None
226            };
227
228            for (request, resolved_value) in cached {
229                match vm_guard.apply_resolver_result(bytecode, &request.cache_key, resolved_value) {
230                    Ok(mut new_mutations) => mutations.append(&mut new_mutations),
231                    Err(err) => {
232                        tracing::warn!(
233                            cache_key = %request.cache_key,
234                            error = %err,
235                            "Failed to apply cached resolver result"
236                        );
237                        failed.push(request);
238                    }
239                }
240            }
241            let apply_cached_ms = elapsed_ms(apply_cached_started_at);
242
243            let apply_pending_started_at = Instant::now();
244            match resolved_map {
245                Ok(resolved_map) => {
246                    for entry in pending {
247                        match resolved_map.get(entry.backend_request.key()) {
248                            Some(resolved_value) => match vm_guard.apply_resolver_result(
249                                bytecode,
250                                &entry.request.cache_key,
251                                resolved_value.clone(),
252                            ) {
253                                Ok(mut new_mutations) => mutations.append(&mut new_mutations),
254                                Err(err) => {
255                                    tracing::warn!(
256                                        cache_key = %entry.request.cache_key,
257                                        error = %err,
258                                        "Failed to apply resolver result"
259                                    );
260                                    failed.push(entry.request);
261                                }
262                            },
263                            None => match entry.backend_request {
264                                RuntimeResolverRequest::TokenMetadata { .. } => {
265                                    match vm_guard
266                                        .handle_token_resolver_miss(&entry.request.cache_key)
267                                    {
268                                        TokenResolverMissAction::RetryScheduled => {
269                                            retry_scheduled_count += 1;
270                                        }
271                                        TokenResolverMissAction::NegativeCached => {
272                                            negative_cached_miss_count += 1;
273                                        }
274                                        TokenResolverMissAction::Dropped => {}
275                                    }
276                                }
277                                RuntimeResolverRequest::UrlJson { .. } => {
278                                    failed.push(entry.request)
279                                }
280                            },
281                        }
282                    }
283                }
284                Err(err) => {
285                    tracing::warn!(error = %err, "Runtime resolver backend request failed");
286                    for entry in pending {
287                        match entry.backend_request {
288                            RuntimeResolverRequest::TokenMetadata { .. } => {
289                                match vm_guard.handle_token_resolver_miss(&entry.request.cache_key)
290                                {
291                                    TokenResolverMissAction::RetryScheduled => {
292                                        retry_scheduled_count += 1;
293                                    }
294                                    TokenResolverMissAction::NegativeCached => {
295                                        negative_cached_miss_count += 1;
296                                    }
297                                    TokenResolverMissAction::Dropped => {}
298                                }
299                            }
300                            RuntimeResolverRequest::UrlJson { .. } => failed.push(entry.request),
301                        }
302                    }
303                }
304            }
305            let apply_pending_ms = elapsed_ms(apply_pending_started_at);
306
307            let failed_count = failed.len();
308            let restore_failed_started_at = Instant::now();
309            if !failed.is_empty() {
310                vm_guard.restore_resolver_requests(failed);
311            }
312            let restore_failed_ms = elapsed_ms(restore_failed_started_at);
313            if apply_context.is_some() {
314                vm_guard.set_current_context(previous_context);
315            }
316
317            let total_ms = elapsed_ms(total_started_at);
318            if should_log_runtime_resolver_profile(total_ms) {
319                tracing::info!(
320                    stage = "runtime.resolver_apply",
321                    request_count,
322                    cached_count,
323                    pending_count,
324                    invalid_count,
325                    negative_cached_count,
326                    unique_request_count,
327                    failed_count,
328                    retry_scheduled_count,
329                    negative_cached_miss_count,
330                    mutation_count = mutations.len(),
331                    total_ms,
332                    classify_ms,
333                    resolve_batch_ms,
334                    apply_cached_ms,
335                    apply_pending_ms,
336                    restore_failed_ms,
337                    "Runtime resolver profile"
338                );
339            }
340
341            mutations
342        })
343    }
344}
345
346pub struct InProcessResolver {
347    token_client: Option<TokenMetadataResolverClient>,
348    url_client: UrlResolverClient,
349}
350
351impl InProcessResolver {
352    pub fn from_env() -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
353        Ok(Self {
354            token_client: TokenMetadataResolverClient::from_env()?,
355            url_client: UrlResolverClient::new(),
356        })
357    }
358
359    pub fn new(
360        token_client: Option<TokenMetadataResolverClient>,
361        url_client: UrlResolverClient,
362    ) -> Self {
363        Self {
364            token_client,
365            url_client,
366        }
367    }
368
369    pub async fn resolve_batch_internal(
370        &self,
371        requests: &[RuntimeResolverRequest],
372    ) -> ResolverBatchResult {
373        let total_started_at = Instant::now();
374        let mut results = HashMap::new();
375        let mut token_requests = Vec::new();
376        let mut url_requests = Vec::new();
377
378        let classify_started_at = Instant::now();
379        for request in requests {
380            match request {
381                RuntimeResolverRequest::TokenMetadata { key, mint } => {
382                    token_requests.push((key.clone(), mint.clone()));
383                }
384                RuntimeResolverRequest::UrlJson { key, url, method } => {
385                    url_requests.push((key.clone(), url.clone(), method.clone()));
386                }
387            }
388        }
389        let classify_ms = elapsed_ms(classify_started_at);
390
391        let token_request_count = token_requests.len();
392        let url_request_count = url_requests.len();
393        let mut token_batch_ms = 0.0_f64;
394        let mut url_dedupe_ms = 0.0_f64;
395        let mut url_batch_ms = 0.0_f64;
396        let mut url_unique_request_count = 0usize;
397
398        if !token_requests.is_empty() {
399            let token_started_at = Instant::now();
400            if let Some(token_client) = &self.token_client {
401                let mints: Vec<String> = token_requests
402                    .iter()
403                    .map(|(_, mint)| mint.clone())
404                    .collect();
405                match token_client.resolve_token_metadata(&mints).await {
406                    Ok(resolved) => {
407                        for (key, mint) in token_requests {
408                            if let Some(value) = resolved.get(&mint) {
409                                results.insert(key, value.clone());
410                            }
411                        }
412                    }
413                    Err(err) => {
414                        tracing::warn!(error = %err, "Failed to resolve token metadata batch");
415                    }
416                }
417            } else {
418                tracing::warn!(
419                    count = token_requests.len(),
420                    "DAS_API_ENDPOINT not set; token resolver requests will be re-queued"
421                );
422            }
423            token_batch_ms = elapsed_ms(token_started_at);
424        }
425
426        if !url_requests.is_empty() {
427            let url_dedupe_started_at = Instant::now();
428            let mut unique = HashMap::new();
429            for (key, url, method) in &url_requests {
430                unique
431                    .entry((url.clone(), method.clone()))
432                    .or_insert_with(Vec::new)
433                    .push(key.clone());
434            }
435            url_unique_request_count = unique.len();
436
437            let batch_input: Vec<(String, HttpMethod)> = unique.keys().cloned().collect();
438            url_dedupe_ms = elapsed_ms(url_dedupe_started_at);
439            let url_batch_started_at = Instant::now();
440            let resolved = self.url_client.resolve_batch(&batch_input).await;
441            url_batch_ms = elapsed_ms(url_batch_started_at);
442
443            for ((url, method), keys) in unique {
444                if let Some(value) = resolved.get(&(url, method)) {
445                    for key in keys {
446                        results.insert(key, value.clone());
447                    }
448                }
449            }
450        }
451
452        let total_ms = elapsed_ms(total_started_at);
453        if should_log_runtime_resolver_profile(total_ms) {
454            tracing::info!(
455                stage = "runtime.resolver_batch",
456                request_count = requests.len(),
457                token_request_count,
458                url_request_count,
459                url_unique_request_count,
460                resolved_count = results.len(),
461                total_ms,
462                classify_ms,
463                token_batch_ms,
464                url_dedupe_ms,
465                url_batch_ms,
466                "Runtime resolver batch profile"
467            );
468        }
469
470        Ok(results)
471    }
472}
473
474impl RuntimeResolver for InProcessResolver {
475    fn resolve_batch<'a>(
476        &'a self,
477        requests: &'a [RuntimeResolverRequest],
478    ) -> ResolverBatchFuture<'a> {
479        Box::pin(async move { self.resolve_batch_internal(requests).await })
480    }
481}
482
483pub fn runtime_resolver_cache_key(resolver: &ResolverType, input: &Value) -> String {
484    crate::vm::resolver_cache_key(resolver, input)
485}
486
487fn runtime_request_from_vm_request(request: &ResolverRequest) -> Option<RuntimeResolverRequest> {
488    match &request.resolver {
489        ResolverType::Token => extract_mint_from_input(&request.input).map(|mint| {
490            RuntimeResolverRequest::TokenMetadata {
491                key: request.cache_key.clone(),
492                mint,
493            }
494        }),
495        ResolverType::Url(config) => match &request.input {
496            Value::String(url) if !url.is_empty() => Some(RuntimeResolverRequest::UrlJson {
497                key: request.cache_key.clone(),
498                url: url.clone(),
499                method: config.method.clone(),
500            }),
501            _ => None,
502        },
503    }
504}
505
506fn extract_mint_from_input(input: &Value) -> Option<String> {
507    match input {
508        Value::String(value) if !value.is_empty() => Some(value.clone()),
509        Value::Object(map) => map
510            .get("mint")
511            .and_then(|value| value.as_str())
512            .filter(|value| !value.is_empty())
513            .map(str::to_string),
514        _ => None,
515    }
516}
517
518#[cfg(test)]
519mod tests {
520    use super::*;
521
522    #[test]
523    fn token_request_extracts_mint_from_object_input() {
524        let request = ResolverRequest {
525            cache_key: "token:mint".to_string(),
526            resolver: ResolverType::Token,
527            input: serde_json::json!({ "mint": "abc" }),
528        };
529
530        let runtime_request = runtime_request_from_vm_request(&request).unwrap();
531        assert_eq!(
532            runtime_request,
533            RuntimeResolverRequest::TokenMetadata {
534                key: "token:mint".to_string(),
535                mint: "abc".to_string(),
536            }
537        );
538    }
539
540    #[test]
541    fn url_request_uses_existing_cache_key() {
542        let request = ResolverRequest {
543            cache_key: "url:get:https://example.com".to_string(),
544            resolver: ResolverType::Url(crate::ast::UrlResolverConfig {
545                url_source: crate::ast::UrlSource::FieldPath("metadata_url".to_string()),
546                method: HttpMethod::Get,
547                extract_path: None,
548            }),
549            input: serde_json::json!("https://example.com"),
550        };
551
552        let runtime_request = runtime_request_from_vm_request(&request).unwrap();
553        assert_eq!(
554            runtime_request,
555            RuntimeResolverRequest::UrlJson {
556                key: "url:get:https://example.com".to_string(),
557                url: "https://example.com".to_string(),
558                method: HttpMethod::Get,
559            }
560        );
561    }
562}