From b715bde58b43f2c8555e7aeaae405d369acd14f7 Mon Sep 17 00:00:00 2001
From: Bugen Zhao <i@bugenzhao.com>
Date: Tue, 19 Sep 2023 15:08:02 +0800
Subject: [PATCH] refactor(udf): connect lazily to get rid of `block_in_place`
 (#12413)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>
---
 src/expr/src/expr/expr_udf.rs | 4 +---
 src/udf/src/error.rs          | 2 +-
 src/udf/src/external.rs       | 7 +++++++
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/expr/src/expr/expr_udf.rs b/src/expr/src/expr/expr_udf.rs
index d0ecf58c3af9..b7a03a6611f5 100644
--- a/src/expr/src/expr/expr_udf.rs
+++ b/src/expr/src/expr/expr_udf.rs
@@ -164,9 +164,7 @@ pub(crate) fn get_or_create_client(link: &str) -> Result<Arc<ArrowFlightUdfClien
         Ok(client)
     } else {
         // create new client
-        let client = Arc::new(tokio::task::block_in_place(|| {
-            tokio::runtime::Handle::current().block_on(ArrowFlightUdfClient::connect(link))
-        })?);
+        let client = Arc::new(ArrowFlightUdfClient::connect_lazy(link)?);
         clients.insert(link.into(), Arc::downgrade(&client));
         Ok(client)
     }
diff --git a/src/udf/src/error.rs b/src/udf/src/error.rs
index f20816ee5b2c..98b9faa57171 100644
--- a/src/udf/src/error.rs
+++ b/src/udf/src/error.rs
@@ -23,7 +23,7 @@ pub enum Error {
     #[error("failed to connect to UDF service: {0}")]
     Connect(#[from] tonic::transport::Error),
 
-    #[error("failed to check UDF: {0}")]
+    #[error("failed to send requests to UDF service: {0}")]
     Tonic(#[from] Box<tonic::Status>),
 
     #[error("failed to call UDF: {0}")]
diff --git a/src/udf/src/external.rs b/src/udf/src/external.rs
index 585adc7ebec5..d2e6afe71233 100644
--- a/src/udf/src/external.rs
+++ b/src/udf/src/external.rs
@@ -38,6 +38,13 @@ impl ArrowFlightUdfClient {
         Ok(Self { client })
     }
 
+    /// Connect to a UDF service lazily (i.e. only when the first request is sent).
+    pub fn connect_lazy(addr: &str) -> Result<Self> {
+        let conn = tonic::transport::Endpoint::new(addr.to_string())?.connect_lazy();
+        let client = FlightServiceClient::new(conn);
+        Ok(Self { client })
+    }
+
     /// Check if the function is available and the schema is match.
     pub async fn check(&self, id: &str, args: &Schema, returns: &Schema) -> Result<()> {
         let descriptor = FlightDescriptor::new_path(vec![id.into()]);