rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-13 22:36:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(meta): configure drain parallelism and strategy over HTTP (#10218)

Browse Source 
* feat(meta): configure drain parallelism and strategy over HTTP

Signed-off-by: WenyXu <wenymedia@gmail.com>

* feat(meta): configure drain timeout with human-readable durations

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(meta): reject overflowing drain deadlines

Signed-off-by: WenyXu <wenymedia@gmail.com>

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2026-11-07 04:49:42 +00:00
committed by GitHub
parent 4688fa4be8
commit 2fefad6a9e
2 changed files with 217 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/meta-srv/src/procedure/drain_datanode/manager.rs
+42 -17
View File
@@ -21,7 +21,7 @@ use common_meta::key::TableMetadataManagerRef;
use common_procedure::{ProcedureContext, ProcedureId, ProcedureManagerRef, ProcedureWithId};
use common_telemetry::info;
use common_time::util::current_time_millis;
use snafu::{ResultExt, ensure};
use snafu::{OptionExt, ResultExt, ensure};
use tokio::sync::Mutex;
use crate::error::{self, Result};
@@ -161,6 +161,14 @@ impl DrainDatanodeManager {
.fail();
}
// The budget starts at submission and must fit the persisted absolute deadline.
let deadline_ms = i64::try_from(timeout.as_millis())
.ok()
.and_then(|timeout_ms| current_time_millis().checked_add(timeout_ms))
.context(error::InvalidArgumentsSnafu {
err_msg: "the timeout exceeds the supported deadline range",
})?;
self.ensure_ready()?;
let Some(guard) = self.tracker.insert_running_procedure(&task) else {
return error::DrainDatanodeRunningSnafu {
@@ -176,7 +184,7 @@ impl DrainDatanodeManager {
strategy: task.strategy,
// The budget starts when the task is submitted, so neither a slow start nor a
// restart extends it.
deadline_ms: current_time_millis() + timeout.as_millis() as i64,
deadline_ms,
};
let procedure = DrainDatanodeProcedure::new(self.context(persistent_ctx), vec![guard]);
let procedure_with_id =
@@ -698,21 +706,38 @@ mod tests {
assert!(matches!(error, error::Error::InvalidArguments { .. }));
}
// The timeout must be greater than zero.
let error = manager
.submit_procedure(
common_procedure::ProcedureContext::default(),
DrainDatanodeProcedureTask {
node_id: 1,
target_node_id: Some(2),
parallelism: DEFAULT_PARALLELISM,
strategy: DrainDatanodeStrategy::FailFast,
timeout: Some(Duration::ZERO),
},
)
.await
.unwrap_err();
assert!(matches!(error, error::Error::InvalidArguments { .. }));
// Reject zero, millisecond conversion overflow, and absolute deadline overflow.
for timeout in [
Duration::ZERO,
Duration::from_millis(i64::MAX as u64 + 1),
Duration::from_millis(i64::MAX as u64),
] {
let error = manager
.submit_procedure(
common_procedure::ProcedureContext::default(),
DrainDatanodeProcedureTask {
timeout: Some(timeout),
..task()
},
)
.await
.unwrap_err();
assert!(matches!(error, error::Error::InvalidArguments { .. }));
assert!(!manager.tracker.contains(1));
}
assert!(
env.procedure_manager()
.list_procedures()
.await
.unwrap()
.is_empty()
);
assert!(
env.node_scheduling_manager
.is_schedulable(NodeSchedulingRole::Datanode, 1)
.await
.unwrap()
);
}
#[tokio::test]
src/meta-srv/src/service/admin/datanode.rs
+175 -11
View File
@@ -13,6 +13,7 @@
// limitations under the License.
use std::sync::Arc;
use std::time::Duration;
use axum::Json;
use axum::extract::{Query, State};
@@ -24,7 +25,7 @@ use servers::http::result::error_result::ErrorResponse;
use crate::key::NodeSchedulingRole;
use crate::metasrv::Metasrv;
use crate::procedure::drain_datanode::{
DEFAULT_PARALLELISM, DrainDatanodeProcedureTask, DrainDatanodeStrategy,
DEFAULT_DRAIN_TIMEOUT, DEFAULT_PARALLELISM, DrainDatanodeProcedureTask, DrainDatanodeStrategy,
};
use crate::service::admin::util;
@@ -32,6 +33,45 @@ use crate::service::admin::util;
pub(crate) struct DrainDatanodeRequest {
node_id: u64,
target_node_id: Option<u64>,
/// Concurrent region budget; defaults to [`DEFAULT_PARALLELISM`].
#[serde(default = "default_parallelism")]
parallelism: usize,
/// Failure handling; defaults to [`DrainDatanodeStrategy::TryAll`].
#[serde(default)]
strategy: DrainStrategy,
/// Whole-drain budget; defaults to [`DEFAULT_DRAIN_TIMEOUT`].
#[serde(default = "default_timeout", with = "humantime_serde")]
timeout: Duration,
}
/// HTTP strategy names follow the admin API's snake_case convention.
#[derive(Debug, Default, Deserialize)]
#[serde(rename_all = "snake_case")]
enum DrainStrategy {
FailFast,
#[default]
TryAll,
}
impl From<DrainStrategy> for DrainDatanodeStrategy {
fn from(strategy: DrainStrategy) -> Self {
match strategy {
DrainStrategy::FailFast => Self::FailFast,
DrainStrategy::TryAll => Self::TryAll,
}
}
}
impl From<DrainDatanodeRequest> for DrainDatanodeProcedureTask {
fn from(request: DrainDatanodeRequest) -> Self {
Self {
node_id: request.node_id,
target_node_id: request.target_node_id,
parallelism: request.parallelism,
strategy: request.strategy.into(),
timeout: Some(request.timeout),
}
}
}
#[derive(Debug, Deserialize)]
@@ -50,6 +90,14 @@ pub(crate) struct DatanodeSchedulingResponse {
schedulable: bool,
}
fn default_parallelism() -> usize {
DEFAULT_PARALLELISM
}
fn default_timeout() -> Duration {
DEFAULT_DRAIN_TIMEOUT
}
/// Submits an asynchronous leader-region drain.
/// Rejects requests for a datanode already being drained.
#[axum_macros::debug_handler]
@@ -61,16 +109,7 @@ pub(crate) async fn drain(
util::ensure_leader(&metasrv)?;
let procedure_id = metasrv
.drain_datanode_manager()
.submit_procedure(
ProcedureContext::default(),
DrainDatanodeProcedureTask {
node_id: request.node_id,
target_node_id: request.target_node_id,
parallelism: DEFAULT_PARALLELISM,
strategy: DrainDatanodeStrategy::default(),
timeout: None,
},
)
.submit_procedure(ProcedureContext::default(), request.into())
.await?;
Ok::<_, crate::error::Error>(DrainDatanodeResponse {
procedure_id: procedure_id.to_string(),
@@ -191,6 +230,131 @@ mod tests {
(status, serde_json::from_slice(&bytes).unwrap())
}
#[test]
fn test_drain_request_options() {
for (body, parallelism, strategy) in [
(json!({"node_id": 1}), 4, DrainDatanodeStrategy::TryAll),
(
json!({"node_id": 1, "parallelism": 8}),
8,
DrainDatanodeStrategy::TryAll,
),
(
json!({"node_id": 1, "strategy": "fail_fast"}),
4,
DrainDatanodeStrategy::FailFast,
),
(
json!({"node_id": 1, "parallelism": 1, "strategy": "fail_fast"}),
1,
DrainDatanodeStrategy::FailFast,
),
(
json!({"node_id": 1, "parallelism": 32, "strategy": "try_all"}),
32,
DrainDatanodeStrategy::TryAll,
),
] {
let request: DrainDatanodeRequest = serde_json::from_value(body).unwrap();
let task: DrainDatanodeProcedureTask = request.into();
assert_eq!(parallelism, task.parallelism);
assert_eq!(strategy, task.strategy);
assert_eq!(1, task.node_id);
assert_eq!(None, task.target_node_id);
assert_eq!(Some(Duration::from_secs(20 * 60)), task.timeout);
}
}
#[test]
fn test_drain_request_timeout() {
for (value, expected) in [
(json!("30m"), Some(Duration::from_secs(1800))),
(json!("1h 30m"), Some(Duration::from_secs(5400))),
(json!("500ms"), Some(Duration::from_millis(500))),
] {
let request: DrainDatanodeRequest =
serde_json::from_value(json!({"node_id": 1, "timeout": value})).unwrap();
let task: DrainDatanodeProcedureTask = request.into();
assert_eq!(expected, task.timeout);
}
}
#[tokio::test]
async fn test_http_drain_rejects_invalid_options() {
let metasrv = Arc::new(MetasrvBuilder::new().build().await.unwrap());
metasrv.procedure_manager().start().await.unwrap();
let app = admin_axum_router(metasrv.clone());
for (options, expected) in [
(json!({"timeout": null}), StatusCode::UNPROCESSABLE_ENTITY),
(
json!({"parallelism": null}),
StatusCode::UNPROCESSABLE_ENTITY,
),
(json!({"strategy": null}), StatusCode::BAD_REQUEST),
(json!({"timeout": "0s"}), StatusCode::BAD_REQUEST),
(
json!({"timeout": "9223372036854775s"}),
StatusCode::BAD_REQUEST,
),
(
json!({"timeout": "9223372036854776s"}),
StatusCode::BAD_REQUEST,
),
(
json!({"timeout": "invalid"}),
StatusCode::UNPROCESSABLE_ENTITY,
),
(json!({"timeout": "-1s"}), StatusCode::UNPROCESSABLE_ENTITY),
(json!({"parallelism": 0}), StatusCode::BAD_REQUEST),
(json!({"parallelism": 33}), StatusCode::BAD_REQUEST),
(json!({"parallelism": -1}), StatusCode::UNPROCESSABLE_ENTITY),
(
json!({"parallelism": 1.5}),
StatusCode::UNPROCESSABLE_ENTITY,
),
(
json!({"strategy": "unknown"}),
StatusCode::UNPROCESSABLE_ENTITY,
),
(
json!({"strategy": "FailFast"}),
StatusCode::UNPROCESSABLE_ENTITY,
),
] {
let mut body = options;
body["node_id"] = json!(1);
let response = app
.clone()
.oneshot(
Request::builder()
.method(Method::POST)
.uri("/admin/datanode/drain")
.header("content-type", "application/json")
.body(Body::from(body.to_string()))
.unwrap(),
)
.await
.unwrap();
assert_eq!(expected, response.status(), "{body}");
}
assert!(
metasrv
.procedure_manager()
.list_procedures()
.await
.unwrap()
.is_empty()
);
assert!(
metasrv
.node_scheduling_manager()
.is_schedulable(NodeSchedulingRole::Datanode, 1)
.await
.unwrap()
);
metasrv.procedure_manager().stop().await.unwrap();
}
#[tokio::test]
async fn test_http_scheduling_is_idempotent() {
let metasrv = Arc::new(MetasrvBuilder::new().build().await.unwrap());