servers/http/
workload_scheduler.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::num::NonZeroU32;
16
17use axum::Json;
18use axum::body::Bytes;
19use axum::http::StatusCode;
20use axum::response::IntoResponse;
21use serde::{Deserialize, Serialize};
22use snafu::{ResultExt, ensure};
23
24use crate::error::{InvalidParameterSnafu, ParseJsonSnafu, Result};
25
26#[axum_macros::debug_handler]
27pub(super) async fn set_enabled_handler(body: Bytes) -> Result<impl IntoResponse> {
28    let enabled: bool = serde_json::from_slice(&body).context(ParseJsonSnafu)?;
29    ensure!(
30        common_runtime::set_workload_scheduler_enabled(enabled),
31        InvalidParameterSnafu {
32            reason: "workload scheduler was not constructed at startup",
33        }
34    );
35    let change_note = format!("Workload scheduler enabled={enabled}");
36    Ok((StatusCode::OK, change_note))
37}
38
39#[derive(Debug, Deserialize)]
40struct SchedulerWeightsDto {
41    query: NonZeroU32,
42    write: NonZeroU32,
43}
44
45#[axum_macros::debug_handler]
46pub(super) async fn set_weights_handler(body: Bytes) -> Result<impl IntoResponse> {
47    let weights: SchedulerWeightsDto = serde_json::from_slice(&body).context(ParseJsonSnafu)?;
48    ensure!(
49        common_runtime::set_workload_scheduler_weights(weights.query, weights.write),
50        InvalidParameterSnafu {
51            reason: "workload scheduler was not constructed at startup",
52        }
53    );
54    let change_note = format!(
55        "Workload scheduler weights query={}, write={}",
56        weights.query, weights.write
57    );
58    Ok((StatusCode::OK, change_note))
59}
60
61/// Per-class scheduler status exposed by the HTTP API.
62#[derive(Debug, Serialize)]
63struct ClassStatusDto {
64    weight: u32,
65}
66
67/// Point-in-time workload scheduler status. Scheduler class fields are omitted
68/// when the scheduler was not constructed at startup or the corresponding
69/// class is unavailable.
70#[derive(Debug, Serialize)]
71struct SchedulerStatusDto {
72    enabled: bool,
73    #[serde(skip_serializing_if = "Option::is_none")]
74    query: Option<ClassStatusDto>,
75    #[serde(skip_serializing_if = "Option::is_none")]
76    write: Option<ClassStatusDto>,
77}
78
79/// Returns the current workload scheduler state and query/write weights.
80/// Always returns 200, with `enabled=false` when the scheduler is dynamically
81/// disabled.
82#[axum_macros::debug_handler]
83pub(super) async fn get_status_handler() -> Result<impl IntoResponse> {
84    let enabled = common_runtime::workload_scheduler_enabled();
85    let Some(stats) = common_runtime::workload_scheduler_stats() else {
86        return Ok(Json(SchedulerStatusDto {
87            enabled: false,
88            query: None,
89            write: None,
90        }));
91    };
92
93    let mut query = None;
94    let mut write = None;
95    for (class, class_stats) in &stats.classes {
96        let status = ClassStatusDto {
97            weight: class_stats.weight,
98        };
99        match class.id() {
100            1 => query = Some(status),
101            2 => write = Some(status),
102            _ => {}
103        }
104    }
105
106    Ok(Json(SchedulerStatusDto {
107        enabled,
108        query,
109        write,
110    }))
111}
servers/http/workload_scheduler.rs

servers/http/
workload_scheduler.rs
