Automatically retry on transaction failure for Neo4j.

Author: badmonster0

Cargo.toml

@@ -102,3 +102,4 @@ tokio-stream = "0.1.17"
 async-stream = "0.3.6"
 neo4rs = "0.8.0"
 bytes = "1.10.1"
+rand = "0.9.0"

src/base/value.rs

@@ -533,6 +533,23 @@ impl From<KeyValue> for Value {
     }
 }
 
+impl From<&KeyValue> for Value {
+    fn from(value: &KeyValue) -> Self {
+        match value {
+            KeyValue::Bytes(v) => Value::Basic(BasicValue::Bytes(v.clone())),
+            KeyValue::Str(v) => Value::Basic(BasicValue::Str(v.clone())),
+            KeyValue::Bool(v) => Value::Basic(BasicValue::Bool(*v)),
+            KeyValue::Int64(v) => Value::Basic(BasicValue::Int64(*v)),
+            KeyValue::Range(v) => Value::Basic(BasicValue::Range(*v)),
+            KeyValue::Uuid(v) => Value::Basic(BasicValue::Uuid(*v)),
+            KeyValue::Date(v) => Value::Basic(BasicValue::Date(*v)),
+            KeyValue::Struct(v) => Value::Struct(FieldValues {
+                fields: v.iter().map(Value::from).collect(),
+            }),
+        }
+    }
+}
+
 impl From<FieldValues> for Value {
     fn from(value: FieldValues) -> Self {
         Value::Struct(value)

src/ops/storages/neo4j.rs

@@ -92,6 +92,16 @@ impl GraphRelationship {
     }
 }
 
+impl retriable::IsRetryable for neo4rs::Error {
+    fn is_retryable(&self) -> bool {
+        match self {
+            neo4rs::Error::ConnectionError => true,
+            neo4rs::Error::Neo4j(e) => e.kind() == neo4rs::Neo4jErrorKind::Transient,
+            _ => false,
+        }
+    }
+}
+
 #[derive(Default)]
 pub struct GraphPool {
     graphs: Mutex<HashMap<GraphKey, Arc<OnceCell<Arc<Graph>>>>>,
@@ -172,7 +182,7 @@ fn json_value_to_bolt_value(value: &serde_json::Value) -> Result<BoltType> {
     Ok(bolt_value)
 }
 
-fn key_to_bolt(key: KeyValue, schema: &schema::ValueType) -> Result<BoltType> {
+fn key_to_bolt(key: &KeyValue, schema: &schema::ValueType) -> Result<BoltType> {
     value_to_bolt(&key.into(), schema)
 }
 
@@ -362,18 +372,18 @@ FINISH
             tgt_fields,
         })
     }
-}
 
