fix: fix incorrect matches (#2430)

* fix: fix incorrect matches

* fix: fix incorrect status code

src/client/src/error.rs

@@ -72,6 +72,9 @@ pub enum Error {
         source: common_grpc::error::Error,
     },
 
+    #[snafu(display("Failed to request RegionServer, code: {}, source: {}", code, source))]
+    RegionServer { code: Code, source: BoxedError },
+
     // Server error carried in Tonic Status's metadata.
     #[snafu(display("{}", msg))]
     Server { code: StatusCode, msg: String },
@@ -95,9 +98,9 @@ impl ErrorExt for Error {
             | Error::ClientStreaming { .. } => StatusCode::Internal,
 
             Error::Server { code, .. } => *code,
-            Error::FlightGet { source, .. } | Error::HandleRequest { source, .. } => {
-                source.status_code()
-            }
+            Error::FlightGet { source, .. }
+            | Error::HandleRequest { source, .. }
+            | Error::RegionServer { source, .. } => source.status_code(),
             Error::CreateChannel { source, .. } | Error::ConvertFlightData { source, .. } => {
                 source.status_code()
             }

src/client/src/region.rs

@@ -29,7 +29,7 @@ use prost::Message;
 use snafu::{location, Location, OptionExt, ResultExt};
 use tokio_stream::StreamExt;
 
-use crate::error::Error::FlightGet;
+use crate::error::Error::RegionServer;
 use crate::error::{
     self, ConvertFlightDataSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu,
     MissingFieldSnafu, Result, ServerSnafu,
@@ -45,7 +45,7 @@ pub struct RegionRequester {
 impl Datanode for RegionRequester {
     async fn handle(&self, request: RegionRequest) -> MetaResult<AffectedRows> {
         self.handle_inner(request).await.map_err(|err| {
-            if matches!(err, FlightGet { .. }) {
+            if matches!(err, RegionServer { .. }) {
                 meta_error::Error::RetryLater {
                     source: BoxedError::new(err),
                 }
@@ -163,7 +163,19 @@ impl RegionRequester {
         let RegionResponse {
             header,
             affected_rows,
-        } = client.handle(request).await?.into_inner();
+        } = client
+            .handle(request)
+            .await
+            .map_err(|e| {
+                let code = e.code();
+                let err: error::Error = e.into();
+                // Uses `Error::RegionServer` instead of `Error::Server`
+                error::Error::RegionServer {
+                    code,
+                    source: BoxedError::new(err),
+                }
+            })?
+            .into_inner();
 
         check_response_header(header)?;
 

src/datanode/src/alive_keeper.rs

@@ -518,8 +518,8 @@ mod test {
         // wait for countdown task to finish
         let before_await = Instant::now();
         let (finish_instant, result) = rx.await.unwrap();
-        // the mock region server cannot close the region
-        assert_eq!(result, Some(false));
+        // it returns `RegionNotFound`
+        assert_eq!(result, Some(true));
         // this task should be finished after 5 * heartbeat_interval_millis
         // we assert 4 times here
         assert!(

src/datanode/src/error.rs

@@ -518,11 +518,12 @@ impl ErrorExt for Error {
             | MissingInsertBody { .. }
             | ShutdownInstance { .. }
             | JoinTask { .. }
-            | RegionNotFound { .. }
             | RegionEngineNotFound { .. }
             | UnsupportedOutput { .. }
             | GetRegionMetadata { .. } => StatusCode::Internal,
 
+            RegionNotFound { .. } => StatusCode::RegionNotFound,
+
             StartServer { source, .. }
             | ShutdownServer { source, .. }
             | WaitForGrpcServing { source, .. } => source.status_code(),