feat(query): add async retry transform

Author: sundy-li

Cargo.lock

@@ -35,6 +35,8 @@ use crate::DataSchema;
 const UDF_TCP_KEEP_ALIVE_SEC: u64 = 30;
 const UDF_HTTP2_KEEP_ALIVE_INTERVAL_SEC: u64 = 60;
 const UDF_KEEP_ALIVE_TIMEOUT_SEC: u64 = 20;
+// 4MB by default, we use 16G
+// max_encoding_message_size is usize::max by default
 const MAX_DECODING_MESSAGE_SIZE: usize = 16 * 1024 * 1024 * 1024;
 
 #[derive(Debug, Clone)]

src/query/expression/src/utils/udf_client.rs

@@ -23,6 +23,8 @@ jsonb = { workspace = true }
 match-template = { workspace = true }
 serde = { workspace = true }
 typetag = { workspace = true }
+tokio = { workspace = true }
+
 
 [dev-dependencies]
 itertools = { workspace = true }

src/query/pipeline/transforms/Cargo.toml

@@ -25,6 +25,7 @@ mod transform_dummy;
 mod transform_multi_sort_merge;
 mod transform_sort_merge_base;
 
+mod transform_retry_async;
 mod transform_sort_merge;
 mod transform_sort_merge_limit;
 pub mod transform_sort_partial;
@@ -38,6 +39,7 @@ pub use transform_blocking::*;
 pub use transform_compact::*;
 pub use transform_dummy::*;
 pub use transform_multi_sort_merge::try_add_multi_sort_merge;
+pub use transform_retry_async::*;
 pub use transform_sort_merge::sort_merge;
 pub use transform_sort_merge::*;
 pub use transform_sort_merge_base::*;

src/query/pipeline/transforms/src/processors/transforms/mod.rs

@@ -0,0 +1,74 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use databend_common_exception::Result;
+use databend_common_expression::DataBlock;
+
+use super::AsyncTransform;
+
+pub trait AsyncRetry: AsyncTransform {
+    fn retry_on(&self, err: &databend_common_exception::ErrorCode) -> bool;
+    fn retry_strategy(&self) -> RetryStrategy;
+}
+
+#[derive(Clone)]
+pub struct RetryStrategy {
+    pub retry_times: usize,
+    pub retry_sleep_duration: Option<tokio::time::Duration>,
+}
+
+pub struct AsyncRetryWrapper<T: AsyncRetry + 'static> {
+    t: T,
+}
+
+impl<T: AsyncRetry + 'static> AsyncRetryWrapper<T> {
+    pub fn create(inner: T) -> Self {
+        Self { t: inner }
+    }
+}
+
+#[async_trait::async_trait]
+impl<T: AsyncRetry + 'static> AsyncTransform for AsyncRetryWrapper<T> {
+    const NAME: &'static str = T::NAME;
+
+    async fn transform(&mut self, data: DataBlock) -> Result<DataBlock> {
+        let strategy = self.t.retry_strategy();
+        for _ in 0..strategy.retry_times {
+            match self.t.transform(data.clone()).await {
+                Ok(v) => return Ok(v),
+                Err(e) => {
+                    if !self.t.retry_on(&e) {
+                        return Err(e);
+                    }
+                    if let Some(duration) = strategy.retry_sleep_duration {
+                        tokio::time::sleep(duration).await;
+                    }
+                }
+            }
+        }
+        self.t.transform(data.clone()).await
+    }
+
+    fn name(&self) -> String {
+        Self::NAME.to_string()
+    }
+
+    async fn on_start(&mut self) -> Result<()> {
+        self.t.on_start().await
+    }
+
+    async fn on_finish(&mut self) -> Result<()> {
+        self.t.on_finish().await
+    }
+}

src/query/pipeline/transforms/src/processors/transforms/transform_retry_async.rs

@@ -24,8 +24,11 @@ use databend_common_expression::DataBlock;
 use databend_common_expression::DataField;
 use databend_common_expression::DataSchema;
 use databend_common_expression::FunctionContext;
+use databend_common_pipeline_transforms::processors::AsyncRetry;
+use databend_common_pipeline_transforms::processors::AsyncRetryWrapper;
 use databend_common_pipeline_transforms::processors::AsyncTransform;
 use databend_common_pipeline_transforms::processors::AsyncTransformer;
+use databend_common_pipeline_transforms::processors::RetryStrategy;
 use databend_common_sql::executor::physical_plans::UdfFunctionDesc;
 
 use crate::pipelines::processors::InputPort;
@@ -44,10 +47,22 @@ impl TransformUdfServer {
         input: Arc<InputPort>,
         output: Arc<OutputPort>,
     ) -> Result<Box<dyn Processor>> {
-        Ok(AsyncTransformer::create(input, output, Self {
-            func_ctx,
-            funcs,
-        }))
+        let s = Self { func_ctx, funcs };
+        let retry_wrapper = AsyncRetryWrapper::create(s);
+        Ok(AsyncTransformer::create(input, output, retry_wrapper))
+    }
+}
+
+impl AsyncRetry for TransformUdfServer {
+    fn retry_on(&self, err: &databend_common_exception::ErrorCode) -> bool {
+        true
+    }
+
+    fn retry_strategy(&self) -> RetryStrategy {
+        RetryStrategy {
+            retry_times: 64,
+            retry_sleep_duration: Some(tokio::time::Duration::from_millis(500)),
+        }
     }
 }