feat: Add cluster maintenance mode

quickwit/quickwit-cli/src/cli.rs

@@ -18,6 +18,7 @@ use quickwit_serve::EnvFilterReloadFn;
 use tracing::Level;
 
 use crate::index::{IndexCliCommand, build_index_command};
+use crate::maintenance::{MaintenanceCliCommand, build_maintenance_command};
 use crate::service::{RunCliCommand, build_run_command};
 use crate::source::{SourceCliCommand, build_source_command};
 use crate::split::{SplitCliCommand, build_split_command};
@@ -47,6 +48,7 @@ pub fn build_cli() -> Command {
         .subcommand(build_source_command().display_order(3))
         .subcommand(build_split_command().display_order(4))
         .subcommand(build_tool_command().display_order(5))
+        .subcommand(build_maintenance_command().display_order(6))
         .arg_required_else_help(true)
         .disable_help_subcommand(true)
         .subcommand_required(true)
@@ -59,6 +61,7 @@ pub enum CliCommand {
     Split(SplitCliCommand),
     Source(SourceCliCommand),
     Tool(ToolCliCommand),
+    Maintenance(MaintenanceCliCommand),
 }
 
 impl CliCommand {
@@ -69,6 +72,7 @@ impl CliCommand {
             CliCommand::Source(_) => Level::ERROR,
             CliCommand::Split(_) => Level::ERROR,
             CliCommand::Tool(_) => Level::ERROR,
+            CliCommand::Maintenance(_) => Level::ERROR,
         }
     }
 
@@ -82,6 +86,9 @@ impl CliCommand {
             "source" => SourceCliCommand::parse_cli_args(submatches).map(CliCommand::Source),
             "split" => SplitCliCommand::parse_cli_args(submatches).map(CliCommand::Split),
             "tool" => ToolCliCommand::parse_cli_args(submatches).map(CliCommand::Tool),
+            "maintenance" => {
+                MaintenanceCliCommand::parse_cli_args(submatches).map(CliCommand::Maintenance)
+            }
             _ => bail!("unknown command `{subcommand}`"),
         }
     }
@@ -93,6 +100,7 @@ impl CliCommand {
             CliCommand::Source(subcommand) => subcommand.execute().await,
             CliCommand::Split(subcommand) => subcommand.execute().await,
             CliCommand::Tool(subcommand) => subcommand.execute().await,
+            CliCommand::Maintenance(subcommand) => subcommand.execute().await,
         }
     }
 }

quickwit/quickwit-cli/src/lib.rs

@@ -50,6 +50,7 @@ pub mod index;
 #[cfg(feature = "jemalloc")]
 pub mod jemalloc;
 pub mod logger;
+pub mod maintenance;
 pub mod metrics;
 pub mod service;
 pub mod source;

quickwit/quickwit-cli/src/maintenance.rs

