From 43c30d14463ba27de36a8c50cfde41d7a6e09323 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 2 Sep 2026 07:02:39 +0000 Subject: [PATCH] feat(runtime): add weighted workload scheduler (#8736) * feat(runtime): add weighted workload scheduler Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(runtime): switch catio to GreptimeTeam fork with admission-wait metrics Use the GreptimeTeam/catio fork (pinned c20eafc) which adds ClassStats::total_admission_wait and ClassStats::admitted, recorded at each QUEUED -> ADMITTED transition. This exposes the scheduler's own admission delay (excluding Tokio queueing and poll execution), enabling admission-wait based fairness gates. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore: bump catio to dynamic-config revision Bump the catio scheduler fork to 9f4b028 which adds Scheduler::set_weight and Scheduler::set_max_concurrent_polls for runtime configuration. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(perf): runtime-adjustable workload scheduler parameters Expose dynamic adjustment of the experimental workload scheduler at runtime: - common-runtime: set_workload_scheduler_weights and set_workload_scheduler_max_concurrent_polls, which forward to the catio scheduler's set_weight/set_max_concurrent_polls when the scheduler is enabled and reject zero values. - servers: /debug/workload_scheduler/weights and /debug/workload_scheduler/max_concurrent_polls POST handlers, so operators can rebalance query/write shares or admission concurrency without restarting the datanode. Both endpoints return 400 with a clear reason when the scheduler is disabled or the requested value is invalid. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(perf): add GET /debug/workload_scheduler status endpoint Returns the current weights (per class), max_concurrent_polls, active_polls and per-class counters (queued, tasks, wakes, polls, completed, cancelled, admitted, total_admission_wait) as JSON. When the scheduler is disabled, returns enabled=false with the other fields omitted, so operators can distinguish 'disabled' from an error. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore: bump catio to time-accounting revision Bump the catio scheduler fork to 257ba56 which replaces admission-count accounting with real execution-time accounting (pass += exec_time / (weight * concurrency)), so CPU share follows the configured weights regardless of poll length. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore: bump catio to lock-free sampling revision Bump the catio scheduler fork to efdc0a4 which adds an optional downsampled clock sampling mode (SchedulerBuilder::sample_every_polls, default off) with a lock-free per-class atomic counter, so the downsampled path costs one fetch_add per poll instead of a global mutex. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore: pin catio to scheduler PR head Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(runtime): add scheduler bypass control Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore: advance catio scheduler fixes Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore: pin merged catio scheduler Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore: regenerate config docs for workload scheduler Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore: pin catio scheduler test fix Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(http): satisfy scheduler lifecycle clippy Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test: add distributed scheduler toggle coverage Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat: finalize workload scheduler runtime controls Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore: pin merged catio atomic weights Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore: preserve unrelated lockfile resolution Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * perf(runtime): downsample scheduler time accounting Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(runtime): verify cross-runtime scheduler progress Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(runtime): configure scheduler poll sampling Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(runtime): clarify scheduler activation Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(runtime): explain scheduler use case Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Co-authored-by: Ruihang Xia --- .github/scripts/workload-scheduler-e2e.sh | 205 ++++++++++ .github/workflows/integration.yml | 9 +- Cargo.lock | 9 + Cargo.toml | 1 + config/config.md | 10 + config/datanode.example.toml | 12 + config/standalone.example.toml | 12 + src/cmd/src/datanode/builder.rs | 1 - src/cmd/src/standalone.rs | 2 +- src/cmd/tests/load_config_test.rs | 45 +++ src/common/runtime/Cargo.toml | 1 + src/common/runtime/src/global.rs | 436 +++++++++++++++++++-- src/common/runtime/src/lib.rs | 7 +- src/common/runtime/src/metrics.rs | 293 ++++++++++++++ src/common/runtime/src/runtime_default.rs | 4 + src/servers/src/http.rs | 13 + src/servers/src/http/workload_scheduler.rs | 111 ++++++ 17 files changed, 1133 insertions(+), 38 deletions(-) create mode 100644 .github/scripts/workload-scheduler-e2e.sh create mode 100644 src/servers/src/http/workload_scheduler.rs diff --git a/.github/scripts/workload-scheduler-e2e.sh b/.github/scripts/workload-scheduler-e2e.sh new file mode 100644 index 0000000000..173353e8a2 --- /dev/null +++ b/.github/scripts/workload-scheduler-e2e.sh @@ -0,0 +1,205 @@ +#!/usr/bin/env bash +# Copyright 2023 Greptime Team +# +# 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. + +set -euo pipefail + +BIN="${GREPTIMEDB_BIN:-${1:-./bins/greptime}}" +RUN_DIR="${WORKLOAD_SCHEDULER_RUN_DIR:-${TMPDIR:-/tmp}/sqlness-workload-scheduler-e2e}" +TIMEOUT="${WORKLOAD_SCHEDULER_READINESS_TIMEOUT:-60}" +PIDS=() +rm -rf -- "${RUN_DIR}" +mkdir -p "${RUN_DIR}" +SQL_LOG="${RUN_DIR}/sql.log" +STATE_LOG="${RUN_DIR}/scheduler-state.log" +: >"${SQL_LOG}"; : >"${STATE_LOG}" + +cleanup() { + set +e + for pid in "${PIDS[@]}"; do kill -TERM -- "-${pid}" 2>/dev/null || true; done + sleep 1 + for pid in "${PIDS[@]}"; do kill -KILL -- "-${pid}" 2>/dev/null || true; wait "${pid}" 2>/dev/null || true; done +} +trap cleanup EXIT +[[ -x "${BIN}" ]] || { echo "greptime binary is not executable: ${BIN}" >&2; exit 1; } + +# Hold all sockets while selecting them; selections are distinct, but a close-to-bind race remains. +mapfile -t PORTS < <(python3 - <<'PY' +import socket +s=[] +try: + for _ in range(8): + x=socket.socket(); x.bind(("127.0.0.1",0)); s.append(x) + print("\n".join(str(x.getsockname()[1]) for x in s)) +finally: + for x in s: x.close() +PY +) +(( ${#PORTS[@]} == 8 )) || { echo "failed to allocate loopback ports" >&2; exit 1; } +META_RPC="127.0.0.1:${PORTS[0]}"; META_HTTP="http://127.0.0.1:${PORTS[1]}" +DATA_RPC="127.0.0.1:${PORTS[2]}"; DATA_HTTP="http://127.0.0.1:${PORTS[3]}" +FRONT_RPC="127.0.0.1:${PORTS[4]}"; FRONT_HTTP="http://127.0.0.1:${PORTS[5]}" +MYSQL="127.0.0.1:${PORTS[6]}"; POSTGRES="127.0.0.1:${PORTS[7]}" + +start() { + local name="$1"; shift; local dir="${RUN_DIR}/${name}"; mkdir -p "${dir}" + echo "starting ${name}; logs retained in ${dir}" >&2 + setsid "${BIN}" "$@" >"${dir}/stdout.log" 2>"${dir}/stderr.log" & PIDS+=("$!") +} +wait_health() { + local name="$1" url="$2" deadline=$((SECONDS + TIMEOUT)) + until curl -fsS --max-time 2 "${url}/health" >/dev/null; do + (( SECONDS < deadline )) || { echo "timed out waiting for ${name} health" >&2; return 1; } + sleep 1 + done +} +wait_lease() { + local deadline=$((SECONDS + TIMEOUT)) + until curl -fsS --max-time 2 "${META_HTTP}/admin/node-lease" | jq -e 'type == "array" and length > 0' >/dev/null; do + (( SECONDS < deadline )) || { echo "timed out waiting for datanode lease" >&2; return 1; } + sleep 1 + done +} +state() { + local out; out="$(curl -fsS --max-time 10 "${DATA_HTTP}/debug/workload_scheduler")" + printf '%s\n' "${out}" >>"${STATE_LOG}"; printf '%s' "${out}" +} +sql() { + local statement="$1" out + out="$(curl -fsS --max-time 30 --data-urlencode "sql=${statement}" --data-urlencode 'db=public' "${FRONT_HTTP}/v1/sql")" + printf '%s\n%s\n' "${statement}" "${out}" >>"${SQL_LOG}" + jq -e '(.error // null) == null' <<<"${out}" >/dev/null || { echo "SQL failed: ${statement}" >&2; return 1; } + printf '%s' "${out}" +} +assert_status() { + local status="$1" enabled="$2" query_weight="$3" write_weight="$4" + jq -e --argjson enabled "${enabled}" --argjson query_weight "${query_weight}" \ + --argjson write_weight "${write_weight}" \ + '.enabled == $enabled and .query.weight == $query_weight and .write.weight == $write_weight' \ + <<<"${status}" >/dev/null || { echo "unexpected scheduler status: ${status}" >&2; return 1; } +} +assert_metrics() { + local metrics + metrics="$(curl -fsS --max-time 10 "${DATA_HTTP}/metrics")" + grep -Eq '^greptime_workload_scheduler_enabled 1$' <<<"${metrics}" + for class in query write; do + expected_weight=3 + [[ "${class}" == write ]] && expected_weight=7 + grep -Eq "^greptime_workload_scheduler_weight\\{workload=\\\"${class}\\\"\\} ${expected_weight}$" <<<"${metrics}" + grep -Eq "^greptime_workload_scheduler_polls_total\\{workload=\\\"${class}\\\"\\} [0-9.e+-]+$" <<<"${metrics}" + done +} +metric_polls() { + local metrics + metrics="$(curl -fsS --max-time 10 "${DATA_HTTP}/metrics")" + awk ' + $1 == "greptime_workload_scheduler_polls_total{workload=\"query\"}" { + workload = "query" + count[workload]++ + value[workload] = $2 + if (NF != 2 || $2 !~ /^[+-]?([0-9]+([.][0-9]*)?|[.][0-9]+)([eE][+-]?[0-9]+)?$/) { + invalid[workload] = $2 + } + } + $1 == "greptime_workload_scheduler_polls_total{workload=\"write\"}" { + workload = "write" + count[workload]++ + value[workload] = $2 + if (NF != 2 || $2 !~ /^[+-]?([0-9]+([.][0-9]*)?|[.][0-9]+)([eE][+-]?[0-9]+)?$/) { + invalid[workload] = $2 + } + } + END { + failed = 0 + for (i = 1; i <= 2; i++) { + workload = (i == 1 ? "query" : "write") + if (count[workload] == 0) { + printf "metric_polls: missing %s Prometheus sample\n", workload > "/dev/stderr" + failed = 1 + } else if (count[workload] != 1) { + printf "metric_polls: expected exactly one %s Prometheus sample, found %d\n", workload, count[workload] > "/dev/stderr" + failed = 1 + } + if (workload in invalid) { + printf "metric_polls: invalid %s Prometheus numeric token: %s\n", workload, invalid[workload] > "/dev/stderr" + failed = 1 + } + } + if (failed) exit 1 + print value["query"] + print value["write"] + } + ' <<<"${metrics}" +} +assert_metric_polls() { + local before="$1" after="$2" phase="$3" comparison="$4" + local -a before_values after_values + mapfile -t before_values <<<"${before}" + mapfile -t after_values <<<"${after}" + for i in 0 1; do + local class=query + [[ "${i}" == 1 ]] && class=write + awk -v before="${before_values[${i}]}" -v after="${after_values[${i}]}" \ + "BEGIN { exit !(after ${comparison} before) }" || { + echo "${phase}: ${class} Prometheus polls did not satisfy ${comparison}" >&2; return 1; + } + done +} +assert_values() { + local response="$1" expected="$2" + jq -e --argjson expected "${expected}" '[.output[0].records.rows[][0]] == $expected' <<<"${response}" >/dev/null +} + +mkdir -p "${RUN_DIR}/datanode" +cat >"${RUN_DIR}/datanode/datanode.toml" </dev/null +assert_status "$(state)" true 3 7 +assert_metrics +sql 'CREATE TABLE scheduler_e2e (ts TIMESTAMP TIME INDEX, v INT)' >/dev/null +before="$(metric_polls)"; sql 'INSERT INTO scheduler_e2e VALUES (1000, 10), (2000, 20)' >/dev/null +assert_values "$(sql 'SELECT v FROM scheduler_e2e ORDER BY v')" '[10,20]'; after="$(metric_polls)" +assert_status "$(state)" true 3 7; assert_metric_polls "${before}" "${after}" 'enabled phase' '>' + +curl -fsS --max-time 10 -X POST -H 'Content-Type: application/json' --data false "${DATA_HTTP}/debug/workload_scheduler/enabled" >/dev/null +before="$(metric_polls)"; assert_status "$(state)" false 3 7 +sql 'INSERT INTO scheduler_e2e VALUES (3000, 30)' >/dev/null +assert_values "$(sql 'SELECT v FROM scheduler_e2e ORDER BY v')" '[10,20,30]'; after="$(metric_polls)" +assert_status "$(state)" false 3 7; assert_metric_polls "${before}" "${after}" 'disabled phase' '==' + +echo "distributed workload scheduler E2E passed; logs retained in ${RUN_DIR}" diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index b8868123a4..d7a91cadd4 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -151,12 +151,19 @@ jobs: run: tar -xvf ./bins.tar.gz - name: Run sqlness run: RUST_BACKTRACE=1 ./bins/sqlness-runner bare ${{ matrix.mode.opts }} -c ./tests/cases --bins-dir ./bins --preserve-state + - if: matrix.mode.name == 'Basic' + name: Run distributed workload scheduler E2E + env: + GREPTIMEDB_BIN: ${{ github.workspace }}/bins/greptime + WORKLOAD_SCHEDULER_RUN_DIR: /tmp/sqlness-workload-scheduler-e2e + run: bash ./.github/scripts/workload-scheduler-e2e.sh - name: Upload sqlness logs if: failure() uses: actions/upload-artifact@v4 with: name: sqlness-logs-${{ matrix.mode.name }} - path: /tmp/sqlness* + path: | + /tmp/sqlness* retention-days: 3 export-import-v2-e2e: diff --git a/Cargo.lock b/Cargo.lock index ff2ce93d77..5efd0bda67 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1712,6 +1712,14 @@ dependencies = [ "tokio-stream", ] +[[package]] +name = "catio" +version = "0.1.0" +source = "git+https://github.com/GreptimeTeam/catio.git?rev=c75be64e7a58f620f6340c24b5b8d76f131af6aa#c75be64e7a58f620f6340c24b5b8d76f131af6aa" +dependencies = [ + "tokio", +] + [[package]] name = "cbc" version = "0.1.2" @@ -2821,6 +2829,7 @@ name = "common-runtime" version = "1.3.0" dependencies = [ "async-trait", + "catio", "clap", "common-error", "common-macro", diff --git a/Cargo.toml b/Cargo.toml index fab7b2aa73..330b6f0640 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -282,6 +282,7 @@ api = { path = "src/api" } auth = { path = "src/auth" } cache = { path = "src/cache" } catalog = { path = "src/catalog" } +catio = { git = "https://github.com/GreptimeTeam/catio.git", rev = "c75be64e7a58f620f6340c24b5b8d76f131af6aa" } cli = { path = "src/cli" } client = { path = "src/client" } cmd = { path = "src/cmd", default-features = false } diff --git a/config/config.md b/config/config.md index 52362b55c4..a1c91a8af8 100644 --- a/config/config.md +++ b/config/config.md @@ -26,6 +26,11 @@ | `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | | `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.experimental_workload_scheduler` | -- | -- | Experimental weighted, work-conserving query/write task scheduler. | +| `runtime.experimental_workload_scheduler.enable` | Bool | `false` | Enable when concurrent queries and writes interfere with each other—for example, when long-running queries increase ingestion latency.
The weights set their relative runtime shares while both are backlogged. Disabled by default. | +| `runtime.experimental_workload_scheduler.query_weight` | Integer | `2` | Relative query share while both query and write workloads are backlogged. | +| `runtime.experimental_workload_scheduler.write_weight` | Integer | `8` | Relative write share while both query and write workloads are backlogged. | +| `runtime.experimental_workload_scheduler.sample_every_polls` | Integer | `16` | Number of polls between scheduler fairness samples. Must be greater than zero. | | `http` | -- | -- | The HTTP server options. | | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. | | `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the
`prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value. | @@ -526,6 +531,11 @@ | `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `runtime.query_rt_size` | Integer | `7` | The number of threads to execute datanode query operations.
Defaults to max(num_cpus - 1, 2). | | `runtime.ingest_rt_size` | Integer | `8` | The number of threads to execute datanode ingestion operations. | +| `runtime.experimental_workload_scheduler` | -- | -- | Experimental weighted, work-conserving query/write task scheduler. | +| `runtime.experimental_workload_scheduler.enable` | Bool | `false` | Enable when concurrent queries and writes interfere with each other—for example, when long-running queries increase ingestion latency.
The weights set their relative runtime shares while both are backlogged. Disabled by default. | +| `runtime.experimental_workload_scheduler.query_weight` | Integer | `2` | Relative query share while both query and write workloads are backlogged. | +| `runtime.experimental_workload_scheduler.write_weight` | Integer | `8` | Relative write share while both query and write workloads are backlogged. | +| `runtime.experimental_workload_scheduler.sample_every_polls` | Integer | `16` | Number of polls between scheduler fairness samples. Must be greater than zero. | | `meta_client` | -- | -- | The metasrv client options. | | `meta_client.metasrv_addrs` | Array | -- | The addresses of the metasrv. | | `meta_client.timeout` | String | `3s` | Operation timeout. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 4846ebcea1..a99c817ae2 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -92,6 +92,18 @@ watch = false ## The number of threads to execute datanode ingestion operations. #+ ingest_rt_size = 8 +## Experimental weighted, work-conserving query/write task scheduler. +#+ [runtime.experimental_workload_scheduler] +## Enable when concurrent queries and writes interfere with each other—for example, when long-running queries increase ingestion latency. +## The weights set their relative runtime shares while both are backlogged. Disabled by default. +#+ enable = false +## Relative query share while both query and write workloads are backlogged. +#+ query_weight = 2 +## Relative write share while both query and write workloads are backlogged. +#+ write_weight = 8 +## Number of polls between scheduler fairness samples. Must be greater than zero. +#+ sample_every_polls = 16 + ## The metasrv client options. [meta_client] ## The addresses of the metasrv. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index b08251b259..1026ddc925 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -60,6 +60,18 @@ max_concurrent_queries = 0 ## Defaults to max(num_cpus / 2, 2). #+ compact_rt_max_blocking_threads = 4 +## Experimental weighted, work-conserving query/write task scheduler. +#+ [runtime.experimental_workload_scheduler] +## Enable when concurrent queries and writes interfere with each other—for example, when long-running queries increase ingestion latency. +## The weights set their relative runtime shares while both are backlogged. Disabled by default. +#+ enable = false +## Relative query share while both query and write workloads are backlogged. +#+ query_weight = 2 +## Relative write share while both query and write workloads are backlogged. +#+ write_weight = 8 +## Number of polls between scheduler fairness samples. Must be greater than zero. +#+ sample_every_polls = 16 + ## The HTTP server options. [http] ## The address to bind the HTTP server. diff --git a/src/cmd/src/datanode/builder.rs b/src/cmd/src/datanode/builder.rs index 3dfa03e52e..810dfeedc7 100644 --- a/src/cmd/src/datanode/builder.rs +++ b/src/cmd/src/datanode/builder.rs @@ -70,7 +70,6 @@ impl InstanceBuilder { None, ); - common_runtime::init_global_runtimes(&opts.runtime); common_runtime::init_datanode_runtimes(&opts.runtime); crate::options::flush_dropped_plugin_warnings(); diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index d628461ba7..6c5ac25f64 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -429,7 +429,7 @@ impl StartCommand { Some(&opts.component.slow_query), ); - common_runtime::init_global_runtimes(&opts.runtime); + common_runtime::init_standalone_runtimes(&opts.runtime); crate::options::flush_dropped_plugin_warnings(); log_versions(verbose_version(), short_version(), APP_NAME); diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 8049e8b8ab..67e2e573f3 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::io::Write; +use std::num::{NonZeroU32, NonZeroUsize}; use std::time::Duration; use cmd::options::GreptimeOptions; @@ -50,6 +51,12 @@ fn test_load_datanode_runtime_options_from_runtime_section() { compact_rt_max_blocking_threads = 6 ingest_rt_size = 8 query_rt_size = 7 + + [runtime.experimental_workload_scheduler] + enable = true + query_weight = 1 + write_weight = 4 + sample_every_polls = 32 "#; let options: GreptimeOptions = toml::from_str(toml).unwrap(); @@ -59,6 +66,44 @@ fn test_load_datanode_runtime_options_from_runtime_section() { assert_eq!(6, options.runtime.compact_rt_max_blocking_threads); assert_eq!(8, options.runtime.ingest_rt_size); assert_eq!(7, options.runtime.query_rt_size); + assert!(options.runtime.experimental_workload_scheduler.enable); + assert_eq!( + NonZeroU32::new(1).unwrap(), + options.runtime.experimental_workload_scheduler.query_weight + ); + assert_eq!( + NonZeroU32::new(4).unwrap(), + options.runtime.experimental_workload_scheduler.write_weight + ); + assert_eq!( + NonZeroUsize::new(32).unwrap(), + options + .runtime + .experimental_workload_scheduler + .sample_every_polls + ); +} + +#[test] +fn test_load_runtime_options_rejects_zero_scheduler_sampling() { + let toml = r#" + [runtime.experimental_workload_scheduler] + sample_every_polls = 0 + "#; + + let result = toml::from_str::>(toml); + assert!(result.is_err()); +} + +#[test] +fn test_load_runtime_options_rejects_zero_scheduler_weight() { + let toml = r#" + [runtime.experimental_workload_scheduler] + query_weight = 0 + "#; + + let result = toml::from_str::>(toml); + assert!(result.is_err()); } #[test] diff --git a/src/common/runtime/Cargo.toml b/src/common/runtime/Cargo.toml index 0b5ba5f096..927aa17b9c 100644 --- a/src/common/runtime/Cargo.toml +++ b/src/common/runtime/Cargo.toml @@ -16,6 +16,7 @@ workspace = true [dependencies] async-trait.workspace = true +catio.workspace = true clap.workspace = true common-error.workspace = true common-macro.workspace = true diff --git a/src/common/runtime/src/global.rs b/src/common/runtime/src/global.rs index 827832fca2..ec804a7f28 100644 --- a/src/common/runtime/src/global.rs +++ b/src/common/runtime/src/global.rs @@ -13,14 +13,19 @@ // limitations under the License. //! Global runtimes +use std::collections::BTreeMap; use std::future::Future; +use std::num::{NonZeroU32, NonZeroUsize}; use std::sync::{Mutex, Once}; -use common_telemetry::info; +use catio::{Scheduler, SchedulerStats, TaskClass}; +use common_telemetry::{info, warn}; use once_cell::sync::Lazy; use paste::paste; use serde::{Deserialize, Serialize}; +use tokio::runtime::Handle; +use crate::metrics::register_workload_scheduler_metrics; use crate::runtime::{BuilderBuild, RuntimeTrait}; use crate::{Builder, JoinHandle, Runtime}; @@ -30,6 +35,34 @@ const HB_WORKERS: usize = 2; /// The minimum number of worker threads for runtimes sized by CPU count. /// A single-threaded runtime can easily deadlock in async code. const MIN_RUNTIME_THREADS: usize = 2; +pub(crate) const QUERY_TASK_CLASS: TaskClass = TaskClass::new(1); +pub(crate) const WRITE_TASK_CLASS: TaskClass = TaskClass::new(2); + +/// Experimental options for sharing Tokio capacity between query and write +/// workloads. +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] +#[serde(default)] +pub struct WorkloadSchedulerOptions { + /// Enables policy-controlled query and write task spawning. + pub enable: bool, + /// Relative share for query polls while writes are also backlogged. + pub query_weight: NonZeroU32, + /// Relative share for write polls while queries are also backlogged. + pub write_weight: NonZeroU32, + /// Number of polls between scheduler fairness samples. + pub sample_every_polls: NonZeroUsize, +} + +impl Default for WorkloadSchedulerOptions { + fn default() -> Self { + Self { + enable: false, + query_weight: NonZeroU32::new(2).unwrap(), + write_weight: NonZeroU32::new(8).unwrap(), + sample_every_polls: NonZeroUsize::new(16).unwrap(), + } + } +} /// The options for the global runtimes. #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] @@ -45,6 +78,8 @@ pub struct RuntimeOptions { pub query_rt_size: usize, /// The number of threads to execute datanode ingestion operations. pub ingest_rt_size: usize, + /// Experimental weighted scheduler for query and write workloads. + pub experimental_workload_scheduler: WorkloadSchedulerOptions, } impl RuntimeOptions { @@ -56,6 +91,7 @@ impl RuntimeOptions { compact_rt_max_blocking_threads: usize::max(cpus / 2, MIN_RUNTIME_THREADS), query_rt_size: usize::max(cpus.saturating_sub(1), MIN_RUNTIME_THREADS), ingest_rt_size: cpus, + experimental_workload_scheduler: WorkloadSchedulerOptions::default(), } } } @@ -103,6 +139,9 @@ struct GlobalRuntimes { hb_runtime: Runtime, query_runtime: Runtime, ingest_runtime: Runtime, + query_handle: Handle, + ingest_handle: Handle, + workload_scheduler: Option, } macro_rules! define_spawn { @@ -132,12 +171,45 @@ macro_rules! define_spawn { }; } +macro_rules! define_scheduled_spawn { + ($type: ident, $class: ident) => { + paste! { + fn [](&self, future: F) -> JoinHandle + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + match &self.workload_scheduler { + Some(scheduler) => scheduler.spawn_in_on( + &self.[<$type _handle>], + $class, + future, + ), + None => self.[<$type _runtime>].spawn(future), + } + } + + fn [](&self, future: F) -> JoinHandle + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + self.[<$type _runtime>].spawn_blocking(future) + } + + fn [](&self, future: F) -> F::Output { + self.[<$type _runtime>].block_on(future) + } + } + }; +} + impl GlobalRuntimes { define_spawn!(global); define_spawn!(compact); define_spawn!(hb); - define_spawn!(query); - define_spawn!(ingest); + define_scheduled_spawn!(query, QUERY_TASK_CLASS); + define_scheduled_spawn!(ingest, WRITE_TASK_CLASS); fn new( global: Option, @@ -145,12 +217,14 @@ impl GlobalRuntimes { heartbeat: Option, query: Option, ingest: Option, + workload_scheduler: Option, ) -> Self { let global_runtime = global.unwrap_or_else(|| create_runtime("global", "global-worker", GLOBAL_WORKERS)); let query_runtime = query.unwrap_or_else(|| global_runtime.clone()); let ingest_runtime = ingest.unwrap_or_else(|| global_runtime.clone()); - + let query_handle = query_runtime.handle(); + let ingest_handle = ingest_runtime.handle(); Self { global_runtime, compact_runtime: compact.unwrap_or_else(|| { @@ -167,6 +241,9 @@ impl GlobalRuntimes { .unwrap_or_else(|| create_runtime("heartbeat", "hb-worker", HB_WORKERS)), query_runtime, ingest_runtime, + query_handle, + ingest_handle, + workload_scheduler, } } } @@ -178,6 +255,7 @@ struct ConfigRuntimes { hb_runtime: Option, query_runtime: Option, ingest_runtime: Option, + workload_scheduler: Option, already_init: bool, } @@ -188,49 +266,70 @@ static GLOBAL_RUNTIMES: Lazy = Lazy::new(|| { let heartbeat = c.hb_runtime.take(); let query = c.query_runtime.take(); let ingest = c.ingest_runtime.take(); + let workload_scheduler = c.workload_scheduler.take(); c.already_init = true; - GlobalRuntimes::new(global, compact, heartbeat, query, ingest) + GlobalRuntimes::new( + global, + compact, + heartbeat, + query, + ingest, + workload_scheduler, + ) }); static CONFIG_RUNTIMES: Lazy> = Lazy::new(|| Mutex::new(ConfigRuntimes::default())); +static START: Once = Once::new(); -/// Initialize the global runtimes +/// Initialize runtimes for frontend, metasrv, and flownode processes. /// -/// # Panics -/// Panics when the global runtimes are already initialized. -/// You should call this function before using any runtime functions. +/// Query and ingest work share the global runtime and no workload scheduler is +/// constructed for these process roles. pub fn init_global_runtimes(options: &RuntimeOptions) { - static START: Once = Once::new(); - START.call_once(move || { + START.call_once(|| { let mut c = CONFIG_RUNTIMES.lock().unwrap(); assert!(!c.already_init, "Global runtimes already initialized"); - c.global_runtime = Some(create_runtime( - "global", - "global-worker", - options.global_rt_size, - )); - c.compact_runtime = Some(create_compact_runtime( - "compact", - "compact-worker", - options.compact_rt_size, - options.compact_rt_max_blocking_threads, - )); - c.hb_runtime = Some(create_runtime("heartbeat", "hb-worker", HB_WORKERS)); + init_common_runtimes(&mut c, options); + c.already_init = true; }); } -/// Initialize the datanode-specific global runtimes. +/// Initialize runtimes for a standalone process. /// -/// # Panics -/// Panics when the global runtimes are already initialized. -/// You should call this function before using any runtime functions. -pub fn init_datanode_runtimes(options: &RuntimeOptions) { - static START: Once = Once::new(); - START.call_once(move || { +/// Query and ingest work share the global runtime. The scheduler is always +/// constructed, with the global runtime size as its internal bound, and starts +/// enabled according to configuration. +pub fn init_standalone_runtimes(options: &RuntimeOptions) { + START.call_once(|| { let mut c = CONFIG_RUNTIMES.lock().unwrap(); assert!(!c.already_init, "Global runtimes already initialized"); + init_common_runtimes(&mut c, options); + c.workload_scheduler = Some(create_workload_scheduler(options, options.global_rt_size)); + c.already_init = true; + }); +} + +/// Initialize runtimes for a datanode process. +/// +/// Query and ingest use dedicated runtimes. The scheduler is always +/// constructed with their checked combined size as its internal bound, and +/// starts enabled according to configuration. +/// +/// # Panics +/// +/// Panics if the configured query and ingest runtime sizes overflow `usize` +/// when combined. +pub fn init_datanode_runtimes(options: &RuntimeOptions) { + let capacity = options + .query_rt_size + .checked_add(options.ingest_rt_size) + .expect("datanode workload scheduler runtime capacity overflowed usize"); + START.call_once(|| { + let mut c = CONFIG_RUNTIMES.lock().unwrap(); + assert!(!c.already_init, "Global runtimes already initialized"); + init_common_runtimes(&mut c, options); c.query_runtime = Some(create_runtime( "query", "query-worker", @@ -241,9 +340,53 @@ pub fn init_datanode_runtimes(options: &RuntimeOptions) { "ingest-worker", options.ingest_rt_size, )); + c.workload_scheduler = Some(create_workload_scheduler(options, capacity)); + c.already_init = true; }); } +fn init_common_runtimes(c: &mut ConfigRuntimes, options: &RuntimeOptions) { + c.global_runtime = Some(create_runtime( + "global", + "global-worker", + options.global_rt_size, + )); + c.compact_runtime = Some(create_compact_runtime( + "compact", + "compact-worker", + options.compact_rt_size, + options.compact_rt_max_blocking_threads, + )); + c.hb_runtime = Some(create_runtime("heartbeat", "hb-worker", HB_WORKERS)); +} + +fn create_workload_scheduler(options: &RuntimeOptions, capacity: usize) -> Scheduler { + assert!( + capacity > 0, + "experimental workload scheduler capacity must be greater than zero" + ); + let scheduler_options = &options.experimental_workload_scheduler; + let scheduler = Scheduler::builder() + // This is deliberately an internal scheduler bound, not public config. + .max_concurrent_polls(capacity) + .sample_every_polls(scheduler_options.sample_every_polls.get()) + .weight(QUERY_TASK_CLASS, scheduler_options.query_weight.get()) + .weight(WRITE_TASK_CLASS, scheduler_options.write_weight.get()) + .build(); + scheduler.set_enabled(scheduler_options.enable); + register_workload_scheduler_metrics(scheduler.clone()); + info!( + "Constructed the experimental workload scheduler: internal_capacity={}, \ + query_weight={}, write_weight={}, sample_every_polls={}, enabled={}", + capacity, + scheduler_options.query_weight, + scheduler_options.write_weight, + scheduler_options.sample_every_polls, + scheduler_options.enable + ); + scheduler +} + macro_rules! define_global_runtime_spawn { ($type: ident) => { paste! { @@ -284,15 +427,104 @@ define_global_runtime_spawn!(hb); define_global_runtime_spawn!(query); define_global_runtime_spawn!(ingest); +/// Returns whether the experimental workload scheduler is currently enabled. +/// Returns `false` when no scheduler was constructed at startup. +pub fn workload_scheduler_enabled() -> bool { + GLOBAL_RUNTIMES + .workload_scheduler + .as_ref() + .is_some_and(Scheduler::is_enabled) +} + +/// Enables or disables the experimental workload scheduler for new spawn +/// submissions. Returns `false` when no scheduler was constructed at startup. +pub fn set_workload_scheduler_enabled(enabled: bool) -> bool { + let Some(scheduler) = GLOBAL_RUNTIMES.workload_scheduler.as_ref() else { + warn!( + "The experimental workload scheduler was not constructed at startup; ignoring enabled={enabled}" + ); + return false; + }; + + scheduler.set_enabled(enabled); + info!("Experimental workload scheduler enabled={enabled}"); + true +} + +/// Sets the query and write weights atomically. Returns `false` when no +/// scheduler was constructed at startup. +pub fn set_workload_scheduler_weights(query: NonZeroU32, write: NonZeroU32) -> bool { + let Some(scheduler) = GLOBAL_RUNTIMES.workload_scheduler.as_ref() else { + warn!( + "The experimental workload scheduler was not constructed at startup; ignoring query_weight={query}, write_weight={write}" + ); + return false; + }; + + let weights = BTreeMap::from([(QUERY_TASK_CLASS, query), (WRITE_TASK_CLASS, write)]); + scheduler.set_weights(&weights); + info!("Experimental workload scheduler weights query={query}, write={write}"); + true +} + +/// Returns scheduler counters when the experimental workload scheduler was +/// constructed at startup, including while it is dynamically disabled. +pub fn workload_scheduler_stats() -> Option { + GLOBAL_RUNTIMES + .workload_scheduler + .as_ref() + .map(Scheduler::stats) +} + #[cfg(test)] mod tests { - use std::sync::mpsc; - use std::time::Duration; + use std::future::Future; + use std::pin::Pin; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::{Arc, mpsc}; + use std::task::{Context, Poll}; + use std::time::{Duration, Instant}; use tokio_test::assert_ok; use super::*; + struct CooperativePolls { + stop: Arc, + } + + impl Future for CooperativePolls { + type Output = (); + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let deadline = Instant::now() + Duration::from_micros(100); + while Instant::now() < deadline { + std::hint::spin_loop(); + } + + if self.stop.load(Ordering::Relaxed) { + Poll::Ready(()) + } else { + cx.waker().wake_by_ref(); + Poll::Pending + } + } + } + + fn wait_until(description: &str, condition: F) + where + F: Fn() -> bool, + { + let deadline = Instant::now() + Duration::from_secs(5); + while !condition() { + assert!( + Instant::now() < deadline, + "timed out waiting for {description}" + ); + std::thread::sleep(Duration::from_millis(1)); + } + } + #[test] fn test_datanode_runtime_options_default() { let options = RuntimeOptions::default(); @@ -312,6 +544,10 @@ mod tests { options.query_rt_size ); assert_eq!(cpus, options.ingest_rt_size); + assert_eq!( + WorkloadSchedulerOptions::default(), + options.experimental_workload_scheduler + ); } #[test] @@ -354,6 +590,7 @@ mod tests { None, None, None, + None, ); assert_eq!("test-global", runtimes.global_runtime.name()); @@ -400,6 +637,141 @@ mod tests { }); } + #[test] + fn test_workload_scheduler_builds_with_initial_enabled_state() { + let mut options = RuntimeOptions::default(); + options.experimental_workload_scheduler.enable = false; + options.experimental_workload_scheduler.sample_every_polls = NonZeroUsize::new(7).unwrap(); + let scheduler = create_workload_scheduler(&options, options.global_rt_size); + assert_eq!(7, scheduler.stats().sample_every_polls); + assert!(!scheduler.is_enabled()); + + scheduler.set_enabled(true); + assert!(scheduler.is_enabled()); + } + + #[test] + fn test_workload_scheduler_bypasses_disabled_query_and_write_spawns() { + let runtime = create_runtime("test-workload-bypass", "test-workload-bypass-worker", 2); + let scheduler = Scheduler::builder() + .max_concurrent_polls(2) + .weight(QUERY_TASK_CLASS, 2) + .weight(WRITE_TASK_CLASS, 8) + .build(); + scheduler.set_enabled(false); + let runtimes = GlobalRuntimes::new( + Some(runtime.clone()), + Some(runtime.clone()), + Some(runtime.clone()), + Some(runtime.clone()), + Some(runtime.clone()), + Some(scheduler.clone()), + ); + + let query = runtimes.spawn_query(async { "query" }); + let write = runtimes.spawn_ingest(async { "write" }); + let (query, write) = + runtime.block_on(async { (query.await.unwrap(), write.await.unwrap()) }); + + assert_eq!("query", query); + assert_eq!("write", write); + let stats = scheduler.stats(); + for class in [QUERY_TASK_CLASS, WRITE_TASK_CLASS] { + let class_stats = &stats.classes[&class]; + assert_eq!(0, class_stats.tasks); + assert_eq!(0, class_stats.admitted); + assert_eq!(0, class_stats.polls); + } + } + + #[test] + fn test_workload_scheduler_wraps_query_and_write_spawns() { + let runtime = create_runtime("test-workload", "test-workload-worker", 2); + let scheduler = Scheduler::builder() + .max_concurrent_polls(2) + .weight(QUERY_TASK_CLASS, 2) + .weight(WRITE_TASK_CLASS, 8) + .build(); + let runtimes = GlobalRuntimes::new( + Some(runtime.clone()), + Some(runtime.clone()), + Some(runtime.clone()), + Some(runtime.clone()), + Some(runtime.clone()), + Some(scheduler.clone()), + ); + + let query = runtimes.spawn_query(async { "query" }); + let write = runtimes.spawn_ingest(async { "write" }); + let (query, write) = + runtime.block_on(async { (query.await.unwrap(), write.await.unwrap()) }); + + assert_eq!("query", query); + assert_eq!("write", write); + let stats = scheduler.stats(); + assert_eq!(1, stats.classes[&QUERY_TASK_CLASS].polls); + assert_eq!(1, stats.classes[&WRITE_TASK_CLASS].polls); + } + + #[test] + fn test_datanode_query_backlog_does_not_starve_ingest() { + let query_runtime = create_runtime("test-datanode-query", "test-query-worker", 1); + let ingest_runtime = create_runtime("test-datanode-ingest", "test-ingest-worker", 1); + let scheduler = Scheduler::builder() + .max_concurrent_polls(2) + .sample_every_polls(16) + .weight(QUERY_TASK_CLASS, 2) + .weight(WRITE_TASK_CLASS, 8) + .build(); + let runtimes = GlobalRuntimes::new( + Some(query_runtime.clone()), + Some(query_runtime.clone()), + Some(query_runtime.clone()), + Some(query_runtime), + Some(ingest_runtime), + Some(scheduler.clone()), + ); + + let stop_queries = Arc::new(AtomicBool::new(false)); + let query_tasks = (0..3) + .map(|_| { + runtimes.spawn_query(CooperativePolls { + stop: stop_queries.clone(), + }) + }) + .collect::>(); + + // Establish the actual datanode topology: two query polls occupy the + // scheduler's capacity while the third self-waking query is queued. + wait_until("two active query polls and a queued query", || { + let stats = scheduler.stats(); + stats.active_polls == 2 + && stats + .classes + .get(&QUERY_TASK_CLASS) + .is_some_and(|class| class.tasks == 3 && class.queued >= 1) + }); + + let write = runtimes.spawn_ingest(async {}); + wait_until("write body", || write.is_finished()); + + stop_queries.store(true, Ordering::Relaxed); + for query in &query_tasks { + query.abort(); + } + runtimes.block_on_query(async { + for query in query_tasks { + let _ = query.await; + } + }); + runtimes.block_on_ingest(async { + write.await.unwrap(); + }); + wait_until("scheduler polls to drain", || { + scheduler.stats().active_polls == 0 + }); + } + #[test] fn test_datanode_runtime_spawn_block_on() { let handle = spawn_query(async { 1 + 1 }); diff --git a/src/common/runtime/src/lib.rs b/src/common/runtime/src/lib.rs index 533fb57c6a..998b73197a 100644 --- a/src/common/runtime/src/lib.rs +++ b/src/common/runtime/src/lib.rs @@ -23,9 +23,10 @@ pub mod runtime_throttleable; pub use global::{ block_on_compact, block_on_global, block_on_ingest, block_on_query, compact_runtime, create_runtime, global_runtime, ingest_runtime, init_datanode_runtimes, init_global_runtimes, - query_runtime, spawn_blocking_compact, spawn_blocking_global, spawn_blocking_hb, - spawn_blocking_ingest, spawn_blocking_query, spawn_compact, spawn_global, spawn_hb, - spawn_ingest, spawn_query, + init_standalone_runtimes, query_runtime, set_workload_scheduler_enabled, + set_workload_scheduler_weights, spawn_blocking_compact, spawn_blocking_global, + spawn_blocking_hb, spawn_blocking_ingest, spawn_blocking_query, spawn_compact, spawn_global, + spawn_hb, spawn_ingest, spawn_query, workload_scheduler_enabled, workload_scheduler_stats, }; pub use crate::repeated_task::{BoxedTaskFunction, RepeatedTask, TaskFunction}; diff --git a/src/common/runtime/src/metrics.rs b/src/common/runtime/src/metrics.rs index c332ccf6e6..759f8bbe91 100644 --- a/src/common/runtime/src/metrics.rs +++ b/src/common/runtime/src/metrics.rs @@ -13,9 +13,18 @@ // limitations under the License. //! Runtime metrics +use std::collections::BTreeMap; +use std::sync::Mutex; +use std::time::Duration; + +use catio::Scheduler; use lazy_static::lazy_static; +use prometheus::core::{Collector, Desc}; +use prometheus::proto::MetricFamily; use prometheus::*; +use crate::global::{QUERY_TASK_CLASS, WRITE_TASK_CLASS}; + pub const THREAD_NAME_LABEL: &str = "thread_name"; lazy_static! { @@ -32,3 +41,287 @@ lazy_static! { ) .unwrap(); } + +#[derive(Clone, Default)] +struct ClassSnapshot { + polls: u64, + total_admission_wait: Duration, +} + +struct WorkloadSchedulerCollector { + scheduler: Scheduler, + enabled: IntGauge, + active: IntGauge, + weight: IntGaugeVec, + queued: IntGaugeVec, + polls: IntCounterVec, + total_admission_wait: CounterVec, + snapshots: Mutex>, +} + +impl WorkloadSchedulerCollector { + fn new(scheduler: Scheduler) -> Self { + let workload_label = &["workload"]; + Self { + scheduler, + enabled: IntGauge::new( + "greptime_workload_scheduler_enabled", + "Whether the workload scheduler is enabled", + ) + .unwrap(), + active: IntGauge::new( + "greptime_workload_scheduler_active_polls", + "Task polls admitted to Tokio but not yet completed", + ) + .unwrap(), + weight: IntGaugeVec::new( + Opts::new( + "greptime_workload_scheduler_weight", + "Configured workload scheduler weight", + ), + workload_label, + ) + .unwrap(), + queued: IntGaugeVec::new( + Opts::new( + "greptime_workload_scheduler_queued_tasks", + "Tasks queued in the workload scheduler", + ), + workload_label, + ) + .unwrap(), + polls: IntCounterVec::new( + Opts::new( + "greptime_workload_scheduler_polls_total", + "Cumulative task polls admitted by the workload scheduler", + ), + workload_label, + ) + .unwrap(), + total_admission_wait: CounterVec::new( + Opts::new( + "greptime_workload_scheduler_admission_wait_seconds_total", + "Cumulative workload scheduler admission wait time in seconds", + ), + workload_label, + ) + .unwrap(), + snapshots: Mutex::new(BTreeMap::new()), + } + } + + fn update_locked(&self, snapshots: &mut BTreeMap<&'static str, ClassSnapshot>) { + let stats = self.scheduler.stats(); + self.enabled.set(i64::from(self.scheduler.is_enabled())); + self.active.set( + stats + .active_polls + .min(i64::MAX as usize) + .try_into() + .unwrap_or(i64::MAX), + ); + + for (class, workload) in [(QUERY_TASK_CLASS, "query"), (WRITE_TASK_CLASS, "write")] { + let class_stats = stats.classes.get(&class).cloned().unwrap_or_default(); + let labels = &[workload]; + self.weight + .with_label_values(labels) + .set(i64::from(class_stats.weight)); + self.queued.with_label_values(labels).set( + class_stats + .queued + .min(i64::MAX as usize) + .try_into() + .unwrap_or(i64::MAX), + ); + let previous = snapshots.entry(workload).or_default(); + self.polls + .with_label_values(labels) + .inc_by(class_stats.polls.saturating_sub(previous.polls)); + self.total_admission_wait.with_label_values(labels).inc_by( + class_stats + .total_admission_wait + .saturating_sub(previous.total_admission_wait) + .as_secs_f64(), + ); + *previous = ClassSnapshot { + polls: class_stats.polls, + total_admission_wait: class_stats.total_admission_wait, + }; + } + } +} + +impl Collector for WorkloadSchedulerCollector { + fn desc(&self) -> Vec<&Desc> { + let mut desc = self.enabled.desc(); + desc.extend(self.active.desc()); + desc.extend(self.weight.desc()); + desc.extend(self.queued.desc()); + desc.extend(self.polls.desc()); + desc.extend(self.total_admission_wait.desc()); + desc + } + + fn collect(&self) -> Vec { + let mut snapshots = self.snapshots.lock().unwrap(); + self.update_locked(&mut snapshots); + let mut families = self.enabled.collect(); + families.extend(self.active.collect()); + families.extend(self.weight.collect()); + families.extend(self.queued.collect()); + families.extend(self.polls.collect()); + families.extend(self.total_admission_wait.collect()); + families + } +} + +pub(crate) fn register_workload_scheduler_metrics(scheduler: Scheduler) { + register(Box::new(WorkloadSchedulerCollector::new(scheduler))) + .expect("workload scheduler metrics collector registration must succeed"); +} + +#[cfg(test)] +mod tests { + use std::collections::{BTreeMap, BTreeSet}; + use std::time::Duration; + + use prometheus::proto::{MetricFamily, MetricType}; + + use super::*; + + fn family<'a>(families: &'a [MetricFamily], name: &str) -> &'a MetricFamily { + families + .iter() + .find(|family| family.name() == name) + .unwrap_or_else(|| panic!("missing metric family {name}")) + } + + fn counter_values(families: &[MetricFamily]) -> BTreeMap> { + [ + "greptime_workload_scheduler_polls_total", + "greptime_workload_scheduler_admission_wait_seconds_total", + ] + .into_iter() + .map(|name| { + let values = family(families, name) + .get_metric() + .iter() + .map(|metric| { + ( + metric.get_label()[0].value().to_string(), + metric.get_counter().value(), + ) + }) + .collect(); + (name.to_string(), values) + }) + .collect() + } + + #[test] + fn workload_scheduler_collector_reports_class_metrics_and_deltas() { + let scheduler = Scheduler::builder() + .max_concurrent_polls(1) + .weight(QUERY_TASK_CLASS, 2) + .weight(WRITE_TASK_CLASS, 3) + .build(); + scheduler.set_enabled(true); + let collector = WorkloadSchedulerCollector::new(scheduler.clone()); + + let first = collector.collect(); + let expected = [ + ( + "greptime_workload_scheduler_enabled", + MetricType::GAUGE, + false, + ), + ( + "greptime_workload_scheduler_active_polls", + MetricType::GAUGE, + false, + ), + ( + "greptime_workload_scheduler_weight", + MetricType::GAUGE, + true, + ), + ( + "greptime_workload_scheduler_queued_tasks", + MetricType::GAUGE, + true, + ), + ( + "greptime_workload_scheduler_polls_total", + MetricType::COUNTER, + true, + ), + ( + "greptime_workload_scheduler_admission_wait_seconds_total", + MetricType::COUNTER, + true, + ), + ]; + let expected_names: BTreeSet<_> = expected.iter().map(|(name, _, _)| *name).collect(); + let actual_names: BTreeSet<_> = first.iter().map(MetricFamily::name).collect(); + assert_eq!(expected_names, actual_names); + for (name, metric_type, has_workload_label) in expected { + let metric_family = family(&first, name); + assert_eq!(metric_type, metric_family.get_field_type(), "{name}"); + let workloads: BTreeSet<_> = metric_family + .get_metric() + .iter() + .flat_map(|metric| { + assert_eq!(has_workload_label, !metric.get_label().is_empty()); + metric + .get_label() + .iter() + .map(|label| { + assert_eq!("workload", label.name()); + label.value() + }) + .collect::>() + }) + .collect(); + if has_workload_label { + assert_eq!(BTreeSet::from(["query", "write"]), workloads); + } else { + assert!(workloads.is_empty()); + } + } + + let second = collector.collect(); + assert_eq!(counter_values(&first), counter_values(&second)); + + scheduler.set_enabled(true); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap(); + runtime.block_on(async { + let query = scheduler.spawn_in(QUERY_TASK_CLASS, async { + tokio::time::sleep(Duration::from_millis(1)).await; + }); + let write = scheduler.spawn_in(WRITE_TASK_CLASS, async { + tokio::time::sleep(Duration::from_millis(1)).await; + }); + tokio::time::timeout(Duration::from_secs(1), async { + query.await.unwrap(); + write.await.unwrap(); + }) + .await + .expect("scheduled test tasks did not complete"); + }); + + let third = collector.collect(); + let before = counter_values(&second); + let after = counter_values(&third); + for workload in ["query", "write"] { + let metric = "greptime_workload_scheduler_polls_total"; + assert!( + after[metric][workload] > before[metric][workload], + "{metric} did not increase for {workload}" + ); + } + } +} diff --git a/src/common/runtime/src/runtime_default.rs b/src/common/runtime/src/runtime_default.rs index eedf410097..9e14221ab8 100644 --- a/src/common/runtime/src/runtime_default.rs +++ b/src/common/runtime/src/runtime_default.rs @@ -38,6 +38,10 @@ impl DefaultRuntime { _dropper: dropper, } } + + pub(crate) fn handle(&self) -> tokio::runtime::Handle { + self.handle.clone() + } } impl RuntimeTrait for DefaultRuntime { diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index c0f6af8184..b82ad4633a 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -108,6 +108,7 @@ pub mod result; pub mod splunk; mod timeout; pub mod utils; +mod workload_scheduler; use result::HttpOutputWriter; pub(crate) use timeout::DynamicTimeoutLayer; @@ -1077,6 +1078,18 @@ impl HttpServer { Router::new() // handler for changing log level dynamically .route("/log_level", routing::post(dyn_log::dyn_log_handler)) + .route( + "/workload_scheduler", + routing::get(workload_scheduler::get_status_handler), + ) + .route( + "/workload_scheduler/enabled", + routing::post(workload_scheduler::set_enabled_handler), + ) + .route( + "/workload_scheduler/weights", + routing::post(workload_scheduler::set_weights_handler), + ) .route("/enable_trace", routing::post(dyn_trace::dyn_trace_handler)) .nest( "/prof", diff --git a/src/servers/src/http/workload_scheduler.rs b/src/servers/src/http/workload_scheduler.rs new file mode 100644 index 0000000000..f9eea3452d --- /dev/null +++ b/src/servers/src/http/workload_scheduler.rs @@ -0,0 +1,111 @@ +// Copyright 2023 Greptime Team +// +// 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 std::num::NonZeroU32; + +use axum::Json; +use axum::body::Bytes; +use axum::http::StatusCode; +use axum::response::IntoResponse; +use serde::{Deserialize, Serialize}; +use snafu::{ResultExt, ensure}; + +use crate::error::{InvalidParameterSnafu, ParseJsonSnafu, Result}; + +#[axum_macros::debug_handler] +pub(super) async fn set_enabled_handler(body: Bytes) -> Result { + let enabled: bool = serde_json::from_slice(&body).context(ParseJsonSnafu)?; + ensure!( + common_runtime::set_workload_scheduler_enabled(enabled), + InvalidParameterSnafu { + reason: "workload scheduler was not constructed at startup", + } + ); + let change_note = format!("Workload scheduler enabled={enabled}"); + Ok((StatusCode::OK, change_note)) +} + +#[derive(Debug, Deserialize)] +struct SchedulerWeightsDto { + query: NonZeroU32, + write: NonZeroU32, +} + +#[axum_macros::debug_handler] +pub(super) async fn set_weights_handler(body: Bytes) -> Result { + let weights: SchedulerWeightsDto = serde_json::from_slice(&body).context(ParseJsonSnafu)?; + ensure!( + common_runtime::set_workload_scheduler_weights(weights.query, weights.write), + InvalidParameterSnafu { + reason: "workload scheduler was not constructed at startup", + } + ); + let change_note = format!( + "Workload scheduler weights query={}, write={}", + weights.query, weights.write + ); + Ok((StatusCode::OK, change_note)) +} + +/// Per-class scheduler status exposed by the HTTP API. +#[derive(Debug, Serialize)] +struct ClassStatusDto { + weight: u32, +} + +/// Point-in-time workload scheduler status. Scheduler class fields are omitted +/// when the scheduler was not constructed at startup or the corresponding +/// class is unavailable. +#[derive(Debug, Serialize)] +struct SchedulerStatusDto { + enabled: bool, + #[serde(skip_serializing_if = "Option::is_none")] + query: Option, + #[serde(skip_serializing_if = "Option::is_none")] + write: Option, +} + +/// Returns the current workload scheduler state and query/write weights. +/// Always returns 200, with `enabled=false` when the scheduler is dynamically +/// disabled. +#[axum_macros::debug_handler] +pub(super) async fn get_status_handler() -> Result { + let enabled = common_runtime::workload_scheduler_enabled(); + let Some(stats) = common_runtime::workload_scheduler_stats() else { + return Ok(Json(SchedulerStatusDto { + enabled: false, + query: None, + write: None, + })); + }; + + let mut query = None; + let mut write = None; + for (class, class_stats) in &stats.classes { + let status = ClassStatusDto { + weight: class_stats.weight, + }; + match class.id() { + 1 => query = Some(status), + 2 => write = Some(status), + _ => {} + } + } + + Ok(Json(SchedulerStatusDto { + enabled, + query, + write, + })) +}