address review comments

risingwavelabs · Feb 27, 2024 · 44bd236 · 44bd236
1 parent 5adfcec
commit 44bd236
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 26 deletions.
diff --git a/e2e_test/udf/always_retry_python.slt b/e2e_test/udf/always_retry_python.slt
@@ -69,4 +69,7 @@ statement ok
 DROP FUNCTION sleep_always_retry;
 
 statement ok
-DROP FUNCTION sleep_no_retry;
+DROP FUNCTION sleep_no_retry;
+
+statement ok
+DROP TABLE t CASCADE;
diff --git a/src/expr/core/src/expr/expr_udf.rs b/src/expr/core/src/expr/expr_udf.rs
@@ -115,35 +115,38 @@ impl UserDefinedFunction {
                 let disable_retry_count = self.disable_retry_count.load(Ordering::Relaxed);
                 let result = if self.always_retry_on_network_error {
                     client
-                        .call_with_always_retry_on_network_error(&self.identifier, input)
-                        .instrument_await(self.span.clone())
-                        .await
-                } else if disable_retry_count != 0 {
-                    client
-                        .call(&self.identifier, arrow_input)
+                        .call_with_always_retry_on_network_error(&self.identifier, arrow_input)
                         .instrument_await(self.span.clone())
                         .await
                 } else {
-                    client
-                        .call_with_retry(&self.identifier, arrow_input)
-                        .instrument_await(self.span.clone())
-                        .await
+                    let result = if disable_retry_count != 0 {
+                        client
+                            .call(&self.identifier, arrow_input)
+                            .instrument_await(self.span.clone())
+                            .await
+                    } else {
+                        client
+                            .call_with_retry(&self.identifier, arrow_input)
+                            .instrument_await(self.span.clone())
+                            .await
+                    };
+                    let disable_retry_count = self.disable_retry_count.load(Ordering::Relaxed);
+                    let connection_error = matches!(&result, Err(e) if e.is_connection_error());
+                    if connection_error && disable_retry_count != INITIAL_RETRY_COUNT {
+                        // reset count on connection error
+                        self.disable_retry_count
+                            .store(INITIAL_RETRY_COUNT, Ordering::Relaxed);
+                    } else if !connection_error && disable_retry_count != 0 {
+                        // decrease count on success, ignore if exchange failed
+                        _ = self.disable_retry_count.compare_exchange(
+                            disable_retry_count,
+                            disable_retry_count - 1,
+                            Ordering::Relaxed,
+                            Ordering::Relaxed,
+                        );
+                    }
+                    result
                 };
-                let disable_retry_count = self.disable_retry_count.load(Ordering::Relaxed);
-                let connection_error = matches!(&result, Err(e) if e.is_connection_error());
-                if connection_error && disable_retry_count != INITIAL_RETRY_COUNT {
-                    // reset count on connection error
-                    self.disable_retry_count
-                        .store(INITIAL_RETRY_COUNT, Ordering::Relaxed);
-                } else if !connection_error && disable_retry_count != 0 {
-                    // decrease count on success, ignore if exchange failed
-                    _ = self.disable_retry_count.compare_exchange(
-                        disable_retry_count,
-                        disable_retry_count - 1,
-                        Ordering::Relaxed,
-                        Ordering::Relaxed,
-                    );
-                }
                 result?
             }
         };