feat: enhance maintenance mode API and handling (#6022)

* feat(meta): enhance maintenance mode API and handling

* chore: apply suggestions from CR

src/meta-srv/src/region/supervisor.rs

@@ -359,7 +359,11 @@ impl RegionSupervisor {
         match self.is_maintenance_mode_enabled().await {
             Ok(false) => {}
             Ok(true) => {
-                info!("Maintenance mode is enabled, skip failover");
+                warn!(
+                    "Skipping failover since maintenance mode is enabled. Removing failure detectors for regions: {:?}",
+                    regions
+                );
+                self.deregister_failure_detectors(regions).await;
                 return;
             }
             Err(err) => {

src/meta-srv/src/service/admin.rs

@@ -202,8 +202,14 @@ fn boxed(body: String) -> BoxBody {
 
 #[cfg(test)]
 mod tests {
+    use common_meta::kv_backend::memory::MemoryKvBackend;
+    use common_meta::kv_backend::KvBackendRef;
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
     use super::*;
-    use crate::error;
+    use crate::metasrv::builder::MetasrvBuilder;
+    use crate::metasrv::MetasrvOptions;
+    use crate::{bootstrap, error};
 
     struct MockOkHandler;
 
@@ -308,4 +314,102 @@ mod tests {
 
         assert_eq!(http::StatusCode::INTERNAL_SERVER_ERROR, res.status());
     }
+
+    async fn test_metasrv(kv_backend: KvBackendRef) -> Metasrv {
+        let opts = MetasrvOptions::default();
+        let builder = MetasrvBuilder::new()
+            .options(opts)
+            .kv_backend(kv_backend.clone());
+
+        let metasrv = builder.build().await.unwrap();
+        metasrv
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_metasrv_maintenance_mode() {
+        common_telemetry::init_default_ut_logging();
+        let kv_backend = Arc::new(MemoryKvBackend::new());
+        let metasrv = test_metasrv(kv_backend).await;
+        metasrv.try_start().await.unwrap();
+
+        let (mut client, server) = tokio::io::duplex(1024);
+        let metasrv = Arc::new(metasrv);
+        let service = metasrv.clone();
+        let _handle = tokio::spawn(async move {
+            let router = bootstrap::router(service);
+            router
+                .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]))
+                .await
+        });
+
+        // Get maintenance mode
+        let http_request = b"GET /admin/maintenance HTTP/1.1\r\nHost: localhost\r\n\r\n";
+        client.write_all(http_request).await.unwrap();
+        let mut buf = vec![0; 1024];
+        let n = client.read(&mut buf).await.unwrap();
+        let response = String::from_utf8_lossy(&buf[..n]);
+        assert!(response.contains(r#"{"enabled":false}"#));
+        assert!(response.contains("200 OK"));
+
+        // Set maintenance mode to true
+        let http_post = b"POST /admin/maintenance?enable=true HTTP/1.1\r\nHost: localhost\r\nContent-Length: 0\r\n\r\n";
+        client.write_all(http_post).await.unwrap();
+        let mut buf = vec![0; 1024];
+        let n = client.read(&mut buf).await.unwrap();
+        let response = String::from_utf8_lossy(&buf[..n]);
+        assert!(response.contains(r#"{"enabled":true}"#));
+        assert!(response.contains("200 OK"));
+
+        let enabled = metasrv
+            .maintenance_mode_manager()
+            .maintenance_mode()
+            .await
+            .unwrap();
+        assert!(enabled);
+
+        // Get maintenance mode again
+        let http_request = b"GET /admin/maintenance HTTP/1.1\r\nHost: localhost\r\n\r\n";
+        client.write_all(http_request).await.unwrap();
+        let mut buf = vec![0; 1024];
+        let n = client.read(&mut buf).await.unwrap();
+        let response = String::from_utf8_lossy(&buf[..n]);
+        assert!(response.contains(r#"{"enabled":true}"#));
+        assert!(response.contains("200 OK"));
+
+        // Set maintenance mode to false
+        let http_post = b"POST /admin/maintenance?enable=false HTTP/1.1\r\nHost: localhost\r\nContent-Length: 0\r\n\r\n";
+        client.write_all(http_post).await.unwrap();
+        let mut buf = vec![0; 1024];
+        let n = client.read(&mut buf).await.unwrap();
+        let response = String::from_utf8_lossy(&buf[..n]);
+        assert!(response.contains(r#"{"enabled":false}"#));
+        assert!(response.contains("200 OK"));
+
+        let enabled = metasrv
+            .maintenance_mode_manager()
+            .maintenance_mode()
+            .await
+            .unwrap();
+        assert!(!enabled);
+
+        // Set maintenance mode to true via GET request
+        let http_request =
+            b"GET /admin/maintenance?enable=true HTTP/1.1\r\nHost: localhost\r\n\r\n";
+        client.write_all(http_request).await.unwrap();
+        let mut buf = vec![0; 1024];
+        let n = client.read(&mut buf).await.unwrap();
+        let response = String::from_utf8_lossy(&buf[..n]);
+        assert!(response.contains(r#"{"enabled":true}"#));
+        assert!(response.contains("200 OK"));
+
+        // Set maintenance mode to false via GET request
+        let http_request =
+            b"PUT /admin/maintenance?enable=false HTTP/1.1\r\nHost: localhost\r\n\r\n";
+        client.write_all(http_request).await.unwrap();
+        let mut buf = vec![0; 1024];
+        let n = client.read(&mut buf).await.unwrap();
+        let response = String::from_utf8_lossy(&buf[..n]);
+        assert!(response.contains(r#"{"enabled":false}"#));
+        assert!(response.contains("200 OK"));
+    }
 }

src/meta-srv/src/service/admin/maintenance.rs

@@ -15,13 +15,15 @@
 use std::collections::HashMap;
 
 use common_meta::key::maintenance::MaintenanceModeManagerRef;
+use common_telemetry::{info, warn};
+use serde::{Deserialize, Serialize};
 use snafu::{OptionExt, ResultExt};
 use tonic::codegen::http;
 use tonic::codegen::http::Response;
 
 use crate::error::{
-    InvalidHttpBodySnafu, MaintenanceModeManagerSnafu, MissingRequiredParameterSnafu,
-    ParseBoolSnafu, UnsupportedSnafu,
+    self, InvalidHttpBodySnafu, MaintenanceModeManagerSnafu, MissingRequiredParameterSnafu,
+    ParseBoolSnafu, Result, UnsupportedSnafu,
 };
 use crate::service::admin::HttpHandler;
 
@@ -30,6 +32,21 @@ pub struct MaintenanceHandler {
     pub manager: MaintenanceModeManagerRef,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+struct MaintenanceResponse {
+    enabled: bool,
+}
+
+impl TryFrom<MaintenanceResponse> for String {
+    type Error = error::Error;
+
+    fn try_from(response: MaintenanceResponse) -> Result<Self> {
+        serde_json::to_string(&response).context(error::SerializeToJsonSnafu {
+            input: format!("{response:?}"),
+        })
+    }
+}
+
 impl MaintenanceHandler {
     async fn get_maintenance(&self) -> crate::Result<Response<String>> {
         let enabled = self
@@ -37,14 +54,10 @@ impl MaintenanceHandler {
             .maintenance_mode()
             .await
             .context(MaintenanceModeManagerSnafu)?;
-        let response = if enabled {
-            "Maintenance mode is enabled"
-        } else {
-            "Maintenance mode is disabled"
-        };
+        let response = MaintenanceResponse { enabled }.try_into()?;
         http::Response::builder()
             .status(http::StatusCode::OK)
-            .body(response.into())
+            .body(response)
             .context(InvalidHttpBodySnafu)
     }
 
@@ -60,23 +73,24 @@ impl MaintenanceHandler {
                 err_msg: "'enable' must be 'true' or 'false'",
             })?;
 
-        let response = if enable {
+        if enable {
             self.manager
                 .set_maintenance_mode()
                 .await
                 .context(MaintenanceModeManagerSnafu)?;
-            "Maintenance mode enabled"
+            info!("Enable the maintenance mode.");
         } else {
             self.manager
                 .unset_maintenance_mode()
                 .await
                 .context(MaintenanceModeManagerSnafu)?;
-            "Maintenance mode disabled"
+            info!("Disable the maintenance mode.");
         };
 
+        let response = MaintenanceResponse { enabled: enable }.try_into()?;
         http::Response::builder()
             .status(http::StatusCode::OK)
-            .body(response.into())
+            .body(response)
             .context(InvalidHttpBodySnafu)
     }
 }
@@ -90,8 +104,23 @@ impl HttpHandler for MaintenanceHandler {
         params: &HashMap<String, String>,
     ) -> crate::Result<Response<String>> {
         match method {
-            http::Method::GET => self.get_maintenance().await,
-            http::Method::PUT => self.set_maintenance(params).await,
+            http::Method::GET => {
+                if params.is_empty() {
+                    self.get_maintenance().await
+                } else {
+                    warn!(
+                        "Found URL parameters in '/admin/maintenance' request, it's deprecated, will be removed in the future"
+                    );
+                    // The old version operator will send GET request with URL parameters,
+                    // so we need to support it.
+                    self.set_maintenance(params).await
+                }
+            }
+            http::Method::PUT => {
+                warn!("Found PUT request to '/admin/maintenance', it's deprecated, will be removed in the future");
+                self.set_maintenance(params).await
+            }
+            http::Method::POST => self.set_maintenance(params).await,
             _ => UnsupportedSnafu {
                 operation: format!("http method {method}"),
             }