Read/write cache in evaluator.

Author: badmonster0

src/builder/analyzer.rs

@@ -8,6 +8,7 @@ use crate::setup::{
     self, DesiredMode, FlowSetupMetadata, FlowSetupState, ResourceIdentifier, SourceSetupState,
     TargetSetupState, TargetSetupStateCommon,
 };
+use crate::utils::fingerprint::Fingerprinter;
 use crate::{
     api_bail, api_error,
     base::{
@@ -698,24 +699,32 @@ impl<'a> AnalyzerContext<'a> {
                         let output = scope
                             .data
                             .add_field(reactive_op.name.clone(), &output_type)?;
-                        let op_name = reactive_op.name.clone();
+                        let reactive_op = reactive_op.clone();
                         async move {
                             let executor = executor.await.with_context(|| {
-                                format!("Failed to build executor for transform op: {op_name}")
+                                format!("Failed to build executor for transform op: {}", reactive_op.name)
                             })?;
+                            let behavior_version = executor.behavior_version();
                             let function_exec_info = AnalyzedFunctionExecInfo {
                                 enable_caching: executor.enable_caching(),
-                                behavior_version: executor.behavior_version(),
+                                behavior_version,
+                                fingerprinter: Fingerprinter::default()
+                                    .with(&reactive_op.name)?
+                                    .with(&reactive_op.spec)?
+                                    .with(&behavior_version)?
+                                    .with(&output_type.without_attrs())?,
+                                output_type: output_type.typ.clone(),
                             };
                             if function_exec_info.enable_caching
                                 && function_exec_info.behavior_version.is_some()
                             {
                                 api_bail!(
-                                    "When caching is enabled, behavior version must be specified for transform op: {op_name}",
+                                    "When caching is enabled, behavior version must be specified for transform op: {}",
+                                    reactive_op.name
                                 );
                             }
                             Ok(AnalyzedReactiveOp::Transform(AnalyzedTransformOp {
-                                name: op_name,
+                                name: reactive_op.name,
                                 inputs: input_value_mappings,
                                 function_exec_info,
                                 executor,

src/builder/plan.rs

@@ -6,6 +6,7 @@ use crate::base::schema::ValueType;
 use crate::base::value;
 use crate::execution::db_tracking_setup;
 use crate::ops::interface::*;
+use crate::utils::fingerprint::Fingerprinter;
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 pub struct AnalyzedLocalFieldReference {
@@ -65,6 +66,10 @@ pub struct AnalyzedSourceOp {
 pub struct AnalyzedFunctionExecInfo {
     pub enable_caching: bool,
     pub behavior_version: Option<u32>,
+
+    /// Fingerprinter of the function's behavior.
+    pub fingerprinter: Fingerprinter,
+    pub output_type: ValueType,
 }
 
 pub struct AnalyzedTransformOp {

src/execution/evaluator.rs

@@ -304,7 +304,29 @@ async fn evaluate_op_scope(
         match reactive_op {
             AnalyzedReactiveOp::Transform(op) => {
                 let input_values = assemble_input_values(&op.inputs, scoped_entries);
-                let output_value = op.executor.evaluate(input_values).await?;
+                let output_value = if let Some(cache) = op
+                    .function_exec_info
+                    .enable_caching
+                    .then_some(cache)
+                    .flatten()
+                {
+                    let key = op
+                        .function_exec_info
+                        .fingerprinter
+                        .clone()
+                        .with(&input_values)?
+                        .to_fingerprint();
+                    cache
+                        .evaluate(
+                            key,
+                            &op.function_exec_info.output_type,
+                            /*ttl=*/ None,
+                            move || op.executor.evaluate(input_values),
+                        )
+                        .await?
+                } else {
+                    op.executor.evaluate(input_values).await?
+                };
                 head_scope.define_field(&op.output, output_value)?;
             }
 

src/execution/memoization.rs

@@ -2,11 +2,13 @@ use anyhow::Result;
 use serde::{Deserialize, Serialize};
 use std::{
     collections::HashMap,
+    future::Future,
     sync::{Arc, Mutex},
 };
 
 use crate::{
     base::{schema, value},
+    service::error::{SharedError, SharedResultExt},
     utils::fingerprint::Fingerprint,
 };
 
@@ -37,7 +39,7 @@ enum EvaluationCacheData {
     /// Existing entry in previous runs, but not in current run yet.
     Previous(serde_json::Value),
     /// Value appeared in current run.
-    Current(Arc<async_lock::OnceCell<value::Value>>),
+    Current(Arc<async_lock::OnceCell<Result<value::Value, SharedError>>>),
 }
 
 pub struct EvaluationCache {
@@ -79,25 +81,28 @@ impl EvaluationCache {
             .into_iter()
             .filter_map(|(k, e)| match e.data {
                 EvaluationCacheData::Previous(_) => None,
-                EvaluationCacheData::Current(entry) => entry.get().map(|v| {
-                    Ok((
-                        k,
-                        CacheEntry {
-                            time_sec: e.time.timestamp(),
-                            value: serde_json::to_value(v)?,
-                        },
-                    ))
-                }),
+                EvaluationCacheData::Current(entry) => match entry.get() {
+                    Some(Ok(v)) => Some(serde_json::to_value(v).map(|value| {
+                        (
+                            k,
+                            CacheEntry {
+                                time_sec: e.time.timestamp(),
+                                value,
+                            },
+                        )
+                    })),
+                    _ => None,
+                },
             })
-            .collect::<Result<_>>()?)
+            .collect::<Result<_, _>>()?)
     }
 
     pub fn get(
         &self,
         key: Fingerprint,
         typ: &schema::ValueType,
         ttl: Option<chrono::Duration>,
-    ) -> Result<Arc<async_lock::OnceCell<value::Value>>> {
+    ) -> Result<Arc<async_lock::OnceCell<Result<value::Value, SharedError>>>> {
         let mut cache = self.cache.lock().unwrap();
         let result = {
             match cache.entry(key) {
@@ -110,7 +115,7 @@ impl EvaluationCache {
                     match &mut entry_mut.data {
                         EvaluationCacheData::Previous(value) => {
                             let value = value::Value::from_json(std::mem::take(value), typ)?;
-                            let cell = Arc::new(async_lock::OnceCell::from(value));
+                            let cell = Arc::new(async_lock::OnceCell::from(Ok(value)));
                             let time = entry_mut.time;
                             entry.insert(EvaluationCacheEntry {
                                 time,
@@ -133,4 +138,24 @@ impl EvaluationCache {
         };
         Ok(result)
     }
+
+    pub async fn evaluate<Fut>(
+        &self,
+        key: Fingerprint,
+        typ: &schema::ValueType,
+        ttl: Option<chrono::Duration>,
+        compute: impl FnOnce() -> Fut,
+    ) -> Result<value::Value>
+    where
+        Fut: Future<Output = Result<value::Value>>,
+    {
+        let cell = self.get(key, typ, ttl)?;
+        let result = cell
+            .get_or_init(|| {
+                let fut = compute();
+                async move { fut.await.map_err(SharedError::new) }
+            })
+            .await;
+        Ok(result.clone().std_result()?)
+    }
 }