@@ -0,0 +1,149 @@
+// Copyright 2021-Present Datadog, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use anyhow::{Context, bail};
+use clap::{ArgMatches, Command};
+use colored::Colorize;
+use tracing::debug;
+
+use crate::checklist::{GREEN_COLOR, RED_COLOR};
+use crate::{ClientArgs, client_args};
+
+pub fn build_maintenance_command() -> Command {
+    Command::new("maintenance")
+        .about("Manages cluster maintenance mode for safe rolling upgrades.")
+        .args(client_args())
+        .subcommand(Command::new("enable").about(
+            "Enables maintenance mode. Freezes the indexing plan; metadata mutations are accepted \
+             but the plan is not rebuilt.",
+        ))
+        .subcommand(
+            Command::new("disable")
+                .about("Disables maintenance mode and triggers a full indexing plan rebuild."),
+        )
+        .subcommand(Command::new("status").about("Shows the current maintenance mode status."))
+        .subcommand_required(true)
+        .arg_required_else_help(true)
+}
+
+#[derive(Debug, PartialEq)]
+pub struct EnableMaintenanceArgs {
+    pub client_args: ClientArgs,
+}
+
+#[derive(Debug, PartialEq)]
+pub struct DisableMaintenanceArgs {
+    pub client_args: ClientArgs,
+}
+
+#[derive(Debug, PartialEq)]
+pub struct MaintenanceStatusArgs {
+    pub client_args: ClientArgs,
+}
+
+#[derive(Debug, PartialEq)]
+pub enum MaintenanceCliCommand {
+    Enable(EnableMaintenanceArgs),
+    Disable(DisableMaintenanceArgs),
+    Status(MaintenanceStatusArgs),
+}
+
+impl MaintenanceCliCommand {
+    pub fn parse_cli_args(mut matches: ArgMatches) -> anyhow::Result<Self> {
+        let (subcommand, submatches) = matches
+            .remove_subcommand()
+            .context("failed to parse maintenance subcommand")?;
+        match subcommand.as_str() {
+            "enable" => Self::parse_enable_args(submatches),
+            "disable" => Self::parse_disable_args(submatches),
+            "status" => Self::parse_status_args(submatches),
+            _ => bail!("unknown maintenance subcommand `{subcommand}`"),
+        }
+    }
+
+    fn parse_enable_args(mut matches: ArgMatches) -> anyhow::Result<Self> {
+        let client_args = ClientArgs::parse(&mut matches)?;
+        Ok(Self::Enable(EnableMaintenanceArgs { client_args }))
+    }
+
+    fn parse_disable_args(mut matches: ArgMatches) -> anyhow::Result<Self> {
+        let client_args = ClientArgs::parse(&mut matches)?;
+        Ok(Self::Disable(DisableMaintenanceArgs { client_args }))
+    }
+
+    fn parse_status_args(mut matches: ArgMatches) -> anyhow::Result<Self> {
+        let client_args = ClientArgs::parse(&mut matches)?;
+        Ok(Self::Status(MaintenanceStatusArgs { client_args }))
+    }
+
+    pub fn default_log_level(&self) -> tracing::Level {
+        tracing::Level::ERROR
+    }
+
+    pub async fn execute(self) -> anyhow::Result<()> {
+        match self {
+            Self::Enable(args) => enable_maintenance_cli(args).await,
+            Self::Disable(args) => disable_maintenance_cli(args).await,
+            Self::Status(args) => maintenance_status_cli(args).await,
+        }
+    }
+}
+
+async fn enable_maintenance_cli(args: EnableMaintenanceArgs) -> anyhow::Result<()> {
+    debug!(args=?args, "enable-maintenance");
+    println!("❯ Enabling maintenance mode...");
+    let qw_client = args.client_args.client();
+    let response = qw_client.maintenance().enable().await?;
+    println!(
+        "{} Maintenance mode enabled. Indexing plan frozen.",
+        "✔".color(GREEN_COLOR)
+    );
+    debug!(frozen_plan_json_len = response.frozen_plan_json.len());
+    Ok(())
+}
+
+async fn disable_maintenance_cli(args: DisableMaintenanceArgs) -> anyhow::Result<()> {
+    debug!(args=?args, "disable-maintenance");
+    println!("❯ Disabling maintenance mode...");
+    let qw_client = args.client_args.client();
+    qw_client.maintenance().disable().await?;
+    println!(
+        "{} Maintenance mode disabled. Indexing plan rebuild triggered.",
+        "✔".color(GREEN_COLOR)
+    );
+    Ok(())
+}
+
+async fn maintenance_status_cli(args: MaintenanceStatusArgs) -> anyhow::Result<()> {
+    debug!(args=?args, "maintenance-status");
+    let qw_client = args.client_args.client();
+    let status = qw_client.maintenance().status().await?;
+    if status.is_maintenance_mode {
+        println!(
+            "{} Maintenance mode is {}",
+            "●".color(RED_COLOR),
+            "ENABLED".color(RED_COLOR).bold()
+        );
+        if let Some(enabled_at) = status.enabled_at {
+            println!("  Enabled at: {enabled_at}");
+        }
+    } else {
+        println!(
+            "{} Maintenance mode is {}",
+            "●".color(GREEN_COLOR),
+            "DISABLED".color(GREEN_COLOR).bold()
+        );
+    }
+    Ok(())
+}

quickwit/quickwit-control-plane/Cargo.toml

@@ -13,17 +13,20 @@ license.workspace = true
 [dependencies]
 anyhow = { workspace = true }
 async-trait = { workspace = true }
+base64 = { workspace = true }
 bytesize = { workspace = true }
 fnv = { workspace = true }
 futures = { workspace = true }
 itertools = { workspace = true }
 lru = { workspace = true }
 mockall = { workspace = true, optional = true }
 once_cell = { workspace = true }
+prost = { workspace = true }
 rand = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 smallvec = { workspace = true }
+time = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }
 ulid = { workspace = true }