Evaluator starts to read/write cache if enabled (#12)

* Add a timestamp to cache entry.

* Move `Fingerprinter` to `utils`.

* Add `with()` and `write()` APIs to the `Fingerprinter` for convenience.

* Create a `Fingerprint` type and directly use it as cache key.

* 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/db_tracking.rs

@@ -1,10 +1,9 @@
 use super::{db_tracking_setup::TrackingTableSetupState, memoization::MemoizationInfo};
-use crate::utils::db::WriteAction;
+use crate::utils::{db::WriteAction, fingerprint::Fingerprint};
 use anyhow::Result;
 use sqlx::PgPool;
 
-pub type ValueFingerprint = String;
-pub type TrackedTargetKey = (serde_json::Value, i64, Option<ValueFingerprint>);
+pub type TrackedTargetKey = (serde_json::Value, i64, Option<Fingerprint>);
 pub type TrackedTargetKeyForSource = Vec<(i32, Vec<TrackedTargetKey>)>;
 
 #[derive(sqlx::FromRow, Debug)]

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/indexer.rs

@@ -6,18 +6,16 @@ use log::error;
 use serde::Serialize;
 use sqlx::PgPool;
 
-use super::db_tracking::{self, read_source_tracking_info};
+use super::db_tracking::{self, read_source_tracking_info, TrackedTargetKey};
 use super::db_tracking_setup;
-use super::fingerprint::Fingerprinter;
 use super::memoization::{EvaluationCache, MemoizationInfo};
 use crate::base::schema;
 use crate::base::spec::FlowInstanceSpec;
 use crate::base::value::{self, FieldValues, KeyValue};
 use crate::builder::plan::*;
 use crate::ops::interface::{ExportTargetMutation, ExportTargetUpsertEntry};
 use crate::utils::db::WriteAction;
-
-use self::db_tracking::{TrackedTargetKey, ValueFingerprint};
+use crate::utils::fingerprint::{Fingerprint, Fingerprinter};
 
 use super::evaluator::{evaluate_source_entry, ScopeValueBuilder};
 
@@ -96,8 +94,8 @@ struct TrackingInfoForTarget<'a> {
     // Existing keys info. Keyed by target key.
     // Will be removed after new rows for the same key are added into `new_staging_keys_info` and `mutation.upserts`,
     // hence all remaining ones are to be deleted.
-    existing_staging_keys_info: HashMap<serde_json::Value, Vec<(i64, Option<ValueFingerprint>)>>,
-    existing_keys_info: HashMap<serde_json::Value, Vec<(i64, Option<ValueFingerprint>)>>,
+    existing_staging_keys_info: HashMap<serde_json::Value, Vec<(i64, Option<Fingerprint>)>>,
+    existing_keys_info: HashMap<serde_json::Value, Vec<(i64, Option<Fingerprint>)>>,
 
     // New keys info for staging.
     new_staging_keys_info: Vec<TrackedTargetKey>,
@@ -215,9 +213,11 @@ async fn precommit_source_tracking_info(
                         .fields
                         .push(value.fields[*field as usize].clone());
                 }
-                let mut fingerprinter = Fingerprinter::default();
-                field_values.serialize(&mut fingerprinter)?;
-                let curr_fp = Some(fingerprinter.to_base64());
+                let curr_fp = Some(
+                    Fingerprinter::default()
+                        .with(&field_values)?
+                        .to_fingerprint(),
+                );
 
                 let existing_target_keys = target_info.existing_keys_info.remove(&primary_key_json);
                 let existing_staging_target_keys = target_info
@@ -439,9 +439,8 @@ pub async fn update_source_entry<'a>(
         .map(|info| info.memoization_info.map(|info| info.0))
         .flatten()
         .flatten();
-    let evaluation_cache = memoization_info
-        .map(|info| EvaluationCache::from_stored(info.cache))
-        .unwrap_or_default();
+    let evaluation_cache =
+        EvaluationCache::new(process_timestamp, memoization_info.map(|info| info.cache));
     let value_builder =
         evaluate_source_entry(plan, source_op_idx, schema, key, Some(&evaluation_cache)).await?;
 

src/execution/memoization.rs

@@ -1,42 +1,25 @@
 use anyhow::Result;
 use serde::{Deserialize, Serialize};
-
-use base64::prelude::*;
 use std::{
     collections::HashMap,
+    future::Future,
     sync::{Arc, Mutex},
 };
 
-use crate::base::{schema, value};
-
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
-pub struct CacheKey(Vec<u8>);
-
-impl Serialize for CacheKey {
-    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-    where
-        S: serde::Serializer,
-    {
-        serializer.serialize_str(&BASE64_STANDARD.encode(&self.0))
-    }
-}
+use crate::{
+    base::{schema, value},
+    service::error::{SharedError, SharedResultExt},
+    utils::fingerprint::Fingerprint,
+};
 
-impl<'de> Deserialize<'de> for CacheKey {
-    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
-    where
-        D: serde::Deserializer<'de>,
-    {
-        let s = String::deserialize(deserializer)?;
-        let bytes = BASE64_STANDARD
-            .decode(s)
-            .map_err(serde::de::Error::custom)?;
-        Ok(CacheKey(bytes))
-    }
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CacheEntry {
+    time_sec: i64,
+    value: serde_json::Value,
 }
-
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct MemoizationInfo {
-    pub cache: HashMap<CacheKey, serde_json::Value>,
+    pub cache: HashMap<Fingerprint, CacheEntry>,
 }
 
 impl Default for MemoizationInfo {
@@ -47,66 +30,132 @@ impl Default for MemoizationInfo {
     }
 }
 
-enum EvaluationCacheEntry {
+struct EvaluationCacheEntry {
+    time: chrono::DateTime<chrono::Utc>,
+    data: EvaluationCacheData,
+}
+
+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>>>),
 }
 
-#[derive(Default)]
 pub struct EvaluationCache {
-    cache: Mutex<HashMap<CacheKey, EvaluationCacheEntry>>,
+    current_time: chrono::DateTime<chrono::Utc>,
+    cache: Mutex<HashMap<Fingerprint, EvaluationCacheEntry>>,
 }
 
 impl EvaluationCache {
-    pub fn from_stored(cache: HashMap<CacheKey, serde_json::Value>) -> Self {
+    pub fn new(
+        current_time: chrono::DateTime<chrono::Utc>,
+        existing_cache: Option<HashMap<Fingerprint, CacheEntry>>,
+    ) -> Self {
         Self {
+            current_time,
             cache: Mutex::new(
-                cache
+                existing_cache
                     .into_iter()
-                    .map(|(k, v)| (k, EvaluationCacheEntry::Previous(v)))
+                    .map(|e| e.into_iter())
+                    .flatten()
+                    .map(|(k, e)| {
+                        (
+                            k,
+                            EvaluationCacheEntry {
+                                time: chrono::DateTime::from_timestamp(e.time_sec, 0)
+                                    .unwrap_or(chrono::DateTime::<chrono::Utc>::MIN_UTC),
+                                data: EvaluationCacheData::Previous(e.value),
+                            },
+                        )
+                    })
                     .collect(),
             ),
         }
     }
 
-    pub fn into_stored(self) -> Result<HashMap<CacheKey, serde_json::Value>> {
+    pub fn into_stored(self) -> Result<HashMap<Fingerprint, CacheEntry>> {
         Ok(self
             .cache
             .into_inner()?
             .into_iter()
-            .filter_map(|(k, v)| match v {
-                EvaluationCacheEntry::Previous(_) => None,
-                EvaluationCacheEntry::Current(entry) => {
-                    entry.get().map(|v| Ok((k, serde_json::to_value(v)?)))
-                }
+            .filter_map(|(k, e)| match e.data {
+                EvaluationCacheData::Previous(_) => None,
+                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: CacheKey,
+        key: Fingerprint,
         typ: &schema::ValueType,
-    ) -> Result<Arc<async_lock::OnceCell<value::Value>>> {
+        ttl: Option<chrono::Duration>,
+    ) -> Result<Arc<async_lock::OnceCell<Result<value::Value, SharedError>>>> {
         let mut cache = self.cache.lock().unwrap();
-        let result = match cache.entry(key) {
-            std::collections::hash_map::Entry::Occupied(mut entry) => match &mut entry.get_mut() {
-                EvaluationCacheEntry::Previous(value) => {
-                    let value = value::Value::from_json(std::mem::take(value), typ)?;
-                    let cell = Arc::new(async_lock::OnceCell::from(value));
-                    entry.insert(EvaluationCacheEntry::Current(cell.clone()));
+        let result = {
+            match cache.entry(key) {
+                std::collections::hash_map::Entry::Occupied(mut entry)
+                    if !ttl
+                        .map(|ttl| entry.get().time + ttl < self.current_time)
+                        .unwrap_or(false) =>
+                {
+                    let entry_mut = &mut entry.get_mut();
+                    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(Ok(value)));
+                            let time = entry_mut.time;
+                            entry.insert(EvaluationCacheEntry {
+                                time,
+                                data: EvaluationCacheData::Current(cell.clone()),
+                            });
+                            cell
+                        }
+                        EvaluationCacheData::Current(cell) => cell.clone(),
+                    }
+                }
+                entry => {
+                    let cell = Arc::new(async_lock::OnceCell::new());
+                    entry.insert_entry(EvaluationCacheEntry {
+                        time: self.current_time,
+                        data: EvaluationCacheData::Current(cell.clone()),
+                    });
                     cell
                 }
-                EvaluationCacheEntry::Current(cell) => cell.clone(),
-            },
-            std::collections::hash_map::Entry::Vacant(entry) => {
-                let cell = Arc::new(async_lock::OnceCell::new());
-                entry.insert(EvaluationCacheEntry::Current(cell.clone()));
-                cell
             }
         };
         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()?)
+    }
 }

src/execution/mod.rs

@@ -5,5 +5,4 @@ pub mod query;
 mod db_tracking;
 pub mod db_tracking_setup;
 
-mod fingerprint;
 mod memoization;