refactor(flight): enhance error handling during request serialization and deserialization

Author: dqhl76

src/query/service/src/servers/flight/v1/actions/flight_actions.rs

@@ -69,54 +69,72 @@ impl FlightActions {
         self.actions.insert(
             path.clone(),
             Box::new(move |request| {
-                let mut deserializer = serde_json::Deserializer::from_slice(request);
-                deserializer.disable_recursion_limit();
-
-                let deserializer = serde_stacker::Deserializer {
-                    de: &mut deserializer,
-                    red_zone: recursive::get_minimum_stack_size(),
-                    stack_size: recursive::get_stack_allocation_size(),
-                };
-
-                let request = Req::deserialize(deserializer).map_err(|cause| {
-                    ErrorCode::BadArguments(format!(
-                        "Cannot parse request for {}, cause: {:?}",
-                        path, cause
-                    ))
-                });
+                let request =
+                    match catch_unwind(|| -> std::result::Result<Req, serde_json::Error> {
+                        let mut deserializer = serde_json::Deserializer::from_slice(request);
+                        deserializer.disable_recursion_limit();
+
+                        let deserializer = serde_stacker::Deserializer {
+                            de: &mut deserializer,
+                            red_zone: recursive::get_minimum_stack_size(),
+                            stack_size: recursive::get_stack_allocation_size(),
+                        };
+
+                        Req::deserialize(deserializer)
+                    }) {
+                        Ok(Ok(request)) => Ok(request),
+                        Ok(Err(cause)) => Err(ErrorCode::BadArguments(format!(
+                            "Cannot parse request for {}, len: {}, cause: {:?}",
+                            path,
+                            request.len(),
+                            cause
+                        ))),
+                        Err(cause) => Err(cause.add_message_back(format!(
+                            "(while deserializing flight action request: action={}, len={})",
+                            path,
+                            request.len()
+                        ))),
+                    };
 
                 let path = path.clone();
                 let t = t.clone();
                 Box::pin(async move {
                     let request = request?;
 
-                    let future = catch_unwind(move || t(request))?;
+                    let future = catch_unwind(move || t(request)).map_err(|cause| {
+                        cause.add_message_back(format!(
+                            "(while creating flight action future: action={})",
+                            path
+                        ))
+                    })?;
 
                     let future = CatchUnwindFuture::create(future);
-                    match future
+                    let response = future
                         .await
-                        .with_context(|| "failed to do flight action")
-                        .flatten()
-                    {
-                        Ok(v) => {
-                            let mut out = Vec::with_capacity(512);
-                            let mut serializer = serde_json::Serializer::new(&mut out);
-                            let serializer = serde_stacker::Serializer {
-                                ser: &mut serializer,
-                                red_zone: recursive::get_minimum_stack_size(),
-                                stack_size: recursive::get_stack_allocation_size(),
-                            };
-
-                            v.serialize(serializer).map_err(|cause| {
-                                ErrorCode::BadBytes(format!(
-                                    "Cannot serialize response for {}, cause: {:?}",
-                                    path, cause
-                                ))
-                            })?;
-
-                            Ok(out)
-                        }
-                        Err(err) => Err(err),
+                        .with_context(|| format!("failed to do flight action, action: {}", path))
+                        .flatten()?;
+
+                    match catch_unwind(|| -> std::result::Result<Vec<u8>, serde_json::Error> {
+                        let mut out = Vec::with_capacity(512);
+                        let mut serializer = serde_json::Serializer::new(&mut out);
+                        let serializer = serde_stacker::Serializer {
+                            ser: &mut serializer,
+                            red_zone: recursive::get_minimum_stack_size(),
+                            stack_size: recursive::get_stack_allocation_size(),
+                        };
+
+                        response.serialize(serializer)?;
+                        Ok(out)
+                    }) {
+                        Ok(Ok(out)) => Ok(out),
+                        Ok(Err(cause)) => Err(ErrorCode::BadBytes(format!(
+                            "Cannot serialize response for {}, cause: {:?}",
+                            path, cause
+                        ))),
+                        Err(cause) => Err(cause.add_message_back(format!(
+                            "(while serializing flight action response: action={})",
+                            path
+                        ))),
                     }
                 })
             }),

src/query/service/src/servers/flight/v1/flight_service.rs

@@ -34,6 +34,7 @@ use databend_common_exception::ErrorCode;
 use fastrace::func_path;
 use fastrace::prelude::*;
 use futures_util::stream;
+use log::error;
 use tokio_stream::Stream;
 use tonic::Request;
 use tonic::Response as RawResponse;
@@ -172,7 +173,17 @@ impl FlightService for DatabendQueryFlightService {
             .in_span(root)
             .await
         {
-            Err(cause) => Err(cause.into()),
+            Err(cause) => {
+                error!(
+                    "flight do_action failed, node: {}, action: {}, body_len: {}, code: {}, error: {:?}",
+                    config.query.node_id,
+                    action.r#type,
+                    action.body.len(),
+                    cause.code(),
+                    cause
+                );
+                Err(cause.into())
+            }
             Ok(body) => Ok(RawResponse::new(
                 Box::pin(tokio_stream::once(Ok(FlightResult { body: body.into() })))
                     as FlightStream<FlightResult>,