From 4e940c89c8ae7d70c89678eecfcfa7fa6138e187 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 23 Oct 2026 23:26:13 +0800 Subject: [PATCH] fix: backport partial region write retries to release/5e02efa9 (#9881) * fix: backport partial region write retries to release/5e02efa9 Signed-off-by: WenyXu * ci: fix CI setup steps that fail on Ubuntu 26.04 runners (#9901) * ci: install protoc with taiki-e/install-action Signed-off-by: Dennis Zhuang * ci: run multi-language SDK tests on Python 3.13 Signed-off-by: Dennis Zhuang * ci: drop the unused libfuzzer-14-dev package from fuzz builds Signed-off-by: Dennis Zhuang * ci: build fuzz targets with the bundled libFuzzer Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang (cherry picked from commit 00ab09012d41d53fcf4a3c75fc030642c1c3faa3) Signed-off-by: WenyXu * chore: update greptime-proto to merged release commit Signed-off-by: WenyXu * fix(ci): move monitoring datanodeStorage out of standalone base (#9902) The `datanodeStorage` field belongs directly under `monitoring.standalone` (GreptimeDBStandaloneSpec), not under `monitoring.standalone.base` (PodTemplateSpec). Nesting it under `base` renders an invalid GreptimeDBCluster resource and the CRD schema rejects it during server-side apply: .spec.monitoring.standalone.base.datanodeStorage: field not declared in schema (cherry picked from commit 3483f4a2842a0b4d83a5d3654041fdeee6e8195d) Signed-off-by: WenyXu * fix(ci): specify Kind storage class for datanodes (#9908) Signed-off-by: WenyXu (cherry picked from commit 567849908e4de7c5c7fedcf753b96f0a27427120) Signed-off-by: WenyXu --------- Signed-off-by: WenyXu Signed-off-by: Dennis Zhuang Co-authored-by: dennis zhuang Co-authored-by: liyang --- .../build-windows-artifacts/action.yml | 4 +- .../with-observability.yaml | 10 +- .github/workflows/apidoc.yml | 4 +- .github/workflows/checks.yml | 4 +- .github/workflows/integration.yml | 22 +- .github/workflows/nightly-ci.yml | 8 +- .github/workflows/release.yml | 4 +- .github/workflows/run-multi-lang-tests.yml | 4 +- .github/workflows/rust.yml | 24 +- Cargo.lock | 48 +- Cargo.toml | 2 +- src/api/src/region.rs | 29 +- src/client/src/error.rs | 56 +- src/client/src/lib.rs | 12 +- src/client/src/region.rs | 482 +++++++++++- src/common/error/src/status_code.rs | 7 +- src/common/meta/src/error.rs | 18 + src/datanode/src/error.rs | 5 + src/datanode/src/region_server.rs | 743 +++++++++++++++--- src/meta-srv/src/procedure/utils.rs | 2 + src/metric-engine/AGENTS.md | 1 + src/metric-engine/src/engine.rs | 17 +- src/metric-engine/src/engine/put.rs | 298 ++++--- src/mito2/src/error.rs | 16 + src/mito2/src/flush.rs | 29 +- src/operator/src/write_dispatcher.rs | 114 ++- src/operator/src/write_dispatcher/tests.rs | 321 +++++--- src/servers/src/error.rs | 1 + src/servers/src/grpc/region_server.rs | 52 +- src/servers/src/mysql/writer.rs | 4 +- src/servers/src/postgres/types/error.rs | 3 +- src/standalone/Cargo.toml | 2 + src/standalone/src/datanode_manager.rs | 69 +- src/store-api/src/region_request.rs | 74 +- tests/runner/src/client.rs | 1 + 35 files changed, 1915 insertions(+), 575 deletions(-) diff --git a/.github/actions/build-windows-artifacts/action.yml b/.github/actions/build-windows-artifacts/action.yml index 2c4ae1f7db7..8fe077996d2 100644 --- a/.github/actions/build-windows-artifacts/action.yml +++ b/.github/actions/build-windows-artifacts/action.yml @@ -22,7 +22,9 @@ inputs: runs: using: composite steps: - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 + with: + tool: protoc@3.23.4 - name: Install rust toolchain uses: actions-rust-lang/setup-rust-toolchain@v1 diff --git a/.github/actions/setup-greptimedb-cluster/with-observability.yaml b/.github/actions/setup-greptimedb-cluster/with-observability.yaml index 581f60dea60..6cc923aede3 100644 --- a/.github/actions/setup-greptimedb-cluster/with-observability.yaml +++ b/.github/actions/setup-greptimedb-cluster/with-observability.yaml @@ -1,3 +1,7 @@ +datanode: + storage: + storageClassName: standard + tracing: enabled: true endpoint: "http://my-greptimedb-monitor-standalone.my-greptimedb:4000/v1/otlp/v1/traces" @@ -15,9 +19,9 @@ monitoring: limits: cpu: 500m memory: 256Mi - datanodeStorage: - fs: - storageSize: 5Gi + datanodeStorage: + fs: + storageSize: 5Gi vector: resources: requests: diff --git a/.github/workflows/apidoc.yml b/.github/workflows/apidoc.yml index eb7a720a421..fb5601fdc5b 100644 --- a/.github/workflows/apidoc.yml +++ b/.github/workflows/apidoc.yml @@ -19,9 +19,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - run: cargo doc --workspace --no-deps --document-private-items - run: | diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 682a2e5921a..9f1d79c7140 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -102,9 +102,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Rust Cache uses: Swatinem/rust-cache@v2 diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index d7a91cadd43..eef954305ef 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -37,9 +37,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Build Binaries Cache id: build-binaries-cache @@ -176,9 +176,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Rust Cache uses: Swatinem/rust-cache@v2 @@ -310,9 +310,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 with: cache: false @@ -344,13 +344,11 @@ jobs: if: steps.fuzz-binaries.outputs.cache-hit != 'true' shell: bash run: | - sudo apt-get install -y libfuzzer-14-dev cargo install cargo-fuzz --version "${CARGO_FUZZ_VERSION}" --locked - name: Build fuzz target binaries if: steps.fuzz-binaries.outputs.cache-hit != 'true' shell: bash env: - CUSTOM_LIBFUZZER_PATH: /usr/lib/llvm-14/lib/libFuzzer.a FUZZ_TARGETS: | fuzz_create_database fuzz_create_table @@ -564,9 +562,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Build Profile-CI Cache id: build-greptime-ci-cache @@ -682,9 +680,9 @@ jobs: uses: ./.github/actions/setup-etcd-cluster # Only the Kafka WAL helper needs the Rust toolchain. - if: matrix.mode.kafkaWalHelper - uses: arduino/setup-protoc@v3 + uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - if: matrix.mode.kafkaWalHelper uses: actions-rust-lang/setup-rust-toolchain@v1 with: diff --git a/.github/workflows/nightly-ci.yml b/.github/workflows/nightly-ci.yml index d4afa2778b8..23b78d77cb4 100644 --- a/.github/workflows/nightly-ci.yml +++ b/.github/workflows/nightly-ci.yml @@ -53,9 +53,9 @@ jobs: fetch-depth: 0 persist-credentials: false - uses: ./.github/actions/setup-cyborg - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Rust Cache uses: Swatinem/rust-cache@v2 @@ -83,9 +83,9 @@ jobs: fetch-depth: 0 persist-credentials: false - uses: ./.github/actions/setup-cyborg - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: KyleMayes/install-llvm-action@v1 with: version: "14.0" diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index acf52e241c2..cf4ce574d83 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -452,9 +452,9 @@ jobs: fetch-depth: 0 persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 with: diff --git a/.github/workflows/run-multi-lang-tests.yml b/.github/workflows/run-multi-lang-tests.yml index 52fdca2117c..dad4fbaa080 100644 --- a/.github/workflows/run-multi-lang-tests.yml +++ b/.github/workflows/run-multi-lang-tests.yml @@ -83,10 +83,10 @@ jobs: java-version: '17' cache: 'maven' - - name: Setup Python 3.8 + - name: Setup Python 3.13 uses: actions/setup-python@v5 with: - python-version: '3.8' + python-version: '3.13' - name: Setup uv uses: astral-sh/setup-uv@v3 diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 6d12566b684..f5de10a1e2b 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -34,9 +34,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 with: components: rustfmt @@ -52,9 +52,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 with: components: clippy @@ -78,9 +78,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Rust Cache uses: Swatinem/rust-cache@v2 @@ -105,9 +105,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - name: Install riscv64 cross toolchain run: sudo apt-get update && sudo apt-get install -y gcc-riscv64-linux-gnu g++-riscv64-linux-gnu - uses: actions-rust-lang/setup-rust-toolchain@v1 @@ -142,9 +142,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: rui314/setup-mold@v1 - name: Install toolchain uses: actions-rust-lang/setup-rust-toolchain@v1 @@ -198,9 +198,9 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 + - uses: taiki-e/install-action@v2 with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + tool: protoc@3.23.4 - uses: rui314/setup-mold@v1 - name: Install toolchain uses: actions-rust-lang/setup-rust-toolchain@v1 diff --git a/Cargo.lock b/Cargo.lock index 3e1a9ea9b75..72e72cf07ad 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1630,7 +1630,7 @@ dependencies = [ "maybe-owned", "rustix 1.0.7", "rustix-linux-procfs", - "windows-sys 0.61.2", + "windows-sys 0.60.2", "winx", ] @@ -2280,7 +2280,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] @@ -5186,7 +5186,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.60.2", ] [[package]] @@ -5782,7 +5782,7 @@ checksum = "94e7099f6313ecacbe1256e8ff9d617b75d1bcb16a6fddef94866d225a01a14a" dependencies = [ "io-lifetimes 2.0.4", "rustix 1.0.7", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -6137,7 +6137,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=0c4beffaa2b3968034915867593a717c6ab0c6b4#0c4beffaa2b3968034915867593a717c6ab0c6b4" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=97175efcd40e964b283a21061dce68deaa2a0136#97175efcd40e964b283a21061dce68deaa2a0136" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", @@ -6684,7 +6684,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.4", + "socket2 0.5.10", "tokio", "tower-service", "tracing", @@ -7183,7 +7183,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20fd6de4ccfcc187e38bc21cfa543cb5a302cb86a8b114eb7f0bf0dc9f8ac00f" dependencies = [ "io-lifetimes 3.0.1", - "windows-sys 0.52.0", + "windows-sys 0.60.2", ] [[package]] @@ -7221,7 +7221,7 @@ checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -9235,7 +9235,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -11473,7 +11473,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.4.1", + "heck 0.5.0", "itertools 0.12.1", "log", "multimap", @@ -11493,8 +11493,8 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.4.1", - "itertools 0.14.0", + "heck 0.5.0", + "itertools 0.10.5", "log", "multimap", "once_cell", @@ -11542,7 +11542,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -11555,7 +11555,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -12012,7 +12012,7 @@ dependencies = [ "once_cell", "socket2 0.5.10", "tracing", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -12984,7 +12984,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.4.15", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -12997,7 +12997,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.9.4", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -13088,7 +13088,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.60.2", ] [[package]] @@ -13975,7 +13975,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.117", @@ -14004,7 +14004,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -14399,7 +14399,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -15016,7 +15016,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.0.7", - "windows-sys 0.52.0", + "windows-sys 0.60.2", ] [[package]] @@ -16823,7 +16823,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] @@ -17273,7 +17273,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f3fd376f71958b862e7afb20cfe5a22830e1963462f3a17f49d82a6c1d1f42d" dependencies = [ "bitflags 2.12.1", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index babbf32751e..32acb740f99 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -159,7 +159,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "0c4beffaa2b3968034915867593a717c6ab0c6b4" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "97175efcd40e964b283a21061dce68deaa2a0136" } hex = "0.4" hostname = "0.4.0" http = "1" diff --git a/src/api/src/region.rs b/src/api/src/region.rs index e7000cd7448..fb0d5b5e3f4 100644 --- a/src/api/src/region.rs +++ b/src/api/src/region.rs @@ -15,7 +15,19 @@ use std::collections::HashMap; use common_base::AffectedRows; -use greptime_proto::v1::region::RegionResponse as RegionResponseV1; +use greptime_proto::v1::region::{RegionRequestHeader, RegionResponse as RegionResponseV1}; + +/// Request extension advertising support for per-entry write failures. +pub const ACCEPT_PARTIAL_RESULT: &str = "accept_partial_result"; + +/// Advertises support for per-entry write failures in the request header. +pub fn set_accept_partial_result(header: &mut RegionRequestHeader) { + header + .query_context + .get_or_insert_with(Default::default) + .extensions + .insert(ACCEPT_PARTIAL_RESULT.to_string(), "true".to_string()); +} /// This result struct is derived from [RegionResponseV1] #[derive(Debug)] @@ -25,6 +37,21 @@ pub struct RegionResponse { pub metadata: Vec, } +/// Successful entries and failed-entry errors indexed by input order. +#[derive(Debug)] +pub struct BatchRegionResponse { + pub affected_rows: AffectedRows, + pub extensions: HashMap>, + pub failures: Vec<(usize, E)>, +} + +impl BatchRegionResponse { + /// Returns the failure with the smallest original request index. + pub fn first_failure(&self) -> Option<&(usize, E)> { + self.failures.iter().min_by_key(|(index, _)| *index) + } +} + impl RegionResponse { pub fn from_region_response(region_response: RegionResponseV1) -> Self { Self { diff --git a/src/client/src/error.rs b/src/client/src/error.rs index fdf0df9582f..d9fa8c96050 100644 --- a/src/client/src/error.rs +++ b/src/client/src/error.rs @@ -18,7 +18,7 @@ use common_error::define_from_tonic_status; use common_error::ext::{BoxedError, ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use common_macro::stack_trace_debug; -use snafu::{Location, Snafu, location}; +use snafu::{Location, OptionExt, Snafu, location}; use tonic::Code; use tonic::metadata::errors::InvalidMetadataValue; @@ -26,6 +26,15 @@ use tonic::metadata::errors::InvalidMetadataValue; #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display( + "Region request has {} failed entries; {}", + response.failures.len(), + response.first_failure().map(|(index, error)| format!("entry {index}: {}", error.output_msg())).unwrap_or_default() + ))] + PartialRegion { + response: Box>, + }, + #[snafu(display("Illegal Flight messages, reason: {}", reason))] IllegalFlightMessages { reason: String, @@ -109,11 +118,12 @@ pub enum Error { location: Location, }, - // Server error carried in Tonic Status's metadata. + // Server error carried in a protobuf response status. #[snafu(display("{}", msg))] Server { code: StatusCode, msg: String, + retry_hint: RetryHint, #[snafu(implicit)] location: Location, }, @@ -163,6 +173,9 @@ pub type Result = std::result::Result; impl ErrorExt for Error { fn status_code(&self) -> StatusCode { match self { + Error::PartialRegion { response } => response + .first_failure() + .map_or(StatusCode::Internal, |(_, error)| error.status_code()), Error::IllegalFlightMessages { .. } | Error::MissingField { .. } | Error::IllegalDatabaseResponse { .. } => StatusCode::Internal, @@ -188,7 +201,8 @@ impl ErrorExt for Error { fn retry_hint(&self) -> RetryHint { match self { - Error::Tonic { retry_hint, .. } => *retry_hint, + Error::PartialRegion { .. } => RetryHint::Retryable, + Error::Tonic { retry_hint, .. } | Error::Server { retry_hint, .. } => *retry_hint, Error::FlightGet { source, .. } | Error::FlightStream { source, .. } | Error::FlowServer { source, .. } @@ -220,9 +234,22 @@ define_from_tonic_status!(Error, Tonic); impl From for common_meta::error::Error { fn from(error: Error) -> Self { - Self::External { - source: BoxedError::new(error), - location: location!(), + match error { + Error::PartialRegion { response } => Self::PartialRegion { + response: Box::new(api::region::BatchRegionResponse { + affected_rows: response.affected_rows, + extensions: response.extensions, + failures: response + .failures + .into_iter() + .map(|(index, error)| (index, error.into())) + .collect(), + }), + }, + error => Self::External { + source: BoxedError::new(error), + location: location!(), + }, } } } @@ -258,3 +285,20 @@ impl Error { ) } } + +impl Error { + /// Decodes an application error carried in a protobuf response. + pub(crate) fn try_from_response_status(status: &api::v1::Status) -> Result { + let code = StatusCode::from_u32(status.status_code).with_context(|| { + IllegalDatabaseResponseSnafu { + err_msg: format!("Unknown server status: {}", status.status_code), + } + })?; + Ok(ServerSnafu { + code, + msg: status.err_msg.clone(), + retry_hint: status.retry_hint.parse().unwrap_or(RetryHint::NonRetryable), + } + .build()) + } +} diff --git a/src/client/src/lib.rs b/src/client/src/lib.rs index 4b9d543e35f..208a3233ff7 100644 --- a/src/client/src/lib.rs +++ b/src/client/src/lib.rs @@ -35,7 +35,7 @@ use snafu::OptionExt; pub use self::client::{Client, ClientOptions}; pub use self::database::{Database, OutputMetrics, OutputWithMetrics}; pub use self::error::{Error, Result}; -use crate::error::{IllegalDatabaseResponseSnafu, ServerSnafu}; +use crate::error::IllegalDatabaseResponseSnafu; pub fn from_grpc_response(response: GreptimeResponse) -> Result { let header = response.header.context(IllegalDatabaseResponseSnafu { @@ -53,14 +53,6 @@ pub fn from_grpc_response(response: GreptimeResponse) -> Result { Response::AffectedRows(AffectedRows { value }) => Ok(value), } } else { - let status_code = - StatusCode::from_u32(status.status_code).context(IllegalDatabaseResponseSnafu { - err_msg: format!("invalid status: {:?}", status), - })?; - ServerSnafu { - code: status_code, - msg: status.err_msg, - } - .fail() + Err(Error::try_from_response_status(&status)?) } } diff --git a/src/client/src/region.rs b/src/client/src/region.rs index f9ebbb0d0bf..255e67de2f6 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -18,11 +18,11 @@ use std::borrow::Borrow; use std::sync::Arc; use std::time::Duration; -use api::region::RegionResponse; +use api::region::{BatchRegionResponse, RegionResponse}; use api::v1::ResponseHeader; use api::v1::region::{ - RegionRequest, RegionRequestHeader, RemoteDynFilterRequest, RemoteDynFilterUnregister, - RemoteDynFilterUpdate, region_request, remote_dyn_filter_request, + RegionRequest, RegionRequestHeader, RegionResponse as PbRegionResponse, RemoteDynFilterRequest, + RemoteDynFilterUnregister, RemoteDynFilterUpdate, region_request, remote_dyn_filter_request, }; use arc_swap::ArcSwapOption; use arrow_flight::Ticket; @@ -49,7 +49,7 @@ use tonic::codec::CompressionEncoding; use crate::error::{ self, FlightGetSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu, - MissingFieldSnafu, Result, ServerSnafu, + MissingFieldSnafu, Result, }; use crate::flight::{FlightMessageReader, decode_flight_data}; use crate::region::codec::RegionCodec; @@ -141,7 +141,8 @@ impl RegionRequester { /// /// Accepts both [`RegionRequest`] and [`Arc`]. Callers can retain a /// shared request for another attempt; this method does not retry or determine - /// whether replaying the request is safe. + /// whether replaying the request is safe. Callers accepting partial results in the + /// request header must inspect failed entries before retrying. pub async fn handle(&self, request: R) -> Result where R: Borrow + Send + Sync + 'static, @@ -180,6 +181,11 @@ impl RegionRequester { .body .as_ref() .context(MissingFieldSnafu { field: "body" })?; + let request_count = match request_body { + region_request::Body::Inserts(requests) => Some(requests.requests.len()), + region_request::Body::Deletes(requests) => Some(requests.requests.len()), + _ => None, + }; let is_insert = matches!( request_body, region_request::Body::Inserts(_) | region_request::Body::BulkInsert(_) @@ -199,6 +205,13 @@ impl RegionRequester { } } + let accept_partial_result = request + .borrow() + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .and_then(|context| context.extensions.get(api::region::ACCEPT_PARTIAL_RESULT)) + .is_some_and(|value| value == "true"); let response = async { client .ready() @@ -223,9 +236,10 @@ impl RegionRequester { .map_err(|status| map_region_error(status, addr))? .into_inner(); - check_response_header(&response.header)?; - - Ok(RegionResponse::from_region_response(response)) + match request_count { + Some(count) if accept_partial_result => decode_write_response(response, count), + _ => decode_region_response(response), + } } } @@ -385,6 +399,7 @@ fn build_remote_dyn_filter_request( } } +/// Checks that the response contains a successful application status. pub fn check_response_header(header: &Option) -> Result<()> { let status = header .as_ref() @@ -396,18 +411,88 @@ pub fn check_response_header(header: &Option) -> Result<()> { if StatusCode::is_success(status.status_code) { Ok(()) } else { - let code = - StatusCode::from_u32(status.status_code).context(IllegalDatabaseResponseSnafu { - err_msg: format!("unknown server status: {:?}", status), - })?; - ServerSnafu { - code, - msg: status.err_msg.clone(), - } - .fail() + Err(error::Error::try_from_response_status(status)?) } } +/// Decodes an ordinary region response without per-entry write failures. +pub fn decode_region_response(response: PbRegionResponse) -> Result { + if !response.failures.is_empty() { + return IllegalDatabaseResponseSnafu { + err_msg: "Failure details are only valid for inserts/deletes", + } + .fail(); + } + check_response_header(&response.header)?; + Ok(RegionResponse::from_region_response(response)) +} + +/// Decodes an inserts/deletes response, validating failures against the original entry count. +/// A failed single-entry request returns its underlying error. +fn decode_write_response(response: PbRegionResponse, entry_count: usize) -> Result { + if response.failures.is_empty() { + if response + .header + .as_ref() + .and_then(|header| header.status.as_ref()) + .is_some_and(|status| status.status_code == StatusCode::PartialFailure as u32) + { + return IllegalDatabaseResponseSnafu { + err_msg: "Partial failure has no failed entries", + } + .fail(); + } + return decode_region_response(response); + } + let status = response + .header + .as_ref() + .and_then(|header| header.status.as_ref()) + .context(IllegalDatabaseResponseSnafu { + err_msg: "Missing region response status", + })?; + if status.status_code != StatusCode::PartialFailure as u32 { + return IllegalDatabaseResponseSnafu { + err_msg: "Failure details require a partial-failure status", + } + .fail(); + } + let mut seen = std::collections::HashSet::with_capacity(response.failures.len()); + let mut failures = Vec::with_capacity(response.failures.len()); + for error in response.failures { + let index = error.request_index as usize; + if index >= entry_count || !seen.insert(index) { + return IllegalDatabaseResponseSnafu { + err_msg: "Invalid or duplicate batch error index", + } + .fail(); + } + let status = error.status.context(IllegalDatabaseResponseSnafu { + err_msg: "Missing batch error status", + })?; + if StatusCode::is_success(status.status_code) { + return IllegalDatabaseResponseSnafu { + err_msg: "Batch error contains a success status", + } + .fail(); + } + failures.push((index, error::Error::try_from_response_status(&status)?)); + } + // A single failed entry has no successful peers to preserve. + if entry_count == 1 + && let Some((_, error)) = failures.pop() + { + return Err(error); + } + Err(error::Error::PartialRegion { + response: Box::new(BatchRegionResponse { + affected_rows: response.affected_rows as _, + extensions: response.extensions, + failures, + }), + }) +} + #[cfg(test)] #[allow(deprecated)] mod tests { @@ -416,15 +501,17 @@ mod tests { use api::v1::region::region_server::{Region, RegionServer}; use api::v1::region::{ - BulkInsertRequest, InsertRequest, InsertRequests, RegionResponse as PbRegionResponse, - RemoteDynFilterUnregister, RemoteDynFilterUpdate, region_request, - remote_dyn_filter_request, + BulkInsertRequest, InsertRequest, InsertRequests, RegionRequestFailure, + RegionResponse as PbRegionResponse, RemoteDynFilterUnregister, RemoteDynFilterUpdate, + region_request, remote_dyn_filter_request, }; use api::v1::value::ValueData; use api::v1::{ - ColumnDataType, ColumnSchema as PbColumnSchema, Row, Rows, Status as PbStatus, Value, + ColumnDataType, ColumnSchema as PbColumnSchema, ResponseHeader, Row, Rows, + Status as PbStatus, Value, }; use common_error::ext::RetryHint; + use common_error::status_code::StatusCode; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_recordbatch::adapter::RecordBatchMetrics; use datatypes::arrow::array::Int32Array; @@ -579,6 +666,20 @@ mod tests { } } + fn with_partial_result(mut request: RegionRequest) -> RegionRequest { + request + .header + .get_or_insert_with(Default::default) + .query_context + .get_or_insert_with(Default::default) + .extensions + .insert( + api::region::ACCEPT_PARTIAL_RESULT.to_string(), + "true".to_string(), + ); + request + } + #[tokio::test] async fn test_owned_and_shared_requests_preserve_payload_and_compression() { let server = RegionTestServer::start().await; @@ -636,6 +737,54 @@ mod tests { } } + #[tokio::test] + async fn test_partial_result_requires_explicit_acceptance() { + let server = RegionTestServer::start().await; + let requester = server.create_requester(ChannelConfig::default(), false, false); + for accepted in [false, true] { + let request = if accepted { + with_partial_result(build_insert_request()) + } else { + build_insert_request() + }; + let shared = Arc::new(request); + server + .service + .responses + .lock() + .unwrap() + .push_back(Ok(PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::PartialFailure as u32, + ..Default::default() + }), + }), + failures: vec![RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: StatusCode::RegionNotFound as u32, + retry_hint: RetryHint::RetryAfterRouteRefresh.as_str().to_string(), + ..Default::default() + }), + }], + ..Default::default() + })); + let error = requester.handle(Arc::clone(&shared)).await.unwrap_err(); + if accepted { + assert_eq!(error.status_code(), StatusCode::RegionNotFound); + assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh); + } else { + assert!(matches!(error, Error::IllegalDatabaseResponse { .. })); + } + assert_eq!( + &server.service.requests.lock().unwrap().last().unwrap().0, + shared.as_ref() + ); + assert_eq!(Arc::strong_count(&shared), 1); + } + } + #[tokio::test] async fn test_shared_request_can_be_resent_after_transport_error() { let server = RegionTestServer::start().await; @@ -762,7 +911,10 @@ mod tests { assert_eq!(error.tonic_code(), baseline.tonic_code()); } let datanode: &dyn Datanode = &requester; - let error = datanode.handle(build_insert_request()).await.unwrap_err(); + let error = datanode + .handle_request(Arc::new(build_insert_request())) + .await + .unwrap_err(); assert_eq!(error.status_code(), StatusCode::RegionNotFound); assert_eq!(error.retry_hint(), hint); assert!(matches!(error, meta_error::Error::External { .. })); @@ -832,10 +984,16 @@ mod tests { body: Some(body), }; let shared = Arc::new(request); - let shared_error = datanode.handle(Arc::clone(&shared)).await.unwrap_err(); + let shared_error = datanode + .handle_request(Arc::clone(&shared)) + .await + .unwrap_err(); assert_eq!(Arc::strong_count(&shared), 1); let request = Arc::try_unwrap(shared).unwrap(); - let owned_error = datanode.handle(request).await.unwrap_err(); + let owned_error = datanode + .handle_request(Arc::new(request)) + .await + .unwrap_err(); for error in [owned_error, shared_error] { assert_eq!( error.retry_hint(), @@ -1053,8 +1211,281 @@ mod tests { .unwrap() } + #[tokio::test] + async fn test_partial_response_validation_and_compression() { + let server = RegionTestServer::start().await; + let requester = server.create_requester(ChannelConfig::default(), true, true); + let mut request = build_insert_request(); + if let Some(region_request::Body::Inserts(inserts)) = &mut request.body { + inserts.requests.push(inserts.requests[0].clone()); + } + let request = Arc::new(with_partial_result(request)); + let failure = RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: StatusCode::RegionNotFound as u32, + err_msg: "moved".to_string(), + retry_hint: "retry_after_route_refresh".to_string(), + }), + }; + for (failures, valid) in [ + (vec![], true), + (vec![failure.clone()], true), + (vec![failure.clone(), failure.clone()], false), + ( + vec![RegionRequestFailure { + request_index: u32::MAX, + ..failure.clone() + }], + false, + ), + ( + vec![RegionRequestFailure { + request_index: 0, + status: None, + }], + false, + ), + ( + vec![RegionRequestFailure { + request_index: 0, + status: Some(PbStatus::default()), + }], + false, + ), + ( + vec![RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: u32::MAX, + ..Default::default() + }), + }], + false, + ), + ] { + let has_error = !failures.is_empty(); + server + .service + .responses + .lock() + .unwrap() + .push_back(Ok(PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: if has_error { + StatusCode::PartialFailure + } else { + StatusCode::Success + } as u32, + ..Default::default() + }), + }), + affected_rows: 3, + extensions: std::collections::HashMap::from([("test".to_string(), vec![1, 2])]), + failures, + ..Default::default() + })); + let response = requester.handle_inner(Arc::clone(&request)).await; + if valid { + if has_error { + let error = response.unwrap_err(); + assert!(error.output_msg().contains("entry 0")); + assert!(error.output_msg().contains("moved")); + let error: common_meta::error::Error = error.into(); + assert!(error.output_msg().contains("moved")); + let common_meta::error::Error::PartialRegion { response } = error else { + panic!("Expected partial region failure"); + }; + assert_eq!(response.affected_rows, 3); + assert_eq!(response.extensions["test"], vec![1, 2]); + assert_eq!( + response.failures[0].1.retry_hint(), + RetryHint::RetryAfterRouteRefresh + ); + } else { + let response = response.unwrap(); + assert_eq!(response.affected_rows, 3); + assert_eq!(response.extensions["test"], vec![1, 2]); + } + } else { + assert!(matches!( + response.unwrap_err(), + Error::IllegalDatabaseResponse { .. } + )); + } + } + for (received, send, accept) in server.service.requests.lock().unwrap().iter() { + assert_eq!(received, request.as_ref()); + assert!(*send && *accept); + } + assert_eq!(Arc::strong_count(&request), 1); + } + + #[test] + fn test_partial_response_envelope() { + let failure = RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: StatusCode::RegionNotFound as u32, + retry_hint: RetryHint::RetryAfterRouteRefresh.as_str().to_string(), + ..Default::default() + }), + }; + for (status, failures, count) in [ + (Some(StatusCode::Success), vec![failure.clone()], Some(2)), + (Some(StatusCode::PartialFailure), vec![], Some(2)), + ( + Some(StatusCode::PartialFailure), + vec![failure.clone()], + None, + ), + (None, vec![failure.clone()], Some(2)), + ] { + let response = PbRegionResponse { + header: status.map(|code| ResponseHeader { + status: Some(PbStatus { + status_code: code as u32, + ..Default::default() + }), + }), + failures, + ..Default::default() + }; + assert!(matches!( + match count { + Some(count) => decode_write_response(response, count), + None => decode_region_response(response), + }, + Err(Error::IllegalDatabaseResponse { .. }) + )); + } + let response = PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::PartialFailure as u32, + ..Default::default() + }), + }), + failures: vec![failure], + ..Default::default() + }; + // A partial envelope must not be interpreted as success by the header-only decoder. + assert!(check_response_header(&response.header).is_err()); + let error = decode_write_response(response, 1).unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RegionNotFound); + assert_eq!(error.retry_hint(), RetryHint::RetryAfterRouteRefresh); + } + + #[test] + fn test_partial_error_preserves_status_and_retry_hint() { + let expected_message = "Region request has 2 failed entries; entry 0: first input failed"; + for code in [StatusCode::StorageUnavailable, StatusCode::InvalidArguments] { + let response = PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::PartialFailure as u32, + ..Default::default() + }), + }), + failures: vec![ + RegionRequestFailure { + request_index: 1, + status: Some(PbStatus { + status_code: StatusCode::Internal as u32, + err_msg: "later input failed".to_string(), + ..Default::default() + }), + }, + RegionRequestFailure { + request_index: 0, + status: Some(PbStatus { + status_code: code as u32, + err_msg: "first input failed".to_string(), + retry_hint: RetryHint::NonRetryable.as_str().to_string(), + }), + }, + ], + ..Default::default() + }; + let error = decode_write_response(response, 2).unwrap_err(); + assert_eq!(code, error.status_code()); + assert_eq!(expected_message, error.to_string()); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + let error: common_meta::error::Error = error.into(); + assert_eq!(code, error.status_code()); + assert_eq!(expected_message, error.to_string()); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + let common_meta::error::Error::PartialRegion { response } = error else { + panic!("expected partial result") + }; + assert!( + response + .failures + .iter() + .all(|(_, error)| error.retry_hint() == RetryHint::NonRetryable) + ); + } + } + + #[tokio::test] + async fn test_region_rpc_errors() { + let server = RegionTestServer::start().await; + let requester = server.create_requester(ChannelConfig::default(), false, false); + let request = Arc::new(build_insert_request()); + let mut business_error = Status::unimplemented("unsupported write"); + business_error.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_CODE, + (StatusCode::Unsupported as u32) + .to_string() + .parse() + .unwrap(), + ); + for status in [ + Status::unimplemented("unknown method"), + business_error, + Status::unavailable("lost response"), + Status::deadline_exceeded("timeout"), + ] { + let code = status.code(); + server + .service + .responses + .lock() + .unwrap() + .push_back(Err(status)); + let error = requester + .handle_inner(Arc::clone(&request)) + .await + .unwrap_err(); + assert!(matches!(error, Error::RegionServer { code: actual, .. } if actual == code)); + } + assert_eq!(server.service.requests.lock().unwrap().len(), 4); + } + #[test] fn test_check_response_header() { + for (hint, expected) in [ + ("retryable", RetryHint::Retryable), + ( + "retry_after_route_refresh", + RetryHint::RetryAfterRouteRefresh, + ), + ("non_retryable", RetryHint::NonRetryable), + ("", RetryHint::NonRetryable), + ("future_hint", RetryHint::NonRetryable), + ] { + let error = check_response_header(&Some(ResponseHeader { + status: Some(PbStatus { + status_code: StatusCode::RegionNotFound as u32, + err_msg: "region moved".to_string(), + retry_hint: hint.to_string(), + }), + })) + .unwrap_err(); + assert_eq!(error.retry_hint(), expected, "{hint}"); + } + let result = check_response_header(&None); assert!(matches!( result.unwrap_err(), @@ -1071,6 +1502,7 @@ mod tests { status: Some(PbStatus { status_code: StatusCode::Success as u32, err_msg: String::default(), + ..Default::default() }), })); assert!(result.is_ok()); @@ -1079,6 +1511,7 @@ mod tests { status: Some(PbStatus { status_code: u32::MAX, err_msg: String::default(), + ..Default::default() }), })); assert!(matches!( @@ -1090,6 +1523,7 @@ mod tests { status: Some(PbStatus { status_code: StatusCode::Internal as u32, err_msg: "blabla".to_string(), + ..Default::default() }), })); let Error::Server { code, msg, .. } = result.unwrap_err() else { diff --git a/src/common/error/src/status_code.rs b/src/common/error/src/status_code.rs index b2c99c8e06b..592e40fec45 100644 --- a/src/common/error/src/status_code.rs +++ b/src/common/error/src/status_code.rs @@ -45,6 +45,8 @@ pub enum StatusCode { DeadlineExceeded = 1008, /// Service got suspended for various reason. For example, resources exceed limit. Suspended = 1009, + /// One or more entries failed; inspect per-entry results before retrying. + PartialFailure = 1010, // ====== End of common status code ================ // ====== Begin of SQL related status code ========= @@ -150,7 +152,8 @@ impl StatusCode { | StatusCode::RuntimeResourcesExhausted | StatusCode::External => true, - StatusCode::Success + StatusCode::PartialFailure + | StatusCode::Success | StatusCode::Unsupported | StatusCode::InvalidArguments | StatusCode::InvalidSyntax @@ -300,7 +303,7 @@ macro_rules! define_into_tonic_status { pub fn status_to_tonic_code(status_code: StatusCode) -> Code { match status_code { StatusCode::Success => Code::Ok, - StatusCode::Unknown | StatusCode::External => Code::Unknown, + StatusCode::Unknown | StatusCode::External | StatusCode::PartialFailure => Code::Unknown, StatusCode::Unsupported => Code::Unimplemented, StatusCode::Unexpected | StatusCode::IllegalState diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index de3831589e7..0d74992a5e0 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -41,6 +41,18 @@ pub use retry_hint::{retry_hint_from_postgres_error, retry_hint_from_postgres_po #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display("Failed to execute batch request"))] + BatchRequest { source: Arc }, + + #[snafu(display( + "Region request has {} failed entries; {}", + response.failures.len(), + response.first_failure().map(|(index, error)| format!("entry {index}: {}", error.output_msg())).unwrap_or_default() + ))] + PartialRegion { + response: Box>, + }, + #[snafu(display("Empty key is not allowed"))] EmptyKey { #[snafu(implicit)] @@ -1167,6 +1179,10 @@ impl ErrorExt for Error { fn status_code(&self) -> StatusCode { use Error::*; match self { + Error::PartialRegion { response } => response + .first_failure() + .map_or(StatusCode::Internal, |(_, error)| error.status_code()), + Error::BatchRequest { source } => source.status_code(), IllegalServerState { .. } | EtcdTxnOpResponse { .. } | EtcdFailed { .. } @@ -1330,6 +1346,8 @@ impl ErrorExt for Error { use Error::*; match self { + Error::PartialRegion { .. } => RetryHint::Retryable, + Error::BatchRequest { source } => source.retry_hint(), RetryLater { source, .. } => match source.retry_hint() { RetryHint::RetryAfterRouteRefresh => RetryHint::RetryAfterRouteRefresh, RetryHint::Retryable | RetryHint::NonRetryable => RetryHint::Retryable, diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index 96ccd7ef354..4ab7f9c7104 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -30,6 +30,9 @@ use tokio::time::error::Elapsed; #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display("Failed to write a batch group"))] + BatchWrite { source: std::sync::Arc }, + #[snafu(display("Failed to execute async task"))] AsyncTaskExecute { #[snafu(implicit)] @@ -446,6 +449,7 @@ impl ErrorExt for Error { fn status_code(&self) -> StatusCode { use Error::*; match self { + Self::BatchWrite { source } => source.status_code(), NewPlanDecoder { source, .. } | ExecuteLogicalPlan { source, .. } => { source.status_code() } @@ -524,6 +528,7 @@ impl ErrorExt for Error { use Error::*; match self { + Self::BatchWrite { source } => source.retry_hint(), RegionNotReady { .. } | RegionNotFound { .. } => RetryHint::RetryAfterRouteRefresh, RegionBusy { .. } | ConcurrentQueryLimiterClosed { .. } diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index f3e5e1a0f60..411b13004d7 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -25,11 +25,12 @@ use std::sync::{Arc, RwLock}; use std::task::{Context, Poll}; use std::time::Duration; -use api::region::RegionResponse; +use api::region::{BatchRegionResponse, RegionResponse}; use api::v1::meta::TopicStat; use api::v1::region::sync_request::ManifestInfo; use api::v1::region::{ - ListMetadataRequest, RegionResponse as RegionResponseV1, SyncRequest, region_request, + ListMetadataRequest, RegionRequest as RegionRequestV1, RegionRequestFailure, + RegionResponse as RegionResponseV1, SyncRequest, region_request, }; use api::v1::{ResponseHeader, Status}; use arrow_flight::{FlightData, Ticket}; @@ -51,7 +52,6 @@ use datafusion::datasource::TableProvider; use datafusion_common::tree_node::TreeNode; use datatypes::schema::SchemaRef; use either::Either; -use futures_util::future::try_join_all; use futures_util::{Stream, StreamExt}; use metric_engine::engine::MetricEngine; use mito2::engine::{MITO_ENGINE_NAME, MitoEngine}; @@ -74,7 +74,7 @@ use servers::grpc::region_server::RegionServerHandler; use session::context::{ FLIGHT_METRICS_HEARTBEAT_INTERVAL, QueryContext, QueryContextBuilder, QueryContextRef, }; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; use store_api::metric_engine_consts::{ FILE_ENGINE_NAME, LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, }; @@ -84,8 +84,8 @@ use store_api::region_engine::{ SyncRegionFromRequest, }; use store_api::region_request::{ - AffectedRows, BatchRegionDdlRequest, RegionCatchupRequest, RegionCloseRequest, - RegionOpenRequest, RegionRequest, + self as store_region_request, AffectedRows, BatchRegionDdlRequest, RegionCatchupRequest, + RegionCloseRequest, RegionOpenRequest, RegionRequest, }; use store_api::storage::RegionId; use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; @@ -511,7 +511,7 @@ impl RegionServer { let span = tracing_context.attach(info_span!("RegionServer::handle_batch_ddl_requests")); self.inner - .handle_batch_request(batch_request) + .handle_batch_ddl_requests(batch_request) .trace(span) .await } @@ -519,49 +519,46 @@ impl RegionServer { async fn handle_requests_in_parallel( &self, request: region_request::Body, - ) -> Result { - let requests = - RegionRequest::try_from_request_body(request).context(BuildRegionRequestsSnafu)?; - - // Try to optimize batch Put requests for metric engine - // Returns either Some(response) or None(requests_back) - match self.try_handle_metric_batch_puts(requests).await? { - Either::Left(response) => Ok(response), - Either::Right(requests) => { - // Fallback: original parallel processing - let tracing_context = TracingContext::from_current_span(); - let join_tasks = - requests - .into_iter() - .map(|(region_id, req): (RegionId, RegionRequest)| { - let self_to_move = self; - let span = tracing_context.attach(info_span!( - "RegionServer::handle_region_request", - region_id = region_id.to_string() - )); - async move { - self_to_move - .handle_request(region_id, req) - .trace(span) - .await - } - }); - - let results = try_join_all(join_tasks).await?; - let mut affected_rows = 0; - let mut extensions = HashMap::new(); - for result in results { - affected_rows += result.affected_rows; - extensions.extend(result.extensions); + ) -> Result> { + let requests = prepare_batch_write_requests(request); + let requests = match self.try_handle_metric_batch_puts(requests).await? { + Either::Left(response) => return Ok(response), + Either::Right(requests) => requests, + }; + let tracing_context = TracingContext::from_current_span(); + let results = futures::future::join_all(requests.into_iter().map(|request| { + let span = tracing_context.attach(info_span!( + "RegionServer::handle_region_request", + region_id = request.region_id.to_string() + )); + async move { + ( + request.index, + self.handle_request(request.region_id, request.request) + .trace(span) + .await, + ) + } + })) + .await; + let mut response = BatchRegionResponse { + affected_rows: 0, + extensions: HashMap::new(), + failures: Vec::new(), + }; + for (index, result) in results { + match result { + Ok(result) => { + response.affected_rows += result.affected_rows; + response.extensions.extend(result.extensions); + } + Err(error) => { + log_batch_write_error(&error, index); + response.failures.push((index, error)); } - - Ok(RegionResponse { - affected_rows, - extensions, - metadata: Vec::new(), - }) } } + Ok(response) } async fn handle_requests_in_serial( @@ -592,65 +589,35 @@ impl RegionServer { }) } - /// Attempts to optimize batch Put requests for metric engine. - /// - /// Returns Either::Left(response) if optimization succeeded, - /// or Either::Right(original_requests) to fall back to parallel processing. - /// - /// This avoids cloning requests when optimization cannot be applied. + /// Attempts metric batch writes, returning the requests when ordinary execution is needed. async fn try_handle_metric_batch_puts( &self, - requests: Vec<(RegionId, RegionRequest)>, - ) -> Result>> { - if requests.is_empty() { + requests: Vec, + ) -> Result, Vec>> { + let Some(first) = requests.first() else { return Ok(Either::Right(requests)); - } - - // Quick check: verify first request is Put and is metric engine - if !matches!(requests[0].1, RegionRequest::Put(_)) { - return Ok(Either::Right(requests)); - } - let first_region_id = requests[0].0; - let request_type = requests[0].1.request_type(); - - // SAFETY: If the first request belongs to metric engine, then ALL requests - // in this batch are guaranteed to belong to metric engine. This invariant - // is maintained by the request batching logic upstream. - let engine = match self - .inner - .get_engine(first_region_id, &RegionChange::None)? - { - CurrentEngine::Engine(e) => e, - _ => return Ok(Either::Right(requests)), }; - - if engine.name() != METRIC_ENGINE_NAME { + if !matches!(first.request, RegionRequest::Put(_)) { return Ok(Either::Right(requests)); } - // Check if ALL requests are Put (now we know it's worth checking) - let mut all_puts = true; - for (_, req) in &requests { - if !matches!(req, RegionRequest::Put(_)) { - all_puts = false; - break; - } - } + let first_region_id = first.region_id; + let request_type = first.request.request_type(); - if !all_puts { + // Upstream batching keeps requests from the same engine together. + // Ordinary execution reports engine lookup failures per entry. + let Ok(CurrentEngine::Engine(engine)) = + self.inner.get_engine(first_region_id, &RegionChange::None) + else { + return Ok(Either::Right(requests)); + }; + if engine.name() != METRIC_ENGINE_NAME + || !requests + .iter() + .all(|request| matches!(request.request, RegionRequest::Put(_))) + { return Ok(Either::Right(requests)); } - - // Now extract Put requests by consuming ownership (zero clone!) - let put_requests = requests.into_iter().map(|(region_id, req)| { - if let RegionRequest::Put(put) = req { - (region_id, put) - } else { - unreachable!("Already checked all are Put") - } - }); - - // Downcast to MetricEngine and call batch API let metric_engine = engine .as_any() .downcast_ref::() @@ -659,39 +626,65 @@ impl RegionServer { })? .clone(); + let indices: Vec<_> = requests + .iter() + .map(|request| (request.index, request.region_id)) + .collect(); + let put_requests = requests.into_iter().map(|request| match request.request { + RegionRequest::Put(put) => (request.region_id, put), + _ => unreachable!("Already checked all are Put"), + }); let tracing_context = TracingContext::from_current_span(); let batch_size = put_requests.len(); let span = tracing_context.attach(info_span!( "RegionServer::handle_metric_batch_puts", - batch_size = batch_size, + batch_size = batch_size )); - let result = common_runtime::spawn_ingest(async move { + let results = common_runtime::spawn_ingest(async move { metric_engine .put_regions_batch(put_requests) .trace(span) .await }) .await - .context(RuntimeJoinSnafu { request_type })? - .map_err(BoxedError::new) - .context(HandleRegionRequestSnafu { - region_id: first_region_id, - }); - - match result { - Ok(total_affected) => { - crate::metrics::REGION_CHANGED_ROW_COUNT - .with_label_values(&[request_type]) - .inc_by(total_affected as u64); - Ok(Either::Left(RegionResponse::new(total_affected))) - } - Err(err) => { - crate::metrics::REGION_SERVER_INSERT_FAIL_COUNT - .with_label_values(&[request_type]) - .inc_by(batch_size as u64); - Err(err) + .context(RuntimeJoinSnafu { request_type })?; + let mut response = BatchRegionResponse { + affected_rows: 0, + extensions: HashMap::new(), + failures: Vec::new(), + }; + for group in results { + let region_id = indices[group.request_indices[0]].1; + let result = group + .result + .map_err(BoxedError::new) + .context(HandleRegionRequestSnafu { region_id }) + .map(RegionResponse::new); + match result { + Ok(result) => { + crate::metrics::REGION_CHANGED_ROW_COUNT + .with_label_values(&[request_type]) + .inc_by(result.affected_rows as u64); + response.affected_rows += result.affected_rows; + } + Err(error) => { + crate::metrics::REGION_SERVER_INSERT_FAIL_COUNT + .with_label_values(&[request_type]) + .inc_by(group.request_indices.len() as u64); + log_batch_write_error(&error, indices[group.request_indices[0]].0); + let error = Arc::new(error); + response + .failures + .extend(group.request_indices.into_iter().map(|index| { + ( + indices[index].0, + error::BatchWriteSnafu.into_error(Arc::clone(&error)), + ) + })); + } } } + Ok(Either::Left(response)) } async fn handle_sync_region_request(&self, request: &SyncRequest) -> Result { @@ -914,16 +907,49 @@ impl Stream for RegionWatermarkStream { #[async_trait] impl RegionServerHandler for RegionServer { + async fn handle_region_request( + &self, + request: RegionRequestV1, + ) -> ServerResult { + let accept_partial_result = request + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .and_then(|context| context.extensions.get(api::region::ACCEPT_PARTIAL_RESULT)) + .is_some_and(|value| value == "true"); + let body = request.body.context(servers_error::InvalidQuerySnafu { + reason: "Expecting non-empty region request body.", + })?; + self.handle_inner(body, accept_partial_result).await + } + async fn handle(&self, request: region_request::Body) -> ServerResult { + self.handle_inner(request, false).await + } +} + +impl RegionServer { + async fn handle_inner( + &self, + request: region_request::Body, + accept_partial_result: bool, + ) -> ServerResult { + let request_count = match &request { + region_request::Body::Inserts(requests) => Some(requests.requests.len()), + region_request::Body::Deletes(requests) => Some(requests.requests.len()), + _ => None, + }; + if let Some(request_count) = request_count { + return self + .handle_write_request(request, request_count, accept_partial_result) + .await; + } let failed_requests_cnt = crate::metrics::REGION_SERVER_REQUEST_FAILURE_COUNT .with_label_values(&[request.as_ref()]); let response = match &request { region_request::Body::Creates(_) | region_request::Body::Drops(_) | region_request::Body::Alters(_) => self.handle_batch_ddl_requests(request).await, - region_request::Body::Inserts(_) | region_request::Body::Deletes(_) => { - self.handle_requests_in_parallel(request).await - } region_request::Body::Sync(sync_request) => { self.handle_sync_region_request(sync_request).await } @@ -938,11 +964,8 @@ impl RegionServerHandler for RegionServer { _ => self.handle_requests_in_serial(request).await, } .map_err(BoxedError::new) - .inspect_err(|_| { - failed_requests_cnt.inc(); - }) + .inspect_err(|_| failed_requests_cnt.inc()) .context(ExecuteGrpcRequestSnafu)?; - Ok(RegionResponseV1 { header: Some(ResponseHeader { status: Some(Status { @@ -953,8 +976,35 @@ impl RegionServerHandler for RegionServer { affected_rows: response.affected_rows as _, extensions: response.extensions, metadata: response.metadata, + failures: Vec::new(), }) } + + async fn handle_write_request( + &self, + request: region_request::Body, + entry_count: usize, + accept_partial_result: bool, + ) -> ServerResult { + let failed_requests_cnt = crate::metrics::REGION_SERVER_REQUEST_FAILURE_COUNT + .with_label_values(&[request.as_ref()]); + let mut result = self + .handle_requests_in_parallel(request) + .await + .map_err(BoxedError::new) + .inspect_err(|_| failed_requests_cnt.inc()) + .context(ExecuteGrpcRequestSnafu)?; + if !result.failures.is_empty() { + failed_requests_cnt.inc(); + // Old clients and single-entry writes retain tonic errors and retry metadata. + if (!accept_partial_result || entry_count == 1) + && let Some((_, error)) = result.failures.drain(..).min_by_key(|(index, _)| *index) + { + return Err(BoxedError::new(error)).context(ExecuteGrpcRequestSnafu); + } + } + Ok(encode_batch_write_response(result)) + } } #[async_trait] @@ -1004,6 +1054,97 @@ impl FlightCraft for RegionServer { } } +/// A decoded write retains its protobuf entry index even when preceding entries are empty. +struct IndexedRegionRequest { + index: usize, + region_id: RegionId, + request: RegionRequest, +} + +fn prepare_batch_write_requests(body: region_request::Body) -> Vec { + match body { + region_request::Body::Inserts(inserts) => inserts + .requests + .into_iter() + .enumerate() + .filter_map(|(index, request)| { + store_region_request::make_region_put(request).map(|(region_id, request)| { + IndexedRegionRequest { + index, + region_id, + request: RegionRequest::Put(request), + } + }) + }) + .collect(), + region_request::Body::Deletes(deletes) => deletes + .requests + .into_iter() + .enumerate() + .filter_map(|(index, request)| { + store_region_request::make_region_delete(request).map(|(region_id, request)| { + IndexedRegionRequest { + index, + region_id, + request: RegionRequest::Delete(request), + } + }) + }) + .collect(), + _ => unreachable!("Only inserts/deletes are prepared as batch writes"), + } +} + +fn encode_batch_write_response(result: BatchRegionResponse) -> RegionResponseV1 { + RegionResponseV1 { + header: Some(ResponseHeader { + status: Some(Status { + status_code: if result.failures.is_empty() { + StatusCode::Success + } else { + StatusCode::PartialFailure + } as u32, + err_msg: if let Some((index, error)) = result.first_failure() { + format!( + "{} region requests failed; entry {index}: {}", + result.failures.len(), + error.output_msg() + ) + } else { + String::new() + }, + ..Default::default() + }), + }), + metadata: Vec::new(), + affected_rows: result.affected_rows as _, + extensions: result.extensions, + failures: result + .failures + .into_iter() + .map(|(index, error)| RegionRequestFailure { + request_index: index as u32, + status: Some(Status { + status_code: error.status_code() as u32, + err_msg: error.output_msg(), + retry_hint: error.retry_hint().as_str().to_string(), + }), + }) + .collect(), + } +} + +fn log_batch_write_error(error: &error::Error, index: usize) { + if error.status_code().should_log_error() { + error!(error; "Failed to handle batch entry, request_index: {}", index); + } else { + debug!( + "Failed to handle batch entry, request_index: {}, err: {}", + index, error + ); + } +} + #[derive(Clone)] enum RegionEngineWithStatus { // An opening, or creating region. @@ -1507,7 +1648,7 @@ impl RegionServerInner { // Handle requests in batch. // // limitation: all create requests must be in the same engine. - pub async fn handle_batch_request( + pub async fn handle_batch_ddl_requests( &self, batch_request: BatchRegionDdlRequest, ) -> Result { @@ -1988,6 +2129,7 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; + use api::v1::region::{InsertRequest, InsertRequests}; use api::v1::{Rows, SemanticType}; use common_error::ext::{ErrorExt, RetryHint}; use common_recordbatch::RecordBatches; @@ -1996,7 +2138,8 @@ mod tests { use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::Int32Vector; use futures_util::StreamExt; - use mito2::test_util::CreateRequestBuilder; + use mito2::config::MitoConfig; + use mito2::test_util::{CreateRequestBuilder, TestEnv, rows_schema}; use query::options::FLOW_RETURN_REGION_SEQ; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::region_engine::RegionEngine; @@ -2392,6 +2535,170 @@ mod tests { assert_eq!(err.status_code(), StatusCode::RegionNotReady); } + #[test] + fn test_batch_write_summary_uses_first_input_failure() { + let first = error::RegionNotFoundSnafu { + region_id: RegionId::new(1, 0), + } + .build(); + let expected = format!("2 region requests failed; entry 0: {}", first.output_msg()); + let later = error::RegionBusySnafu { + region_id: RegionId::new(1, 1), + } + .build(); + let response = encode_batch_write_response(BatchRegionResponse { + affected_rows: 0, + extensions: HashMap::new(), + failures: vec![(1, later), (0, first)], + }); + let status = response.header.unwrap().status.unwrap(); + assert_eq!(status.status_code, StatusCode::PartialFailure as u32); + assert_eq!(status.err_msg, expected); + assert_eq!(response.failures[0].request_index, 1); + assert_eq!(response.failures[1].request_index, 0); + } + + #[tokio::test] + async fn test_write_failure_response_compatibility() { + let server = mock_region_server(); + let missing = RegionId::new(1, 1); + let registering = RegionId::new(1, 2); + let (engine, _receiver) = MockRegionEngine::new(MITO_ENGINE_NAME); + server + .inner + .region_map + .insert(registering, RegionEngineWithStatus::Registering(engine)); + for (region_id, code) in [ + (missing, StatusCode::RegionNotFound), + (registering, StatusCode::RegionNotReady), + ] { + for delete in [false, true] { + for accept_partial in [false, true] { + // A preceding empty entry must still count toward the original envelope size. + for with_empty_entry in [false, true] { + let rows = Some(Rows { + schema: Vec::new(), + rows: Vec::new(), + }); + let body = if delete { + let request = api::v1::region::DeleteRequest { + region_id: region_id.as_u64(), + rows, + ..Default::default() + }; + region_request::Body::Deletes(api::v1::region::DeleteRequests { + requests: if with_empty_entry { + vec![Default::default(), request] + } else { + vec![request] + }, + }) + } else { + let request = InsertRequest { + region_id: region_id.as_u64(), + rows, + ..Default::default() + }; + region_request::Body::Inserts(InsertRequests { + requests: if with_empty_entry { + vec![Default::default(), request] + } else { + vec![request] + }, + }) + }; + let mut header = api::v1::region::RegionRequestHeader::default(); + if accept_partial { + header + .query_context + .get_or_insert_with(Default::default) + .extensions + .insert( + api::region::ACCEPT_PARTIAL_RESULT.to_string(), + "true".to_string(), + ); + } + let result = server + .handle_region_request(RegionRequestV1 { + header: Some(header), + body: Some(body), + }) + .await; + if accept_partial && with_empty_entry { + let response = result.unwrap(); + assert_eq!( + StatusCode::PartialFailure as u32, + response.header.unwrap().status.unwrap().status_code + ); + assert_eq!(1, response.failures.len()); + assert_eq!(1, response.failures[0].request_index); + let status = response.failures[0].status.as_ref().unwrap(); + assert_eq!(code as u32, status.status_code); + assert_eq!( + RetryHint::RetryAfterRouteRefresh.as_str(), + status.retry_hint + ); + } else { + let status: tonic::Status = result.unwrap_err().into(); + assert_eq!( + common_error::status_code::status_to_tonic_code(code), + status.code() + ); + assert_eq!( + (code as u32).to_string(), + status + .metadata() + .get(common_error::GREPTIME_DB_HEADER_ERROR_CODE) + .unwrap() + .to_str() + .unwrap() + ); + assert_eq!( + RetryHint::RetryAfterRouteRefresh.as_str(), + status + .metadata() + .get(common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT) + .unwrap() + .to_str() + .unwrap() + ); + } + } + } + } + } + } + + #[tokio::test] + async fn test_legacy_batch_returns_first_input_error() { + let server = mock_region_server(); + let first = RegionId::new(1, 1); + let second = RegionId::new(1, 2); + let (engine, _receiver) = MockRegionEngine::with_custom_apply_fn("mock", |engine| { + engine.handle_request_delay = Some(Duration::from_millis(10)); + engine.handle_request_mock_fn = Some(Box::new(|region_id, _| { + error::RegionBusySnafu { region_id }.fail() + })); + }); + server.register_test_region(first, engine); + let requests = [first, second] + .map(|region_id| InsertRequest { + region_id: region_id.as_u64(), + rows: Some(Rows::default()), + ..Default::default() + }) + .to_vec(); + let error = server + .handle_region_request(RegionRequestV1 { + header: Some(Default::default()), + body: Some(region_request::Body::Inserts(InsertRequests { requests })), + }) + .await + .unwrap_err(); + assert_eq!(StatusCode::RegionBusy, error.status_code()); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + } + #[tokio::test] async fn test_region_request_failed() { common_telemetry::init_default_ut_logging(); @@ -2417,8 +2724,7 @@ mod tests { .await .unwrap_err(); - let status = mock_region_server.inner.region_map.get(®ion_id); - assert!(status.is_none()); + assert!(!mock_region_server.inner.region_map.contains_key(®ion_id)); mock_region_server .inner @@ -2437,8 +2743,79 @@ mod tests { .await .unwrap_err(); - let status = mock_region_server.inner.region_map.get(®ion_id); - assert!(status.is_some()); + assert!(mock_region_server.inner.region_map.contains_key(®ion_id)); + + let insert = api::v1::region::InsertRequest { + region_id: region_id.as_u64(), + rows: Some(Rows { + schema: Vec::new(), + rows: Vec::new(), + }), + ..Default::default() + }; + let success_region = RegionId::new(1, 2); + let (success_engine, _receiver) = + MockRegionEngine::with_custom_apply_fn("mock", |engine| { + engine.handle_request_delay = Some(std::time::Duration::from_millis(10)); + engine.handle_request_mock_fn = Some(Box::new(|_, _| Ok(5))); + }); + mock_region_server.register_test_region(success_region, success_engine); + let successful_insert = api::v1::region::InsertRequest { + region_id: success_region.as_u64(), + ..insert.clone() + }; + let response = mock_region_server + .handle_region_request(RegionRequestV1 { + header: Some(api::v1::region::RegionRequestHeader { + query_context: Some(api::v1::QueryContext { + extensions: HashMap::from([( + api::region::ACCEPT_PARTIAL_RESULT.to_string(), + "true".to_string(), + )]), + ..Default::default() + }), + ..Default::default() + }), + body: Some(region_request::Body::Inserts(InsertRequests { + requests: vec![ + Default::default(), + insert.clone(), + successful_insert, + insert, + ], + })), + }) + .await + .unwrap(); + assert_eq!(response.affected_rows, 5); + assert_eq!( + response + .header + .as_ref() + .unwrap() + .status + .as_ref() + .unwrap() + .status_code, + StatusCode::PartialFailure as u32 + ); + assert_eq!( + response + .failures + .iter() + .map(|error| error.request_index) + .collect::>(), + vec![1, 3] + ); + assert!( + response + .failures + .iter() + .all(|error| error.status.as_ref().is_some_and(|status| { + status.status_code == StatusCode::Unexpected as u32 + && status.retry_hint == RetryHint::NonRetryable.as_str() + })) + ); } #[tokio::test] @@ -2912,4 +3289,112 @@ mod tests { .await .unwrap_err(); } + + #[tokio::test] + async fn test_metric_batch_failure_indices() { + let mut env = TestEnv::with_prefix("region-server-metric-results").await; + let mito = env.create_engine(MitoConfig::default()).await; + let metric = Arc::new( + MetricEngine::try_new(mito, metric_engine::config::EngineConfig::default()).unwrap(), + ); + let mut server = mock_region_server(); + server.register_engine(metric.clone()); + let mut inserts = Vec::new(); + for table in [10, 20] { + let physical = RegionId::new(table, 0); + let logical = RegionId::new(table + 1, 0); + let mut create = CreateRequestBuilder::new().build(); + create.engine = METRIC_ENGINE_NAME.to_string(); + create.column_metadatas = vec![ + store_api::metadata::ColumnMetadata { + column_id: 0, + semantic_type: SemanticType::Timestamp, + column_schema: ColumnSchema::new( + "greptime_timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + }, + store_api::metadata::ColumnMetadata { + column_id: 1, + semantic_type: SemanticType::Field, + column_schema: ColumnSchema::new( + "greptime_value", + ConcreteDataType::float64_datatype(), + false, + ), + }, + ]; + create.primary_key.clear(); + create.table_dir = format!("metric-results/{table}/"); + create.options = HashMap::from([( + store_api::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY.to_string(), + String::new(), + )]); + server + .handle_request(physical, RegionRequest::Create(create.clone())) + .await + .unwrap(); + create.table_dir = format!("metric-results/{}/", table + 1); + create.options = HashMap::from([( + store_api::metric_engine_consts::LOGICAL_TABLE_METADATA_KEY.to_string(), + physical.as_u64().to_string(), + )]); + server + .handle_request(logical, RegionRequest::Create(create.clone())) + .await + .unwrap(); + inserts.push(InsertRequest { + region_id: logical.as_u64(), + rows: Some(Rows { + schema: rows_schema(&create), + rows: vec![api::v1::Row { + values: vec![ + api::v1::Value { + value_data: Some( + api::v1::value::ValueData::TimestampMillisecondValue(1), + ), + }, + api::v1::Value { + value_data: Some(api::v1::value::ValueData::F64Value(1.0)), + }, + ], + }], + }), + ..Default::default() + }); + } + // Both entries of the second physical group fail at write_data. + metric + .set_region_role(RegionId::new(20, 0), RegionRole::Follower) + .unwrap(); + let response = server + .handle_requests_in_parallel(region_request::Body::Inserts(InsertRequests { + requests: vec![ + Default::default(), + inserts[0].clone(), + inserts[1].clone(), + inserts[1].clone(), + ], + })) + .await + .unwrap(); + assert_eq!(1, response.affected_rows); + let mut failures = response.failures; + failures.sort_unstable_by_key(|(index, _)| *index); + assert_eq!( + vec![2, 3], + failures.iter().map(|(index, _)| *index).collect::>() + ); + assert!( + failures + .iter() + .all(|(_, error)| error.status_code() == StatusCode::RegionNotReady) + ); + assert!( + failures + .iter() + .all(|(_, error)| error.retry_hint() == RetryHint::RetryAfterRouteRefresh) + ); + } } diff --git a/src/meta-srv/src/procedure/utils.rs b/src/meta-srv/src/procedure/utils.rs index 49964de24b7..ba8e31d5266 100644 --- a/src/meta-srv/src/procedure/utils.rs +++ b/src/meta-srv/src/procedure/utils.rs @@ -356,11 +356,13 @@ pub mod mock { status: Some(PbStatus { status_code: 0, err_msg: String::default(), + ..Default::default() }), }), affected_rows: 0, extensions: Default::default(), metadata: Vec::new(), + failures: Vec::new(), }) } } diff --git a/src/metric-engine/AGENTS.md b/src/metric-engine/AGENTS.md index d1049478978..9c3820a0cc3 100644 --- a/src/metric-engine/AGENTS.md +++ b/src/metric-engine/AGENTS.md @@ -51,6 +51,7 @@ Mito2 data region. Queries against a physical region pass straight through. ## Public surface - Entry: `MetricEngine` in `src/metric-engine/src/engine.rs`, built via `try_new(mito, config)`. +- `MetricEngine::put_regions_batch` returns each physical group's outcome with its input indices; a failed group does not stop other groups. - Trait: `impl RegionEngine for MetricEngine` (name = `"metric"`). - Depends on `mito2`, `store-api`, and `mito-codec` (sparse primary key codec). diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index 18d6ff6f197..0d5cb28c2cf 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -488,6 +488,17 @@ impl RegionEngine for MetricEngine { } } +/// Outcome of a group of input requests. +/// +/// Requests rejected before physical-region grouping produce singleton failures. +#[derive(Debug)] +pub struct PutGroupResult { + /// Positions in the input batch, including separate entries for the same region. + pub request_indices: Vec, + /// Aggregate rows on success; no other physical group's outcome is implied. + pub result: Result, +} + impl MetricEngine { pub fn try_new(mito: MitoEngine, mut config: EngineConfig) -> Result { let metadata_region = MetadataRegion::new(mito.clone()); @@ -520,13 +531,11 @@ impl MetricEngine { self.inner.mito.clone() } - /// Batch put operation for multiple logical regions. - /// Requests are grouped by physical region; a failure can leave earlier - /// physical-region groups committed. + /// Writes all physical-region groups and retains each group's outcome. pub async fn put_regions_batch( &self, requests: impl ExactSizeIterator, - ) -> Result { + ) -> Vec { self.inner.put_regions_batch(requests).await } diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index d3ce895ff8a..ef18b7f8c39 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/src/metric-engine/src/engine/put.rs @@ -29,7 +29,7 @@ use store_api::region_request::{ }; use store_api::storage::{RegionId, TableId}; -use crate::engine::MetricEngineInner; +use crate::engine::{MetricEngineInner, PutGroupResult}; use crate::error::{ ColumnNotFoundSnafu, CreateDefaultSnafu, ForbiddenPhysicalWriteSnafu, InvalidRequestSnafu, LogicalRegionNotFoundSnafu, PhysicalRegionNotFoundSnafu, Result, UnexpectedRequestSnafu, @@ -39,6 +39,12 @@ use crate::metrics::{FORBIDDEN_OPERATION_COUNT, MITO_OPERATION_ELAPSED}; use crate::row_modifier::{RowsIter, TableIdInput}; use crate::utils::to_data_region_id; +#[derive(Default)] +struct PutGroup { + request_indices: Vec, + requests: Vec<(RegionId, RegionPutRequest)>, +} + impl MetricEngineInner { /// Dispatch region put request pub async fn put_region( @@ -61,67 +67,63 @@ impl MetricEngineInner { } } - /// Batch write multiple logical regions to the same physical region. - /// - /// Dispatch region put requests in batch. - /// - /// Requests may span multiple physical regions. We group them by physical - /// region and write sequentially. This method fails fast on validation or - /// preparation errors within a group and stops at the first failure. - /// Writes in earlier physical-region groups are not rolled back if a later - /// group fails. + /// Retains every group's outcome instead of returning on the first failure. pub async fn put_regions_batch( &self, - requests: impl ExactSizeIterator, - ) -> Result { + mut requests: impl ExactSizeIterator, + ) -> Vec { let len = requests.len(); - if len == 0 { - return Ok(0); + return Vec::new(); } - let _timer = MITO_OPERATION_ELAPSED .with_label_values(&["put_batch"]) .start_timer(); - - // Fast path: single request, no batching overhead if len == 1 { - let (region_id, req) = requests.into_iter().next().unwrap(); - let is_putting_physical_region = - self.state.read().unwrap().exist_physical_region(region_id); - if is_putting_physical_region { - FORBIDDEN_OPERATION_COUNT.inc(); - return ForbiddenPhysicalWriteSnafu.fail(); + let (region_id, request) = requests.next().unwrap(); + return vec![PutGroupResult { + request_indices: vec![0], + result: self.put_region(region_id, request).await, + }]; + } + let mut groups = HashMap::new(); + let mut results = Vec::new(); + for (index, (region_id, request)) in requests.enumerate() { + if let Err(error) = self.group_put_request(&mut groups, index, region_id, request) { + results.push(PutGroupResult { + request_indices: vec![index], + result: Err(error), + }); } - - return self.put_logical_region(region_id, req).await; } - - let mut requests_per_physical: HashMap> = - HashMap::new(); - for (region_id, request) in requests { - let is_putting_physical_region = - self.state.read().unwrap().exist_physical_region(region_id); - if is_putting_physical_region { - FORBIDDEN_OPERATION_COUNT.inc(); - return ForbiddenPhysicalWriteSnafu.fail(); - } - let physical_region_id = self.find_physical_region_id(region_id)?; - requests_per_physical - .entry(physical_region_id) - .or_default() - .push((region_id, request)); + for (physical_region_id, group) in groups { + let result = self + .put_regions_batch_single_physical(physical_region_id, group.requests) + .await; + results.push(PutGroupResult { + request_indices: group.request_indices, + result, + }); } + results + } - let mut total_affected_rows: AffectedRows = 0; - for (physical_region_id, requests) in requests_per_physical { - let affected_rows = self - .put_regions_batch_single_physical(physical_region_id, requests) - .await?; - total_affected_rows += affected_rows; + fn group_put_request( + &self, + groups: &mut HashMap, + index: usize, + region_id: RegionId, + request: RegionPutRequest, + ) -> Result<()> { + if self.state.read().unwrap().exist_physical_region(region_id) { + FORBIDDEN_OPERATION_COUNT.inc(); + return ForbiddenPhysicalWriteSnafu.fail(); } - - Ok(total_affected_rows) + let physical_region_id = self.find_physical_region_id(region_id)?; + let group = groups.entry(physical_region_id).or_default(); + group.request_indices.push(index); + group.requests.push((region_id, request)); + Ok(()) } /// Write a batch of requests that all belong to the same physical region. @@ -999,7 +1001,9 @@ mod tests { }, ) }); - let affected_rows = engine.inner.put_regions_batch(requests).await.unwrap(); + let results = engine.put_regions_batch(requests).await; + assert_eq!(results.len(), 1); + let affected_rows = results.into_iter().next().unwrap().result.unwrap(); assert_eq!(affected_rows, 6); assert_eq!( scan_timestamp_values(&engine, logical_region_id).await, @@ -1289,12 +1293,13 @@ mod tests { for ((_, request), skip_wal) in mixed_requests.iter_mut().zip(policies) { request.skip_wal = skip_wal; } - let err = env + let results = env .metric() - .inner .put_regions_batch(mixed_requests.into_iter()) - .await - .unwrap_err(); + .await; + assert_eq!(results.len(), 1); + assert_eq!(results[0].request_indices, vec![0, 1]); + let err = results.into_iter().next().unwrap().result.unwrap_err(); assert!(err.to_string().contains("inconsistent WAL policy in batch")); for logical_region_id in [logical_region_1, logical_region_2] { assert!( @@ -1305,12 +1310,12 @@ mod tests { } } - let affected_rows = env + let results = env .metric() - .inner .put_regions_batch(build_requests().into_iter()) - .await - .unwrap(); + .await; + assert_eq!(results.len(), 1); + let affected_rows = results.into_iter().next().unwrap().result.unwrap(); assert_eq!(affected_rows, 5); let request = ScanRequest::default(); @@ -1649,11 +1654,9 @@ mod tests { ]; // Batch write - let affected_rows = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap(); + let results = engine.put_regions_batch(requests.into_iter()).await; + assert_eq!(results.len(), 1); + let affected_rows = results.into_iter().next().unwrap().result.unwrap(); assert_eq!(affected_rows, 10); // Verify physical region contains data from all logical regions @@ -1673,72 +1676,70 @@ mod tests { async fn test_batch_write_with_partial_failure() { let env = TestEnv::new().await; env.init_metric_region().await; - let engine = env.metric(); - let physical_region_id = env.default_physical_region_id(); - let logical_region_1 = env.default_logical_region_id(); - let logical_region_2 = RegionId::new(1024, 1); - let nonexistent_region = RegionId::new(9999, 9999); - - env.create_logical_region(physical_region_id, logical_region_2) + let physical = env.default_physical_region_id(); + let logical = env.default_logical_region_id(); + let other_physical = RegionId::new(4096, 0); + let other_logical = RegionId::new(4097, 0); + env.create_physical_region(other_physical, "batch_results", vec![]) .await; - - // Prepare batch with one invalid region - let schema = test_util::row_schema_with_tags(&["job"]); - let requests = vec![ - ( - logical_region_1, - RegionPutRequest { - skip_wal: false, - rows: Rows { - schema: schema.clone(), - rows: test_util::build_rows(1, 3), - }, - hint: None, - partition_expr_version: None, + env.create_logical_region(other_physical, other_logical) + .await; + let request = |index: i64, version| { + let mut rows = test_util::build_rows(1, 1); + rows[0].values[0].value_data = + Some(api::v1::value::ValueData::TimestampMillisecondValue(index)); + RegionPutRequest { + rows: Rows { + schema: test_util::row_schema_with_tags(&["job"]), + rows, }, - ), - ( - nonexistent_region, - RegionPutRequest { - skip_wal: false, - rows: Rows { - schema: schema.clone(), - rows: test_util::build_rows(1, 2), - }, - hint: None, - partition_expr_version: None, - }, - ), - ( - logical_region_2, - RegionPutRequest { - skip_wal: false, - rows: Rows { - schema: schema.clone(), - rows: test_util::build_rows(1, 5), - }, - hint: None, - partition_expr_version: None, - }, - ), - ]; - - // Batch write - let result = engine.inner.put_regions_batch(requests.into_iter()).await; - assert!(result.is_err()); - - // Invalid region is detected before any write, so the physical region remains empty. - // Fail-fast is per physical-region group; cross-group partial success is possible. - let request = ScanRequest::default(); - let stream = env + skip_wal: false, + hint: None, + partition_expr_version: version, + } + }; + let mut results = env .metric() - .scan_to_stream(physical_region_id, request) - .await - .unwrap(); - let batches = RecordBatches::try_collect(stream).await.unwrap(); - - assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 0); + .put_regions_batch( + vec![ + (logical, request(0, None)), + (RegionId::new(9999, 0), request(1, None)), + (other_logical, request(2, Some(1))), + (other_logical, request(3, None)), + (logical, request(4, None)), + (physical, request(5, None)), + ] + .into_iter(), + ) + .await; + results.sort_by_key(|result| result.request_indices[0]); + assert_eq!( + results + .iter() + .map(|result| result.request_indices.clone()) + .collect::>(), + vec![vec![0, 4], vec![1], vec![2, 3], vec![5]] + ); + assert_eq!(results[0].result.as_ref().unwrap(), &2); + assert!(results[1].result.is_err()); + assert!(results[2].result.is_err()); + assert!(matches!( + results[3].result, + Err(crate::error::Error::ForbiddenPhysicalWrite { .. }) + )); + for (region, expected) in [(physical, 2), (other_physical, 0)] { + let stream = env + .metric() + .scan_to_stream(region, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + expected + ); + } } #[tokio::test] @@ -1762,11 +1763,10 @@ mod tests { }, )]; - let err = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap_err(); + let results = engine.put_regions_batch(requests.into_iter()).await; + assert_eq!(results.len(), 1); + assert_eq!(results[0].request_indices, vec![0]); + let err = results.into_iter().next().unwrap().result.unwrap_err(); assert!(matches!( err, @@ -1810,15 +1810,15 @@ mod tests { ), ]; - let err = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap_err(); - + let mut results = engine.put_regions_batch(requests.into_iter()).await; + results.sort_by_key(|group| group.request_indices[0]); + assert_eq!(results.len(), 2); + assert_eq!(results[0].request_indices, vec![0]); + assert_eq!(results[0].result.as_ref().unwrap(), &1); + assert_eq!(results[1].request_indices, vec![1]); assert!(matches!( - err, - crate::error::Error::ForbiddenPhysicalWrite { .. } + results[1].result, + Err(crate::error::Error::ForbiddenPhysicalWrite { .. }) )); } @@ -1845,11 +1845,10 @@ mod tests { }, )]; - let affected_rows = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap(); + let results = engine.put_regions_batch(requests.into_iter()).await; + assert_eq!(results.len(), 1); + assert_eq!(results[0].request_indices, vec![0]); + let affected_rows = results.into_iter().next().unwrap().result.unwrap(); assert_eq!(affected_rows, 5); } @@ -1859,15 +1858,12 @@ mod tests { env.init_metric_region().await; let engine = env.metric(); - // Empty batch should return zero affected rows - let requests = vec![]; - let affected_rows = engine - .inner - .put_regions_batch(requests.into_iter()) - .await - .unwrap(); - - assert_eq!(affected_rows, 0); + assert!( + engine + .put_regions_batch(vec![].into_iter()) + .await + .is_empty() + ); } #[tokio::test] diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 3d4c5f91f38..669fbfc63ce 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -1668,6 +1668,7 @@ impl ErrorExt for Error { match self { PartitionExprVersionMismatch { .. } | RegionNotFound { .. } + | RegionClosed { .. } | RejectWriteInStaging { .. } | RegionState { state: RegionRoleState::Follower, @@ -1805,6 +1806,21 @@ mod tests { StatusCode::RegionNotFound, RetryHint::RetryAfterRouteRefresh, ), + ( + RegionClosedSnafu { region_id }.build(), + StatusCode::Cancelled, + RetryHint::RetryAfterRouteRefresh, + ), + ( + RegionDroppedSnafu { region_id }.build(), + StatusCode::Cancelled, + RetryHint::NonRetryable, + ), + ( + RegionTruncatedSnafu { region_id }.build(), + StatusCode::Cancelled, + RetryHint::NonRetryable, + ), ] { assert_eq!(error.status_code(), status); assert_eq!(error.retry_hint(), hint); diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index 6599ccbe2bb..17c451eceb8 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -1622,7 +1622,7 @@ impl FlushStatus { #[cfg(test)] mod tests { use api::v1::{OpType, Rows}; - use common_error::ext::ErrorExt; + use common_error::ext::{ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use datatypes::arrow::datatypes::Schema; use datatypes::arrow::record_batch::RecordBatch; @@ -2062,28 +2062,39 @@ mod tests { } #[tokio::test] - async fn test_region_closed_notifies_pending_bulk_writes() { + async fn test_region_closed_notifies_pending_requests() { let region_id = RegionId::new(1, 1); let version_control = Arc::new(VersionControlBuilder::new().build()); - let (bulk_req, output_rx) = new_test_bulk_request(region_id); + let (write_req, write_rx) = new_test_write_request(region_id); + let (bulk_req, bulk_rx) = new_test_bulk_request(region_id); + let (ddl_tx, ddl_rx) = oneshot::channel(); let status = FlushStatus { region_id, version_control, state: CancellableTaskState::new(), pending_task: None, closing: false, - pending_ddls: Vec::new(), - pending_writes: Vec::new(), + pending_ddls: vec![SenderDdlRequest { + region_id, + sender: OptionOutputTx::from(ddl_tx), + request: DdlRequest::Truncate( + store_api::region_request::RegionTruncateRequest::All, + ), + }], + pending_writes: vec![write_req], pending_bulk_writes: vec![bulk_req], }; status.on_region_closed(Arc::new(RegionClosedSnafu { region_id }.build())); - let err = output_rx - .await - .expect("pending bulk write must receive explicit error") - .unwrap_err(); + for rx in [write_rx, bulk_rx] { + let err = rx.await.unwrap().unwrap_err(); + assert_eq!(err.status_code(), StatusCode::Cancelled); + assert_eq!(err.retry_hint(), RetryHint::RetryAfterRouteRefresh); + } + let err = ddl_rx.await.unwrap().unwrap_err(); assert_eq!(err.status_code(), StatusCode::Cancelled); + assert_eq!(err.retry_hint(), RetryHint::RetryAfterRouteRefresh); } #[tokio::test] diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs index ed63be04c39..b0af399d87d 100644 --- a/src/operator/src/write_dispatcher.rs +++ b/src/operator/src/write_dispatcher.rs @@ -27,7 +27,9 @@ use backon::{BackoffBuilder, ExponentialBuilder}; use catalog::CatalogManagerRef; use common_error::ext::{BoxedError, ErrorExt, RetryHint}; use common_meta::cache_invalidator::{CacheInvalidatorRef, Context}; -use common_meta::error::{Error, ExternalSnafu, Result, TableInfoNotFoundSnafu, UnexpectedSnafu}; +use common_meta::error::{ + BatchRequestSnafu, Error, ExternalSnafu, Result, TableInfoNotFoundSnafu, UnexpectedSnafu, +}; use common_meta::instruction::CacheIdent; use common_meta::node_manager::NodeManagerRef; use common_meta::peer::{Peer, PeerRef}; @@ -35,7 +37,7 @@ use common_meta::rpc::router; use futures::future; use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; use prometheus::Histogram; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; use store_api::storage::RegionId; use table::table_name::TableName; use tokio_util::sync::CancellationToken; @@ -180,24 +182,17 @@ impl WriteDispatcher { cancellation: Option<&CancellationToken>, ) -> Result { let histogram = histogram.as_ref(); - let writes = if self.retry.is_some() { - writes - .into_iter() - .flat_map(|(peer, request)| { - split_request(request) - .into_iter() - .map(move |request| (Arc::clone(&peer), request)) - }) - .collect() - } else { - writes - }; - let attempts = future::join_all( - writes - .into_iter() - .map(|(peer, request)| self.send(peer, request, histogram)), - ) - .await; + let attempts = future::join_all(writes.into_iter().map(|(peer, request)| async move { + if self.retry.is_some() { + self.send_batch(peer, request, histogram).await + } else { + vec![self.send(peer, request, histogram).await] + } + })) + .await + .into_iter() + .flatten() + .collect(); self.finish(attempts, histogram, cancellation).await } @@ -205,10 +200,13 @@ impl WriteDispatcher { pub(crate) async fn send( &self, peer: PeerRef, - request: RegionRequest, + mut request: RegionRequest, histogram: Option<&Histogram>, ) -> WriteAttempt { if let Some(retry) = &self.retry { + if matches!(&request.body, Some(Body::Inserts(_) | Body::Deletes(_))) { + api::region::set_accept_partial_result(request.header.get_or_insert_default()); + } return retry .send( RoutedWrite { @@ -266,6 +264,80 @@ impl WriteDispatcher { resolve_peers(&self.partition_manager, regions).await } + async fn send_batch( + &self, + peer: PeerRef, + mut request: RegionRequest, + histogram: Option<&Histogram>, + ) -> Vec { + if !matches!(&request.body, Some(Body::Inserts(_) | Body::Deletes(_))) { + return vec![self.send(peer, request, histogram).await]; + } + // Upgrade datanodes first and roll back frontends first: old datanodes + // cannot report which entries succeeded, so retrying a batch may replay writes. + api::region::set_accept_partial_result(request.header.get_or_insert_default()); + let request = Arc::new(request); + let requester = self.node_manager.datanode(&peer).await; + let result = { + let _timer = histogram.map(Histogram::start_timer); + requester.handle(Arc::clone(&request)).await + }; + let error = match result { + Ok(response) => return vec![WriteAttempt::Success(response)], + Err(error) => error, + }; + let request = match Arc::try_unwrap(request) { + Ok(request) => request, + Err(_) => { + return vec![WriteAttempt::Failed { + error: UnexpectedSnafu { + err_msg: "Region request is still shared after its RPC completed", + } + .build(), + request: None, + }]; + } + }; + let requests = split_request(request); + match error { + Error::PartialRegion { response } => { + let mut failures: HashMap<_, _> = response.failures.into_iter().collect(); + let mut attempts = vec![WriteAttempt::Success(RegionResponse { + affected_rows: response.affected_rows, + extensions: response.extensions, + metadata: Vec::new(), + })]; + for (index, request) in requests.into_iter().enumerate() { + if let Some(error) = failures.remove(&index) { + let request = error.retry_hint().is_retryable().then(|| RoutedWrite { + peer: Arc::clone(&peer), + request: Arc::new(request), + }); + attempts.push(WriteAttempt::Failed { error, request }); + } + } + attempts + } + error => { + // A whole-RPC failure cannot confirm success for any entry. + let error = Arc::new(error); + requests + .into_iter() + .map(|request| { + let request = error.retry_hint().is_retryable().then(|| RoutedWrite { + peer: Arc::clone(&peer), + request: Arc::new(request), + }); + WriteAttempt::Failed { + error: BatchRequestSnafu.into_error(Arc::clone(&error)), + request, + } + }) + .collect() + } + } + } + async fn dispatch_inner( &self, requests: Vec, diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs index d9a273d0c30..fd2d5b2452b 100644 --- a/src/operator/src/write_dispatcher/tests.rs +++ b/src/operator/src/write_dispatcher/tests.rs @@ -21,8 +21,8 @@ use std::time::{Duration, Instant}; use api::v1::region::region_server::{Region, RegionServer}; use api::v1::region::{ - AlignedSchemaVersion, BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader, - RegionResponse as PbRegionResponse, bulk_insert_request, + AlignedSchemaVersion, BulkInsertRequest, DeleteRequest, DeleteRequests, InsertRequest, + InsertRequests, RegionRequestHeader, RegionResponse as PbRegionResponse, bulk_insert_request, }; use api::v1::value::ValueData; use api::v1::{ @@ -77,6 +77,8 @@ use crate::tests::{create_partition_rule_manager, new_test_table_info, prepare_m #[derive(Clone, Default)] struct RecordingRegionService { + request_sizes: Arc>>, + rpc_rejections: Arc>>, requests: Arc>>, rejections: Arc>>>, release: Arc>>>, @@ -89,6 +91,70 @@ impl Region for RecordingRegionService { request: Request, ) -> std::result::Result, Status> { let request = request.into_inner(); + if !matches!(&request.body, Some(Body::Inserts(_) | Body::Deletes(_))) { + return self.handle_entry(request).await; + } + let accepts_partial = request + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .and_then(|context| context.extensions.get(api::region::ACCEPT_PARTIAL_RESULT)) + .is_some_and(|value| value == "true"); + let entries = split_request(request); + let return_tonic_error = !accepts_partial || entries.len() == 1; + self.request_sizes.lock().unwrap().push(entries.len()); + if let Some(status) = self.rpc_rejections.lock().unwrap().pop_front() { + self.requests.lock().unwrap().extend(entries); + return Err(status); + } + let mut response = PbRegionResponse { + header: Some(ResponseHeader { + status: Some(PbStatus::default()), + }), + ..Default::default() + }; + let results = + future::join_all(entries.into_iter().map(|entry| self.handle_entry(entry))).await; + for (index, result) in results.into_iter().enumerate() { + match result { + Ok(result) => response.affected_rows += result.into_inner().affected_rows, + Err(status) => { + if return_tonic_error { + return Err(status); + } + let error = client::Error::from(status); + response + .failures + .push(api::v1::region::RegionRequestFailure { + request_index: index as u32, + status: Some(PbStatus { + status_code: error.status_code() as u32, + err_msg: error.output_msg(), + retry_hint: error.retry_hint().as_str().to_string(), + }), + }); + } + } + } + if !response.failures.is_empty() { + response + .header + .as_mut() + .unwrap() + .status + .as_mut() + .unwrap() + .status_code = StatusCode::PartialFailure as u32; + } + Ok(Response::new(response)) + } +} + +impl RecordingRegionService { + async fn handle_entry( + &self, + request: RegionRequest, + ) -> std::result::Result, Status> { let region_id = write_region_id(&request).unwrap().as_u64(); self.requests.lock().unwrap().push(request); let release = self.release.lock().unwrap().clone(); @@ -236,9 +302,19 @@ impl RegionTestServer { hints .into_iter() .map(|hint| { - let mut status = Status::failed_precondition(format!( - "region {region_id} rejected the write" - )); + let code = match hint { + RetryHint::Retryable => StatusCode::RegionBusy, + RetryHint::RetryAfterRouteRefresh => StatusCode::Cancelled, + RetryHint::NonRetryable => StatusCode::InvalidArguments, + }; + let mut status = Status::new( + common_error::status_code::status_to_tonic_code(code), + format!("region {region_id} rejected the write"), + ); + status.metadata_mut().insert( + common_error::GREPTIME_DB_HEADER_ERROR_CODE, + (code as u32).to_string().parse().unwrap(), + ); status.metadata_mut().insert( common_error::GREPTIME_DB_HEADER_ERROR_RETRY_HINT, hint.as_str().parse().unwrap(), @@ -342,19 +418,36 @@ fn build_insert_request(regions: &[u32]) -> RegionRequest { #[tokio::test] async fn test_write_retry_disabled_and_enabled_ablation() { common_telemetry::init_default_ut_logging(); - for run in 0..3 { + for regions in [vec![1], vec![1, 2]] { for max_retries in [0, 1] { let server = RegionTestServer::start().await; server.fail_region(1, [RetryHint::Retryable]); let dispatcher = server.create_dispatcher(max_retries).await; - let result = dispatcher.dispatch(vec![build_insert_request(&[1])]).await; + let result = dispatcher + .dispatch(vec![build_insert_request(®ions)]) + .await; let attempts = server.attempt_counts()[&RegionId::new(1, 1).as_u64()]; common_telemetry::info!( - "write_retry_ablation run={run} max_retries={max_retries} success={} attempts={attempts}", + "write_retry_ablation regions={regions:?} max_retries={max_retries} success={} attempts={attempts}", result.is_ok() ); - assert_eq!(result.is_ok(), max_retries > 0); + if max_retries == 0 { + let error = result.unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RegionBusy); + assert_eq!(error.retry_hint(), RetryHint::Retryable); + } else { + assert_eq!(result.unwrap().affected_rows, regions.len()); + } assert_eq!(attempts, max_retries + 1); + for request in server.service.requests.lock().unwrap().iter() { + let accepts_partial = request + .header + .as_ref() + .and_then(|header| header.query_context.as_ref()) + .and_then(|context| context.extensions.get(api::region::ACCEPT_PARTIAL_RESULT)) + .is_some_and(|value| value == "true"); + assert_eq!(accepts_partial, max_retries > 0); + } assert!(server.invalidator.calls.lock().unwrap().is_empty()); } } @@ -362,22 +455,29 @@ async fn test_write_retry_disabled_and_enabled_ablation() { #[tokio::test] async fn test_write_retry_does_not_replay_successful_regions() { - let server = RegionTestServer::start().await; - server.fail_region(2, [RetryHint::Retryable]); - let dispatcher = server.create_dispatcher(2).await; - let response = dispatcher - .dispatch(vec![build_insert_request(&[1, 2])]) - .await - .unwrap(); - assert_eq!(response.affected_rows, 2); - assert_eq!( - server.attempt_counts(), - HashMap::from([ - (RegionId::new(1, 1).as_u64(), 1), - (RegionId::new(1, 2).as_u64(), 2), - ]) - ); - assert!(server.invalidator.calls.lock().unwrap().is_empty()); + for regions in [vec![1, 2], vec![1, 1, 1]] { + let server = RegionTestServer::start().await; + server.fail_region(1, [RetryHint::Retryable]); + let dispatcher = server.create_dispatcher(2).await; + let response = dispatcher + .dispatch(vec![build_insert_request(®ions)]) + .await + .unwrap(); + assert_eq!(response.affected_rows, regions.len()); + let mut expected = HashMap::new(); + for region in ®ions { + *expected + .entry(RegionId::new(1, *region).as_u64()) + .or_insert(0) += 1; + } + *expected.get_mut(&RegionId::new(1, 1).as_u64()).unwrap() += 1; + assert_eq!(server.attempt_counts(), expected); + assert_eq!( + *server.service.request_sizes.lock().unwrap(), + vec![regions.len(), 1] + ); + assert!(server.invalidator.calls.lock().unwrap().is_empty()); + } } #[tokio::test] @@ -419,10 +519,10 @@ async fn test_write_retry_transport_errors() { } server .service - .rejections + .rpc_rejections .lock() .unwrap() - .insert(region_id, VecDeque::from([status])); + .push_back(status); let dispatcher = server.create_dispatcher(2).await; let attempt = dispatcher .send( @@ -452,12 +552,12 @@ async fn test_write_retry_transport_attempt_budget() { for failures in [1, 3] { let server = RegionTestServer::start().await; let region_id = RegionId::new(1, 1).as_u64(); - server.service.rejections.lock().unwrap().insert( - region_id, - (0..failures) - .map(|_| Status::new(code, "write failed")) - .collect(), - ); + server + .service + .rpc_rejections + .lock() + .unwrap() + .extend((0..failures).map(|_| Status::new(code, "write failed"))); let dispatcher = server.create_dispatcher(2).await; let response = dispatcher.dispatch(vec![build_insert_request(&[1])]).await; let expected_attempts = if failures == 1 { @@ -490,8 +590,11 @@ async fn test_write_retry_stops_when_a_sibling_is_nonretryable() { .await .unwrap_err(); let terminal_region = RegionId::new(1, 2).as_u64(); + assert_eq!(error.retry_hint(), RetryHint::NonRetryable); + assert_eq!(error.status_code(), StatusCode::InvalidArguments); assert!( - format!("{error:?}") + error + .output_msg() .contains(&format!("region {terminal_region} rejected the write")), "{error:?}" ); @@ -726,81 +829,83 @@ async fn test_finish_retries_initial_failures_in_one_round() { #[tokio::test] async fn test_write_retry_refreshes_stale_route_and_sends_only_failed_region_to_new_leader() { - let old_server = RegionTestServer::start().await; - let new_server = RegionTestServer::start().await; - old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); - let backend = prepare_mocked_backend().await; - // Reuse the operator fixture to populate table metadata and partition rules. - create_partition_rule_manager(backend.clone()).await; - update_table_leader(&backend, &old_server.peer).await; + for regions in [vec![2], vec![1, 2]] { + let old_server = RegionTestServer::start().await; + let new_server = RegionTestServer::start().await; + old_server.fail_region(2, [RetryHint::RetryAfterRouteRefresh]); + let backend = prepare_mocked_backend().await; + // Reuse the operator fixture to populate table metadata and partition rules. + create_partition_rule_manager(backend.clone()).await; + update_table_leader(&backend, &old_server.peer).await; - let (partition_manager, table_info_cache, registry, catalog_manager) = - create_routing_caches(backend.clone()); - let failed_region = RegionId::new(1, 2); - assert_eq!( - partition_manager - .find_region_leader(failed_region) + let (partition_manager, table_info_cache, registry, catalog_manager) = + create_routing_caches(backend.clone()); + let failed_region = RegionId::new(1, 2); + assert_eq!( + partition_manager + .find_region_leader(failed_region) + .await + .unwrap(), + old_server.peer + ); + let table_info = table_info_cache.get(1).await.unwrap().unwrap(); + let (_, versions) = partition_manager + .find_table_partition_rule(&table_info) .await - .unwrap(), - old_server.peer - ); - let table_info = table_info_cache.get(1).await.unwrap().unwrap(); - let (_, versions) = partition_manager - .find_table_partition_rule(&table_info) - .await - .unwrap(); + .unwrap(); - // Metadata has migrated while the frontend still holds the old route. - update_table_leader(&backend, &new_server.peer).await; - assert_eq!( - partition_manager - .find_region_leader(failed_region) - .await - .unwrap(), - old_server.peer - ); + // Metadata has migrated while the frontend still holds the old route. + update_table_leader(&backend, &new_server.peer).await; + assert_eq!( + partition_manager + .find_region_leader(failed_region) + .await + .unwrap(), + old_server.peer + ); - let invalidator = Arc::new(RecordingCacheInvalidator { - calls: Mutex::default(), - inner: Some(registry), - ..Default::default() - }); - let clients = Arc::new(NodeClients::new(ChannelConfig::default())); - let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry( - build_retry_options(2), - catalog_manager, - invalidator.clone(), - ); - let mut request = build_insert_request(&[1, 2]); - if let Some(Body::Inserts(inserts)) = &mut request.body { - for insert in &mut inserts.requests { - insert.partition_expr_version = versions - [&RegionId::from_u64(insert.region_id).region_number()] - .map(|value| PartitionExprVersion { value }); + let invalidator = Arc::new(RecordingCacheInvalidator { + calls: Mutex::default(), + inner: Some(registry), + ..Default::default() + }); + let clients = Arc::new(NodeClients::new(ChannelConfig::default())); + let dispatcher = WriteDispatcher::new(clients, partition_manager).with_retry( + build_retry_options(2), + catalog_manager, + invalidator.clone(), + ); + let mut request = build_insert_request(®ions); + if let Some(Body::Inserts(inserts)) = &mut request.body { + for insert in &mut inserts.requests { + insert.partition_expr_version = versions + [&RegionId::from_u64(insert.region_id).region_number()] + .map(|value| PartitionExprVersion { value }); + } } + let response = dispatcher.dispatch(vec![request]).await.unwrap(); + assert_eq!(response.affected_rows, regions.len()); + assert_eq!( + old_server.attempt_counts(), + regions + .iter() + .map(|region| (RegionId::new(1, *region).as_u64(), 1)) + .collect() + ); + assert_eq!( + new_server.attempt_counts(), + HashMap::from([(failed_region.as_u64(), 1)]) + ); + assert_eq!( + *invalidator.calls.lock().unwrap(), + vec![ + vec![CacheIdent::TableId(1)], + vec![CacheIdent::TableName(TableName::new( + "greptime", "public", "table_1" + ))], + ] + ); } - let response = dispatcher.dispatch(vec![request]).await.unwrap(); - assert_eq!(response.affected_rows, 2); - assert_eq!( - old_server.attempt_counts(), - HashMap::from([ - (RegionId::new(1, 1).as_u64(), 1), - (failed_region.as_u64(), 1), - ]) - ); - assert_eq!( - new_server.attempt_counts(), - HashMap::from([(failed_region.as_u64(), 1)]) - ); - assert_eq!( - *invalidator.calls.lock().unwrap(), - vec![ - vec![CacheIdent::TableId(1)], - vec![CacheIdent::TableName(TableName::new( - "greptime", "public", "table_1" - ))], - ] - ); } #[derive(Clone, Copy, Debug)] @@ -1706,11 +1811,11 @@ async fn check_write_retry_rebuilds_logical_rows_in_mixed_table_envelope(kind: W ); let mut requests = server.service.requests.lock().unwrap().clone(); requests.sort_by_key(|request| write_region_id(request).unwrap().as_u64()); - assert_eq!( - requests, - vec![physical, rebuilt_upper, logical, rebuilt_lower], - "{kind:?}" - ); + let mut expected = vec![physical, rebuilt_upper, logical, rebuilt_lower]; + for request in &mut expected { + api::region::set_accept_partial_result(request.header.get_or_insert_default()); + } + assert_eq!(requests, expected, "{kind:?}"); let calls = invalidator.calls.lock().unwrap(); assert_eq!(calls.len(), 2, "{kind:?}"); assert_eq!(calls[0].len(), 2, "{kind:?}"); diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs index 08e3b4f6d31..efff7eab759 100644 --- a/src/servers/src/error.rs +++ b/src/servers/src/error.rs @@ -994,6 +994,7 @@ pub fn status_code_to_http_status(status_code: &StatusCode) -> HttpStatusCode { StatusCode::Internal | StatusCode::Unexpected | StatusCode::IllegalState + | StatusCode::PartialFailure | StatusCode::Unknown | StatusCode::RuntimeResourcesExhausted | StatusCode::EngineExecuteQuery => HttpStatusCode::INTERNAL_SERVER_ERROR, diff --git a/src/servers/src/grpc/region_server.rs b/src/servers/src/grpc/region_server.rs index cb1cbee7891..842e563f873 100644 --- a/src/servers/src/grpc/region_server.rs +++ b/src/servers/src/grpc/region_server.rs @@ -32,6 +32,13 @@ use crate::grpc::{TonicResult, cancellation}; #[async_trait] pub trait RegionServerHandler: Send + Sync { async fn handle(&self, request: region_request::Body) -> Result; + + async fn handle_region_request(&self, request: RegionRequest) -> Result { + let body = request.body.context(InvalidQuerySnafu { + reason: "Expecting non-empty region request body.", + })?; + self.handle(body).await + } } pub type RegionServerHandlerRef = Arc; @@ -51,15 +58,12 @@ impl RegionServerRequestHandler { let tracing_context = TracingContext::from_w3c( &request .header + .as_ref() .context(InvalidQuerySnafu { reason: "Expecting non-empty region request header.", })? .tracing_context, ); - let query = request.body.context(InvalidQuerySnafu { - reason: "Expecting non-empty region request body.", - })?; - let handler = self.handler.clone(); // Executes requests in another runtime to @@ -71,7 +75,7 @@ impl RegionServerRequestHandler { // 2. avoid the handler blocks the gRPC runtime incidentally. let handle = self.runtime.spawn(async move { handler - .handle(query) + .handle_region_request(request) .trace(tracing_context.attach(info_span!("RegionServerRequestHandler::handle"))) .await .map_err(|e| { @@ -113,3 +117,41 @@ impl RegionServer for RegionServerRequestHandler { cancellation::with_cancellation_handler(request_future, cancellation_future).await } } + +#[cfg(test)] +mod tests { + use super::*; + + struct HeaderCheckingHandler { + expected: RegionRequest, + } + + #[async_trait] + impl RegionServerHandler for HeaderCheckingHandler { + async fn handle(&self, _request: region_request::Body) -> Result { + panic!("request header must reach the region handler"); + } + + async fn handle_region_request(&self, request: RegionRequest) -> Result { + assert_eq!(self.expected, request); + Ok(RegionResponse::default()) + } + } + + #[tokio::test] + async fn test_region_request_preserves_partial_result_header() { + let mut header = api::v1::region::RegionRequestHeader::default(); + api::region::set_accept_partial_result(&mut header); + let request = RegionRequest { + header: Some(header), + body: Some(region_request::Body::Inserts(Default::default())), + }; + let handler = RegionServerRequestHandler::new( + Arc::new(HeaderCheckingHandler { + expected: request.clone(), + }), + common_runtime::global_runtime().clone(), + ); + handler.handle(request).await.unwrap(); + } +} diff --git a/src/servers/src/mysql/writer.rs b/src/servers/src/mysql/writer.rs index d9253c5bea7..049f925c650 100644 --- a/src/servers/src/mysql/writer.rs +++ b/src/servers/src/mysql/writer.rs @@ -512,7 +512,9 @@ pub fn create_mysql_column_def(schema: &SchemaRef) -> Result> { fn mysql_error_kind(status_code: &StatusCode) -> ErrorKind { match status_code { StatusCode::Success => ErrorKind::ER_YES, - StatusCode::Unknown | StatusCode::External => ErrorKind::ER_UNKNOWN_ERROR, + StatusCode::PartialFailure | StatusCode::Unknown | StatusCode::External => { + ErrorKind::ER_UNKNOWN_ERROR + } StatusCode::Unsupported => ErrorKind::ER_NOT_SUPPORTED_YET, StatusCode::Cancelled | StatusCode::DeadlineExceeded => ErrorKind::ER_QUERY_INTERRUPTED, StatusCode::RuntimeResourcesExhausted => ErrorKind::ER_OUT_OF_RESOURCES, diff --git a/src/servers/src/postgres/types/error.rs b/src/servers/src/postgres/types/error.rs index a76bb4362c4..00d1bbbb1c1 100644 --- a/src/servers/src/postgres/types/error.rs +++ b/src/servers/src/postgres/types/error.rs @@ -382,7 +382,8 @@ impl From for PgErrorCode { StatusCode::DeadlineExceeded => PgErrorCode::Ec57000, StatusCode::External => PgErrorCode::Ec58000, - StatusCode::Unknown + StatusCode::PartialFailure + | StatusCode::Unknown | StatusCode::Unexpected | StatusCode::IllegalState | StatusCode::Internal => PgErrorCode::EcXX000, diff --git a/src/standalone/Cargo.toml b/src/standalone/Cargo.toml index 3dcb1d5f108..0eae7ea53fa 100644 --- a/src/standalone/Cargo.toml +++ b/src/standalone/Cargo.toml @@ -46,4 +46,6 @@ tokio.workspace = true url.workspace = true [dev-dependencies] +datanode = { workspace = true, features = ["testing"] } +mito2 = { workspace = true, features = ["test"] } toml.workspace = true diff --git a/src/standalone/src/datanode_manager.rs b/src/standalone/src/datanode_manager.rs index a27a94b3c8d..73801e294b3 100644 --- a/src/standalone/src/datanode_manager.rs +++ b/src/standalone/src/datanode_manager.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use api::region::RegionResponse; use api::v1::region::{RegionRequest, RegionResponse as RegionResponseV1}; use async_trait::async_trait; -use client::region::check_response_header; +use client::region; use common_error::ext::BoxedError; use common_meta::error::{self as meta_error, Result as MetaResult}; use common_meta::node_manager::{ @@ -100,10 +100,7 @@ impl Datanode for RegionInvoker { .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu)?; - check_response_header(&response.header) - .map_err(BoxedError::new) - .context(meta_error::ExternalSnafu)?; - Ok(RegionResponse::from_region_response(response)) + region::decode_region_response(response).map_err(Into::into) } async fn handle_query(&self, request: QueryRequest) -> MetaResult { @@ -125,3 +122,65 @@ impl Datanode for RegionInvoker { .context(meta_error::ExternalSnafu) } } + +#[cfg(test)] +mod tests { + use api::v1::Rows; + use api::v1::region::{InsertRequest, InsertRequests, region_request}; + use common_error::ext::ErrorExt; + use common_error::mock::MockError; + use common_error::status_code::StatusCode; + use datanode::tests::{MockRegionEngine, mock_region_server}; + use mito2::engine::MITO_ENGINE_NAME; + use mito2::test_util::CreateRequestBuilder; + use store_api::region_request::RegionRequest as StoreRegionRequest; + use store_api::storage::RegionId; + + use super::*; + + #[tokio::test] + async fn test_write_preserves_original_error() { + for code in [StatusCode::StorageUnavailable, StatusCode::InvalidArguments] { + let mut server = mock_region_server(); + let (engine, _receiver) = MockRegionEngine::with_mock_fn( + MITO_ENGINE_NAME, + Box::new(move |region_id, request| { + if matches!(request, StoreRegionRequest::Put(_)) { + Err(BoxedError::new(MockError::new(code))) + .context(datanode::error::HandleRegionRequestSnafu { region_id }) + } else { + Ok(0) + } + }), + ); + server.register_engine(engine); + let mut requests = Vec::new(); + for number in 0..2 { + let region_id = RegionId::new(1, number); + server + .handle_request( + region_id, + StoreRegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + requests.push(InsertRequest { + region_id: region_id.as_u64(), + rows: Some(Rows::default()), + ..Default::default() + }); + } + let datanode: DatanodeRef = RegionInvoker::arc(server); + let error = datanode + .handle(RegionRequest { + header: Some(Default::default()), + body: Some(region_request::Body::Inserts(InsertRequests { requests })), + }) + .await + .unwrap_err(); + assert_eq!(code, error.status_code()); + assert_eq!(MockError::new(code).retry_hint(), error.retry_hint()); + assert!(!matches!(error, meta_error::Error::PartialRegion { .. })); + } + } +} diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index 6a40f3dbc6b..3a4bc848135 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -209,47 +209,55 @@ impl RegionRequest { } } +/// Converts one insert, treating absent rows as a no-op like vector inserts. +pub fn make_region_put( + request: api::v1::region::InsertRequest, +) -> Option<(RegionId, RegionPutRequest)> { + request.rows.map(|rows| { + ( + request.region_id.into(), + RegionPutRequest { + rows, + hint: None, + skip_wal: request.skip_wal, + partition_expr_version: request.partition_expr_version.map(|version| version.value), + }, + ) + }) +} + fn make_region_puts(inserts: InsertRequests) -> Result> { - let requests = inserts + Ok(inserts .requests .into_iter() - .filter_map(|r| { - let region_id = r.region_id.into(); - r.rows.map(|rows| { - ( - region_id, - RegionRequest::Put(RegionPutRequest { - rows, - hint: None, - skip_wal: r.skip_wal, - partition_expr_version: r.partition_expr_version.map(|v| v.value), - }), - ) - }) - }) - .collect(); - Ok(requests) + .filter_map(make_region_put) + .map(|(region_id, request)| (region_id, RegionRequest::Put(request))) + .collect()) +} + +/// Converts one delete, treating absent rows as a no-op like vector deletes. +pub fn make_region_delete( + request: api::v1::region::DeleteRequest, +) -> Option<(RegionId, RegionDeleteRequest)> { + request.rows.map(|rows| { + ( + request.region_id.into(), + RegionDeleteRequest { + rows, + hint: None, + partition_expr_version: request.partition_expr_version.map(|version| version.value), + }, + ) + }) } fn make_region_deletes(deletes: DeleteRequests) -> Result> { - let requests = deletes + Ok(deletes .requests .into_iter() - .filter_map(|r| { - let region_id = r.region_id.into(); - r.rows.map(|rows| { - ( - region_id, - RegionRequest::Delete(RegionDeleteRequest { - rows, - hint: None, - partition_expr_version: r.partition_expr_version.map(|v| v.value), - }), - ) - }) - }) - .collect(); - Ok(requests) + .filter_map(make_region_delete) + .map(|(region_id, request)| (region_id, RegionRequest::Delete(request))) + .collect()) } fn parse_region_create(create: CreateRequest) -> Result<(RegionId, RegionCreateRequest)> { diff --git a/tests/runner/src/client.rs b/tests/runner/src/client.rs index a595c877940..03f7f7ed1f5 100644 --- a/tests/runner/src/client.rs +++ b/tests/runner/src/client.rs @@ -169,6 +169,7 @@ impl MultiProtocolClient { result = ServerSnafu { code: status_code, msg, + retry_hint: e.retry_hint(), } .fail(); }