-#[async_trait]
-impl ExportTargetExecutor for RelationshipStorageExecutor {
-    async fn apply_mutation(&self, mutation: ExportTargetMutation) -> Result<()> {
+    fn build_queries_to_apply_mutation(
+        &self,
+        mutation: &ExportTargetMutation,
+    ) -> Result<Vec<neo4rs::Query>> {
         let mut queries = vec![];
-        for upsert in mutation.upserts {
-            let rel_id_bolt = key_to_bolt(upsert.key, &self.key_field.value_type.typ)?;
+        for upsert in mutation.upserts.iter() {
+            let rel_id_bolt = key_to_bolt(&upsert.key, &self.key_field.value_type.typ)?;
             queries
                 .push(neo4rs::query(&self.delete_cypher).param(REL_ID_PARAM, rel_id_bolt.clone()));
 
-            let value = upsert.value;
+            let value = &upsert.value;
             let mut insert_cypher = neo4rs::query(&self.insert_cypher)
                 .param(REL_ID_PARAM, rel_id_bolt)
                 .param(
@@ -425,17 +435,31 @@ impl ExportTargetExecutor for RelationshipStorageExecutor {
             }
             queries.push(insert_cypher);
         }
-        for delete_key in mutation.delete_keys {
+        for delete_key in mutation.delete_keys.iter() {
             queries.push(neo4rs::query(&self.delete_cypher).param(
                 REL_ID_PARAM,
                 key_to_bolt(delete_key, &self.key_field.value_type.typ)?,
             ));
         }
+        Ok(queries)
+    }
+}
 
-        let mut txn = self.graph.start_txn().await?;
-        txn.run_queries(queries).await?;
-        txn.commit().await?;
-        Ok(())
+#[async_trait]
+impl ExportTargetExecutor for RelationshipStorageExecutor {
+    async fn apply_mutation(&self, mutation: ExportTargetMutation) -> Result<()> {
+        retriable::run(
+            || async {
+                let queries = self.build_queries_to_apply_mutation(&mutation)?;
+                let mut txn = self.graph.start_txn().await?;
+                txn.run_queries(queries.clone()).await?;
+                txn.commit().await?;
+                retriable::Ok(())
+            },
+            retriable::RunOptions::default(),
+        )
+        .await
+        .map_err(Into::<anyhow::Error>::into)
     }
 }
 

src/prelude.rs

@@ -24,6 +24,7 @@ pub(crate) use crate::lib_context::{get_lib_context, get_runtime, FlowContext, L
 pub(crate) use crate::ops::interface;
 pub(crate) use crate::service::error::ApiError;
 pub(crate) use crate::setup::AuthRegistry;
+pub(crate) use crate::utils::retriable;
 
 pub(crate) use crate::{api_bail, api_error};
 

src/utils/mod.rs

@@ -1,4 +1,5 @@
 pub mod db;
 pub mod fingerprint;
 pub mod immutable;
+pub mod retriable;
 pub mod yaml_ser;

src/utils/retriable.rs

@@ -0,0 +1,116 @@
+use log::trace;
+use std::{future::Future, time::Duration};
+
+pub trait IsRetryable {
+    fn is_retryable(&self) -> bool;
+}
+
+pub struct Error {
+    error: anyhow::Error,
+    is_retryable: bool,
+}
+
+impl std::fmt::Display for Error {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        std::fmt::Display::fmt(&self.error, f)
+    }
+}
+
+impl std::fmt::Debug for Error {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        std::fmt::Debug::fmt(&self.error, f)
+    }
+}
+
+impl IsRetryable for Error {
+    fn is_retryable(&self) -> bool {
+        self.is_retryable
+    }
+}
+
+impl From<anyhow::Error> for Error {
+    fn from(error: anyhow::Error) -> Self {
+        Self {
+            error,
+            is_retryable: false,
+        }
+    }
+}
+
+impl Into<anyhow::Error> for Error {
+    fn into(self) -> anyhow::Error {
+        self.error
+    }
+}
+
+impl<E: IsRetryable + std::error::Error + Send + Sync + 'static> From<E> for Error {
+    fn from(error: E) -> Self {
+        Self {
+            is_retryable: error.is_retryable(),
+            error: anyhow::Error::new(error),
+        }
+    }
+}
+
+pub type Result<T, E = Error> = std::result::Result<T, E>;
+
+#[allow(non_snake_case)]
+pub fn Ok<T>(value: T) -> Result<T> {
+    Result::Ok(value)
+}
+
+pub struct RunOptions {
+    pub max_retries: usize,
+    pub initial_backoff: Duration,
+    pub max_backoff: Duration,
+}
+
+impl Default for RunOptions {
+    fn default() -> Self {
+        Self {
+            max_retries: 5,
+            initial_backoff: Duration::from_millis(100),
+            max_backoff: Duration::from_secs(10),
+        }
+    }
+}
+
+pub async fn run<
+    Ok,
+    Err: std::fmt::Display + IsRetryable,
+    Fut: Future<Output = Result<Ok, Err>>,
+    F: Fn() -> Fut,
+>(
+    f: F,
+    options: RunOptions,
+) -> Result<Ok, Err> {
+    let mut retries = 0;
+    let mut backoff = options.initial_backoff;
+
+    loop {
+        match f().await {
+            Result::Ok(result) => return Result::Ok(result),
+            Result::Err(err) => {
+                if !err.is_retryable() || retries >= options.max_retries {
+                    return Result::Err(err);
+                }
+                retries += 1;
+                trace!(
+                    "Will retry #{} in {}ms for error: {}",
+                    retries,
+                    backoff.as_millis(),
+                    err
+                );
+                tokio::time::sleep(backoff).await;
+                if backoff < options.max_backoff {
+                    backoff = std::cmp::min(
+                        Duration::from_micros(
+                            (backoff.as_micros() * rand::random_range(1618..=2000) / 1000) as u64,
+                        ),
+                        options.max_backoff,
+                    );
+                }
+            }
+        }
+    }
+}