diff --git a/.agents/skills/greptimedb-development-docker-image/SKILL.md b/.agents/skills/greptimedb-development-docker-image/SKILL.md new file mode 100644 index 00000000000..02a9a61d871 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/SKILL.md @@ -0,0 +1,368 @@ +--- +name: greptimedb-development-docker-image +description: Builds a development-only GreptimeDB Docker image from a local debug binary for local-cluster testing, and optionally pushes it to a development registry. Use when the user asks to package, build, tag, publish, or cross-build a non-release GreptimeDB or GreptimeDB Enterprise image for debugging. +compatibility: Requires Docker with Buildx for cross-platform builds; Podman is supported for native-platform builds. +metadata: + protocols: docker buildx podman + platforms: linux-amd64 linux-arm64 +--- + +# GreptimeDB Development Docker Image + +## Goal + +Package a locally built GreptimeDB binary into a **development-only** Docker +image for debugging and local-cluster testing, optionally push it to a +development registry, and retain the non-secret image settings in `.env` for +the next build. It is **not** a release-image workflow and must not be used to +publish a production or release artifact. + +This skill follows the documented development-image build procedure: + +- Build `greptime` with Cargo's `nightly` profile. +- Copy the target binary into the Docker build context as `./greptime`. +- Build the supplied `Dockerfile`, which uses Ubuntu 24.04 only as the runtime + base image and exposes the `greptime` binary as its entrypoint. + +## Inputs + +Collect or discover the following before any build or push: + +| Input | Discovery and rule | +| --- | --- | +| Source repository | Treat the directory in which this skill is invoked as the default workspace. Include it as an editable field in the batch configuration; do not ask a separate confirmation. Do not assume open-source versus Enterprise. | +| Edition and binary | Default to `greptime`; use the binary name the user requests for Enterprise builds. The copied Docker-context filename must always be `greptime`. | +| Cargo profile and binary source | Include this in the batch configuration. Default to rebuilding `nightly`; reuse an existing binary only after the user explicitly accepts that its freshness is unverified. | +| Target platform | Preselect `linux/amd64` unless Docker's server is `linux/arm64`, then preselect `linux/arm64`. Let the user override it in the batch configuration. Ubuntu 24.04 is the runtime base image, not a target-platform choice. Each image has exactly one target platform, but it may differ from the host platform. | +| Build mode | Ask whether the user wants a locally loadable debug image or a registry push. | +| Registry/repository | Read `IMAGE_REGISTRY` and `IMAGE_REPOSITORY` only from the selected workspace's `.env` and prefill them in the batch configuration. Example: `registry.example.com/team` + `greptimedb-dev`. | +| Tag | Read `IMAGE_TAG` only from the selected workspace's `.env` and prefill it in the batch configuration. When the tag exists in the selected registry, preselect an incremented version. | + +Use the image reference `${IMAGE_REGISTRY}/${IMAGE_REPOSITORY}:${IMAGE_TAG}`. +If the registry is intentionally empty, omit its slash rather than producing a +leading slash. + +## Interactive Workflow + +Use the platform's interactive prompt component for every question, selection, +and confirmation. Do not ask an open-ended text question when a single-choice, +multi-select, or confirmation component can represent the decision. If the +platform does not provide an interactive component, use the equivalent numbered +or lettered prompt below and wait for input before continuing. + +### Batch configuration + +Minimize user round trips: run the collector first against the invocation +directory, then use one interactive form or batched prompt to collect workspace, +edition/binary, profile/reuse-or-rebuild choice, target platform, build mode, +registry/repository, tag, and whether to inspect the configured registry tag. +Use the invocation directory as the editable workspace default. Prefill +collector values and mark recommended defaults. If the user changes workspace, +rerun the collector for that workspace without asking another configuration +question. Only ask a follow-up if information is missing or invalid, or a safety +gate is required. Treat unchanged prefilled values as accepted. + +Keep separate confirmation components only for elevated privileges: `sudo`, +package installation, and privileged QEMU setup. Use one final confirmation for +all non-privileged selected work. + +Always offer **Cancel** for an action that can write state, build, push, or +require elevated privileges. Display the relevant preview before the user +confirms it: source/binary path, `file` architecture result, platform(s), image +reference, build mode, and non-secret `.env` values. + +### Batch configuration fields + +Use single-choice controls inside the one batch form for mutually exclusive +choices: + +| Decision | Required choices | +| --- | --- | +| Cargo output | Rebuild `nightly` (recommended), reuse an existing binary only with a freshness-unverified acknowledgement, choose a profile/binary | +| Target platform | `linux/amd64` (default unless Docker server is arm64), `linux/arm64` (default when Docker server is arm64). The selected image may be cross-built with Buildx. | +| Build mode | Load locally (recommended), push to development registry | +| Registry/repository | Use `.env` defaults, enter new values | +| Tag | Use proposed increment when the configured tag exists, enter a version | + +When push is selected and the image configuration is complete, automatically +run the read-only registry-tag check. It may return `unknown`; do not treat that +as a missing tag. + +Use one final confirmation for all selected non-privileged operations: + +- persist changed non-secret image settings to `.env`; +- run the selected Cargo build or reuse the resolved binary; +- prepare the Docker context; +- build, verify, and, when selected, push the image. + +The final confirmation must clearly state: **“This creates a development and +local-cluster test image, not a release or production artifact.”** + +For a real multi-select only, use lettered choices (`[A]`, `[B]`) and accept +`all`, `none`, or `cancel`; otherwise use single-choice components. + +### Preflight context collection + +Before asking the profile, registry, or tag questions, run the bundled, +read-only collector. It works on macOS and Linux, does not invoke `cargo build`, +and emits JSON that can be shown or summarized to the user: + +```bash +python3 /scripts/collect_context.py \ + --source \ + --bin greptime \ + --profile nightly \ + [--target ] +``` + +Use its report to identify: + +- missing `IMAGE_REGISTRY`, `IMAGE_REPOSITORY`, or `IMAGE_TAG` values in the + selected workspace's `.env`; +- Cargo binary targets available in the workspace and whether the requested + `--bin` exists; +- the resolved Cargo target directory, expected profile output path, and whether + that binary already exists; and +- macOS/Linux host architecture plus Docker's server platform when Docker is + available. + +When push is selected and image configuration is complete, rerun the collector +with `--check-registry-tag` and all selected image values. It uses the selected +engine's read-only manifest inspection and returns `exists` or `unknown`; +`unknown` includes an unavailable registry, an absent tag, or missing +authentication and must not be treated as an absent tag. + +```bash +python3 /scripts/collect_context.py \ + --source \ + --bin \ + --profile \ + --engine \ + --registry \ + --repository \ + --tag \ + --check-registry-tag +``` + +When the status is `exists` and the tag ends in a number, preselect the +collector's `candidate_tag` as the next development tag. The user must still +confirm it before `.env` is updated or an image is built. Never use automatic +incrementing as permission to overwrite or push an existing image. + +If Cargo metadata or the requested binary target is unavailable, stop before +building and ask the user to correct the source path or binary target. + +The profile/binary choice is a field in the batch configuration, not a separate +question. Its default is the existing `nightly` output when present: + +```text +Question: Which Cargo build output should package this image? +Options: +- Rebuild nightly (Recommended): run Cargo with `--profile nightly`, then use its output. +- Reuse nightly output: use the existing output only after showing its path, modification time, size, and matching platform; freshness is unverified. +- Choose profile or binary: provide a different Cargo profile or an explicit binary path. +``` + +Do not infer that a binary under `target/debug`, `target/release`, or another +profile is suitable. For an existing binary, show its resolved path and run +`file` before asking for final build confirmation. + +## Environment State + +The image configuration lives only in the selected workspace's `.env`. The +collector reads only the following keys and never reads or displays other `.env` +values: + +```dotenv +IMAGE_REGISTRY=registry.example.com/team +IMAGE_REPOSITORY=greptimedb-dev +IMAGE_TAG=dev-001 +``` + +1. Read `.env` if it exists. Ignore blank lines and `#` comments. +2. Show the discovered registry, repository, and proposed tag to the user. +3. Include registry/repository and tag in the batch configuration even when + values already exist; saved values are defaults, not authorization to reuse + them. +4. Compare the selected configuration with the three current managed values. If + all match, skip `.env` confirmation and do not update the file. Otherwise, + preview only the changed `IMAGE_*` values, request confirmation, and update + the file while preserving unrelated entries and comments. +5. Never put registry credentials, access tokens, passwords, or `docker login` + output in `.env`, build commands, logs, or responses. Ask the user to log in + themselves if a push needs authentication. + +For a repeated build, offer an incremented tag before asking. Use +`scripts/next_image_tag.py --tag ` to calculate it. It increments +the final numeric component while preserving zero padding: + +```text +weny-2025-0715-01 -> weny-2025-0715-02 +v0.1.4 -> v0.1.5 +debug-009 -> debug-010 +``` + +If the saved tag has no trailing number, ask the user for the next tag rather +than inventing a versioning convention. Do not overwrite an existing image tag +without explicit user confirmation. + +Only if a managed value is missing or differs, persist the values with the +bundled helper after confirmation instead of hand-editing the file: + +```bash +python3 /scripts/update_image_env.py \ + --env /.env \ + --registry \ + --repository \ + --tag +``` + +## Build Procedure + +### 1. Inspect the host and container tooling + +Run these checks and report the selected path: + +```bash +uname -s +uname -m +docker version --format '{{.Server.Os}}/{{.Server.Arch}}' +docker buildx version +docker buildx inspect --bootstrap +``` + +For a native-platform build, Podman can be used if Docker is unavailable. For +any non-native request, require Docker Buildx. Do not silently fall back to a +native build when the requested image platform differs. + +If the requested target is unavailable, stop and ask the user to configure a +Buildx builder and cross-compilation toolchain externally. Do not automate +privileged QEMU or binfmt setup, and never run an unpinned privileged image. + +### 2. Build the executable for the image platform + +From the source checkout, build the requested binary. For open-source amd64: + +```bash +/scripts/build_binary.sh \ + --source --package cmd --bin greptime --profile nightly +``` + +For open-source arm64: + +```bash +/scripts/build_binary.sh \ + --source \ + --package cmd \ + --bin greptime \ + --profile nightly \ + --target aarch64-unknown-linux-gnu +``` + +For Enterprise, use the user-provided executable target, for example: + +```bash +/scripts/build_binary.sh \ + --source --bin greptime-ent-cloud --profile nightly +``` + +For a user-selected profile, replace `nightly` in the helper call and resolve +the binary under `target//` (or +`target///` for cross-compilation). Reuse the +existing target file only when the interactive profile question selected reuse; +otherwise invoke the helper to rebuild it. + +Do not use a binary compiled for the host architecture in an image intended for +another architecture. Determine the binary path from the Cargo target and +profile, then copy it into the Docker context as `greptime`: + +```bash +cp /greptime +``` + +Verify it before building the image with `file /greptime`; its +reported architecture must match the requested target platform. + +Create a fresh isolated build context (never use the repository root) and use +the bundled preparation script. The helper rejects an existing Dockerfile or +binary output to prevent accidental context reuse: + +```bash +/scripts/prepare_context.sh \ + --context "$(mktemp -d)" \ + --binary \ + --dockerfile /assets/Dockerfile \ + --platform +``` + +### 3. Build or push the image + +Run the command only after the user confirms the final image reference and +whether it should be pushed. + +Use the bundled build helper rather than spelling out individual Docker or +Podman commands. It follows the existing scripts' Docker-first/Podman-fallback +behavior, uses Buildx for a non-native Docker request, and accepts exactly one +target platform per image: + +```bash +/scripts/build_image.sh \ + --context \ + --image \ + --platform \ + --mode local|push +``` + +`--mode local` builds a native image locally or uses `docker buildx --load` for +a non-native image. `--mode push` pushes the selected single-platform image. +The helper never pushes in `local` mode. + +This is intentionally a runtime image built from a precompiled binary, not a +multi-stage Dockerfile. Ubuntu 24.04 is the runtime base image only; select one +target platform separately as `linux/amd64` or `linux/arm64`. Only introduce a +multi-stage Dockerfile if the user asks to compile within Docker or the local +Rust toolchain cannot build the required target. Keep the final Ubuntu 24.04 +runtime stage and existing entrypoint unless the user asks to change runtime +behavior. + +At the final confirmation, repeat that the selected reference is a +**development/local-cluster test image**, not a release artifact. If the user +requests a release or production image, stop and direct them to the release +process instead. + +### 4. Verify the result + +For a local image, use the selected engine to inspect its architecture and +start it with `--help`: + +```bash +docker image inspect --format '{{.Os}}/{{.Architecture}}' +docker run --rm --help +``` + +For Podman, use `podman image inspect ` and +`podman run --rm --help`. For a pushed Podman image, use +`podman manifest inspect `. + +For a pushed image, inspect its remote manifest: + +```bash +docker buildx imagetools inspect +``` + +Report the final image reference, target platform, source binary path, and +whether the image was loaded locally or pushed. + +## Safety Rules + +- This skill is only for development and local-cluster testing. Never present + its image as a production, release, or officially published artifact. +- Never push by default. Local debug builds should use a local tag and load the + image unless the user explicitly requests a push. +- Do not run `docker login`, privileged QEMU setup, package installation, or + overwrite a tag without confirmation. +- Do not delete existing images, builders, binaries, or `.env` entries as part + of this workflow. +- Stop and explain if Docker's server is unavailable, Buildx cannot support the + requested platform, or the copied binary architecture does not match the + image platform. diff --git a/.agents/skills/greptimedb-development-docker-image/assets/Dockerfile b/.agents/skills/greptimedb-development-docker-image/assets/Dockerfile new file mode 100644 index 00000000000..4e1a7d2c502 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/assets/Dockerfile @@ -0,0 +1,15 @@ +FROM ubuntu:24.04 + +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + ca-certificates \ + curl \ + unzip \ + mysql-client \ + wget \ + && rm -rf /var/lib/apt/lists/* + +COPY ./greptime /greptime/bin/ + +ENV PATH=/greptime/bin/:$PATH + +ENTRYPOINT ["greptime"] diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/binary_platform.py b/.agents/skills/greptimedb-development-docker-image/scripts/binary_platform.py new file mode 100644 index 00000000000..86f855a11cf --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/binary_platform.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Validate that a binary is an executable ELF for a requested Linux platform.""" + +import argparse +import stat +import struct +import sys +from pathlib import Path +from typing import Optional + +PLATFORM_MACHINES = {"linux/amd64": 0x3E, "linux/arm64": 0xB7} + + +def validate_platform(platform: str) -> int: + try: + return PLATFORM_MACHINES[platform] + except KeyError as error: + raise ValueError("platform must be linux/amd64 or linux/arm64") from error + + +def inspect_elf(path: Path) -> int: + executable_bits = stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH + if path.is_symlink() or not path.is_file() or not (path.stat().st_mode & executable_bits): + raise ValueError("binary must be an executable regular file") + try: + with path.open("rb") as handle: + header = handle.read(64) + except OSError as error: + raise ValueError(f"cannot read binary: {error}") from error + if len(header) != 64 or header[:4] != b"\x7fELF": + raise ValueError("binary is not ELF") + if header[4] != 2 or header[5] != 1 or header[6] != 1: + raise ValueError("binary must be a 64-bit little-endian ELF") + if header[7] not in (0, 3): + raise ValueError("binary must use the System V or Linux ELF ABI") + elf_type, machine, version = struct.unpack_from(" None: + expected = validate_platform(platform) + actual = inspect_elf(path) + if actual != expected: + raise ValueError("binary architecture does not match requested platform") + + +def main(argv: Optional[list[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--binary", required=True, type=Path) + parser.add_argument("--platform", required=True) + args = parser.parse_args(argv) + try: + validate_binary(args.binary, args.platform) + except ValueError as error: + print(f"error: {error}", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/build_binary.sh b/.agents/skills/greptimedb-development-docker-image/scripts/build_binary.sh new file mode 100755 index 00000000000..605685fa7c7 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/build_binary.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash + +set -euo pipefail + +usage() { + printf 'Usage: %s --source PATH --bin NAME --profile PROFILE [--target RUST_TARGET] [--package PACKAGE] [--features FEATURES]\n' "$0" >&2 +} + +source_dir="" +binary_name="" +profile="" +target="" +package="" +features="" + +while [[ $# -gt 0 ]]; do + case "$1" in + --source|--bin|--profile|--target|--package|--features) + if [[ $# -lt 2 ]]; then + usage + exit 2 + fi + case "$1" in + --source) source_dir="$2" ;; + --bin) binary_name="$2" ;; + --profile) profile="$2" ;; + --target) target="$2" ;; + --package) package="$2" ;; + --features) features="$2" ;; + esac + shift 2 + ;; + *) + usage + exit 2 + ;; + esac +done + +if [[ -z "$source_dir" || -z "$binary_name" || -z "$profile" ]]; then + usage + exit 2 +fi + +if [[ ! -f "$source_dir/Cargo.toml" ]]; then + printf 'error: not a Cargo workspace: %s\n' "$source_dir" >&2 + exit 1 +fi + +command=(cargo build --locked --bin "$binary_name" --profile "$profile") +if [[ -n "$package" ]]; then + command+=(--package "$package") +fi +if [[ -n "$features" ]]; then + command+=(--features "$features") +fi +if [[ -n "$target" ]]; then + command+=(--target "$target") +fi + +printf 'building %s with profile %s\n' "$binary_name" "$profile" +( + cd "$source_dir" + exec "${command[@]}" +) diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/build_image.sh b/.agents/skills/greptimedb-development-docker-image/scripts/build_image.sh new file mode 100755 index 00000000000..3b480334671 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/build_image.sh @@ -0,0 +1,136 @@ +#!/usr/bin/env bash + +set -euo pipefail + +usage() { + printf 'Usage: %s --context PATH --image REFERENCE --platform PLATFORM --mode local|push [--dockerfile PATH] [--engine auto|docker|podman] [--allow-local-tag-overwrite]\n' "$0" >&2 +} + +context="" +image="" +platform="" +mode="" +dockerfile="" +engine="auto" +allow_local_tag_overwrite=false + +while [[ $# -gt 0 ]]; do + case "$1" in + --context|--image|--platform|--mode|--dockerfile|--engine) + if [[ $# -lt 2 ]]; then + usage + exit 2 + fi + case "$1" in + --context) context="$2" ;; + --image) image="$2" ;; + --platform) platform="$2" ;; + --mode) mode="$2" ;; + --dockerfile) dockerfile="$2" ;; + --engine) engine="$2" ;; + esac + shift 2 + ;; + --allow-local-tag-overwrite) + allow_local_tag_overwrite=true + shift + ;; + *) + usage + exit 2 + ;; + esac +done + +if [[ -z "$context" || -z "$image" || -z "$platform" || -z "$mode" ]]; then + usage + exit 2 +fi + +if [[ "$mode" != "local" && "$mode" != "push" ]]; then + printf 'error: --mode must be local or push\n' >&2 + exit 2 +fi + +if [[ "$engine" != "auto" && "$engine" != "docker" && "$engine" != "podman" ]]; then + printf 'error: --engine must be auto, docker, or podman\n' >&2 + exit 2 +fi + +dockerfile="${dockerfile:-$context/Dockerfile}" +if [[ ! -f "$dockerfile" ]]; then + printf 'error: Dockerfile does not exist: %s\n' "$dockerfile" >&2 + exit 1 +fi + +if [[ "$platform" == *,* ]]; then + printf 'error: only one target platform is supported per image build\n' >&2 + exit 2 +fi + +case "$platform" in + linux/amd64|linux/arm64) ;; + *) + printf 'error: unsupported platform: %s\n' "$platform" >&2 + exit 2 + ;; +esac + +script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +python3 "$script_dir/binary_platform.py" --binary "$context/greptime" --platform "$platform" + +check_collision() { + local runtime="$1" + if [[ "$allow_local_tag_overwrite" == true ]]; then + return 0 + fi + if [[ "$runtime" == docker ]] && docker image inspect "$image" >/dev/null 2>&1; then + printf 'error: local image tag exists: %s (pass --allow-local-tag-overwrite to replace it)\n' "$image" >&2 + exit 3 + fi + if [[ "$runtime" == podman ]] && podman image exists "$image"; then + printf 'error: local image tag exists: %s (pass --allow-local-tag-overwrite to replace it)\n' "$image" >&2 + exit 3 + fi +} + +if [[ "$engine" != podman ]] && command -v docker >/dev/null 2>&1 && server_platform="$(docker version --format '{{.Server.Os}}/{{.Server.Arch}}' 2>/dev/null)"; then + if [[ "$platform" == "$server_platform" ]]; then + check_collision docker + docker build --platform "$platform" -f "$dockerfile" -t "$image" "$context" + if [[ "$mode" == "push" ]]; then + exec docker push "$image" + fi + exit 0 + fi + + if [[ "$mode" == "push" ]]; then + exec docker buildx build --platform "$platform" --push -f "$dockerfile" -t "$image" "$context" + fi + check_collision docker + exec docker buildx build --platform "$platform" --load -f "$dockerfile" -t "$image" "$context" +fi + +if [[ "$engine" != docker ]] && command -v podman >/dev/null 2>&1; then + native_platform="$(podman info --format '{{.Host.Os}}/{{.Host.Arch}}' 2>/dev/null || true)" + case "$native_platform" in + linux/amd64|linux/arm64) ;; + *) + printf 'error: Podman is unavailable or has unsupported platform: %s\n' "$native_platform" >&2 + exit 1 + ;; + esac + if [[ "$platform" != "$native_platform" ]]; then + printf 'error: Docker Buildx is required for non-native platform %s\n' "$platform" >&2 + exit 1 + fi + check_collision podman + podman build -f "$dockerfile" -t "$image" "$context" + if [[ "$mode" == "push" ]]; then + exec podman push "$image" + fi + exit 0 +fi + +printf 'error: Docker or Podman is required\n' >&2 +exit 1 diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/collect_context.py b/.agents/skills/greptimedb-development-docker-image/scripts/collect_context.py new file mode 100755 index 00000000000..214d56ee789 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/collect_context.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Collect read-only build context for a GreptimeDB development image.""" + +import argparse +import json +import os +import platform +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Optional, cast + +from next_image_tag import increment_tag +from binary_platform import PLATFORM_MACHINES, inspect_elf +from image_config import image_reference, parse_env_file, validate_values + + +ENV_KEYS = ("IMAGE_REGISTRY", "IMAGE_REPOSITORY", "IMAGE_TAG") +PROFILE_DIRECTORIES = {"dev": "debug", "test": "debug", "bench": "release"} +PLATFORM_ARCHITECTURES = { + "x86_64": "linux/amd64", + "amd64": "linux/amd64", + "aarch64": "linux/arm64", + "arm64": "linux/arm64", +} + + +def parse_arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", required=True, type=Path) + parser.add_argument("--bin", dest="binary", default="greptime") + parser.add_argument("--profile", default="nightly") + parser.add_argument("--target", help="Rust target triple for cross-compilation") + parser.add_argument("--package", help="Cargo package containing the selected binary") + parser.add_argument("--platform", choices=("linux/amd64", "linux/arm64")) + parser.add_argument("--engine", choices=("auto", "docker", "podman"), default="auto") + parser.add_argument("--registry") + parser.add_argument("--repository") + parser.add_argument("--tag") + parser.add_argument( + "--check-registry-tag", + action="store_true", + help="inspect the configured image tag with Docker without modifying it", + ) + return parser.parse_args() + + +def native_platform(architecture: str) -> Optional[str]: + return PLATFORM_ARCHITECTURES.get(architecture.lower()) + + +def parse_env(env_path: Path) -> dict[str, object]: + values = parse_env_file(env_path) + missing = [key for key in ("IMAGE_REPOSITORY", "IMAGE_TAG") if not values.get(key)] + try: + if not missing: + validate_values(values) + error = None + except ValueError as exc: + error = str(exc) + return {"path": str(env_path), "exists": env_path.is_file(), "values": values, "missing": missing, "error": error} + + +def cargo_metadata(source: Path) -> dict[str, Any]: + command = ["cargo", "metadata", "--locked", "--format-version=1", "--no-deps"] + try: + result = subprocess.run(command, cwd=source, text=True, capture_output=True, check=True) + except FileNotFoundError as error: + raise RuntimeError("cargo is not installed or is not on PATH") from error + except subprocess.CalledProcessError as error: + detail = error.stderr.strip() or error.stdout.strip() + raise RuntimeError(f"cargo metadata failed: {detail}") from error + return json.loads(result.stdout) + + +def binary_targets(metadata: dict[str, Any]) -> list[dict[str, Any]]: + targets = [] + for package in metadata["packages"]: + for target in package["targets"]: + if "bin" in target["kind"]: + targets.append({ + "package": package["name"], + "name": target["name"], + "path": target["src_path"], + "required_features": target.get("required-features", []), + }) + return sorted(targets, key=lambda target: (target["name"], target["package"])) + + +def expected_binary(target_dir: Path, binary: str, profile: str, rust_target: Optional[str]) -> Path: + profile_directory = PROFILE_DIRECTORIES.get(profile, profile) + output_dir = target_dir / rust_target if rust_target else target_dir + return output_dir / profile_directory / binary + + +def engine_status(engine: str) -> dict[str, object]: + if engine == "docker": + command = ["docker", "version", "--format", "{{.Server.Os}}/{{.Server.Arch}}"] + else: + command = ["podman", "info", "--format", "{{.Host.Os}}/{{.Host.Arch}}"] + try: + result = subprocess.run(command, text=True, capture_output=True, check=True) + except (FileNotFoundError, subprocess.CalledProcessError): + return {"available": False, "server_platform": None} + return {"available": True, "server_platform": result.stdout.strip()} + + +def registry_tag_status(values: dict[str, str], enabled: bool, engine: str, source: str) -> dict[str, object]: + if not enabled: + return {"checked": False, "status": "not_checked", "candidate_tag": None} + + try: + reference = image_reference(values) + except ValueError: + return {"checked": False, "status": "not_configured", "candidate_tag": None, "configuration_source": source} + command = ["docker", "buildx", "imagetools", "inspect", reference] if engine == "docker" else ["podman", "manifest", "inspect", reference] + try: + subprocess.run(command, text=True, capture_output=True, check=True) + except FileNotFoundError: + return {"checked": True, "status": "unknown", "candidate_tag": None, "reference": reference, "engine": engine, "configuration_source": source} + except subprocess.CalledProcessError: + return {"checked": True, "status": "unknown", "candidate_tag": None, "reference": reference, "engine": engine, "configuration_source": source} + + try: + candidate_tag = increment_tag(values["IMAGE_TAG"]) + except ValueError: + candidate_tag = None + return {"checked": True, "status": "exists", "candidate_tag": candidate_tag, "reference": reference, "engine": engine, "configuration_source": source} + + +def collect_report(args: argparse.Namespace) -> dict[str, object]: + source = args.source.expanduser().resolve() + if not (source / "Cargo.toml").is_file(): + raise RuntimeError(f"not a Cargo workspace: {source}") + + metadata = cargo_metadata(source) + targets = binary_targets(metadata) + target_dir = Path(metadata["target_directory"]) + expected = expected_binary(target_dir, args.binary, args.profile, args.target) + host_architecture = platform.machine() + selected_target = next((target for target in targets if target["name"] == args.binary), None) + environment = parse_env(source / ".env") + supplied_values = (args.registry, args.repository, args.tag) + if any(value is not None for value in supplied_values) and not all(value is not None for value in supplied_values): + raise RuntimeError("--registry, --repository, and --tag must be supplied together") + values: dict[str, str] = {"IMAGE_REGISTRY": args.registry or "", "IMAGE_REPOSITORY": args.repository or "", "IMAGE_TAG": args.tag or ""} if all(value is not None for value in supplied_values) else cast(dict[str, str], environment["values"]) + configuration_source = "arguments" if all(value is not None for value in supplied_values) else "environment" + docker = engine_status("docker") + podman = engine_status("podman") + engine = args.engine if args.engine != "auto" else "docker" if docker["available"] else "podman" + binary_details: dict[str, object] = {"platform": None, "error": None} + if expected.is_file(): + try: + machine = inspect_elf(expected) + binary_details["platform"] = next( + (name for name, expected_machine in PLATFORM_MACHINES.items() if expected_machine == machine), + None, + ) + if binary_details["platform"] is None: + binary_details["error"] = f"binary has unsupported ELF machine: {machine}" + except ValueError as exc: + binary_details["error"] = str(exc) + return { + "host": {"os": platform.system(), "arch": host_architecture, "native_platform": native_platform(host_architecture)}, + "docker": docker, + "podman": podman, + "environment": environment, + "registry_tag": registry_tag_status(values, args.check_registry_tag, engine, configuration_source), + "cargo": {"workspace_root": metadata["workspace_root"], "target_dir": str(target_dir), "binaries": targets}, + "selected_binary": {"name": args.binary, "package": args.package, "cargo_target_exists": selected_target is not None, "target": selected_target}, + "expected_binary": {"profile": args.profile, "rust_target": args.target, "path": str(expected), "exists": expected.is_file(), "executable": os.access(expected, os.X_OK), "platform": binary_details["platform"], "requested_platform": args.platform, "platform_matches": binary_details["platform"] == args.platform if args.platform and binary_details["platform"] else None, "validation_error": binary_details["error"], "mtime": datetime.fromtimestamp(expected.stat().st_mtime, timezone.utc).isoformat() if expected.exists() else None, "freshness": "unverified"}, + } + + +def main() -> int: + try: + print(json.dumps(collect_report(parse_arguments()), indent=2, sort_keys=True)) + except RuntimeError as error: + print(f"error: {error}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/image_config.py b/.agents/skills/greptimedb-development-docker-image/scripts/image_config.py new file mode 100644 index 00000000000..cb3260d8c8f --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/image_config.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Bounded parsing and safe updating of the image settings in a .env file.""" + +import os +import re +import tempfile +from pathlib import Path +from typing import Mapping + +MANAGED_KEYS = ("IMAGE_REGISTRY", "IMAGE_REPOSITORY", "IMAGE_TAG") +_REGISTRY = re.compile(r"[a-z0-9][a-z0-9.-]*(?::[0-9]+)?(?:/[a-z0-9][a-z0-9._-]*)*") +_REPOSITORY = re.compile(r"[a-z0-9][a-z0-9._-]*(?:/[a-z0-9][a-z0-9._-]*)*") +_TAG = re.compile(r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,127}") + + +def _value(text: str) -> str: + text = text.strip() + if not text: + return "" + quote = text[0] if text[0] in "'\"" else None + if quote: + end = text.find(quote, 1) + if end < 0 or text[end + 1 :].strip() and not text[end + 1 :].lstrip().startswith("#"): + raise ValueError("unmatched quote or trailing value syntax") + return text[1:end] + # A comment is inline only when introduced by whitespace. This keeps + # valid values such as an image tag containing '#'. + match = re.search(r"\s+#", text) + if match: + text = text[: match.start()].rstrip() + if any(character in text for character in "'\";$`"): + raise ValueError("unsupported shell syntax in .env value") + return text + + +def parse_env(text: str) -> dict[str, str]: + """Parse the three managed keys, without validating their values.""" + lines = text.splitlines() + result: dict[str, str] = {} + for number, original in enumerate(lines, 1): + line = original.strip() + if not line or line.startswith("#"): + continue + export = re.match(r"export[ \t]+", line) + if export: + line = line[export.end() :].lstrip() + key, separator, raw = line.partition("=") + if not separator: + continue + key = key.strip() + if key in MANAGED_KEYS: + try: + result[key] = _value(raw) + except ValueError as error: + raise ValueError(f"line {number}: {error}") from error + return result + + +def parse_env_file(path: Path) -> dict[str, str]: + return parse_env(path.read_text() if path.exists() else "") + + +def normalize_values(values: Mapping[str, str]) -> dict[str, str]: + return { + "IMAGE_REGISTRY": values.get("IMAGE_REGISTRY", ""), + "IMAGE_REPOSITORY": values.get("IMAGE_REPOSITORY", ""), + "IMAGE_TAG": values.get("IMAGE_TAG", ""), + } + + +def validate_values(values: Mapping[str, str]) -> None: + values = normalize_values(values) + missing = [key for key in ("IMAGE_REPOSITORY", "IMAGE_TAG") if not values.get(key)] + if missing: + raise ValueError("missing managed values: " + ", ".join(missing)) + registry = values.get("IMAGE_REGISTRY", "") + if registry and ( + not _REGISTRY.fullmatch(registry) + ): + raise ValueError("invalid registry") + if registry: + host = registry.split("/", 1)[0] + if host != "localhost" and "." not in host and ":" not in host: + raise ValueError("ambiguous registry; use an FQDN, explicit port, or localhost") + if not _REPOSITORY.fullmatch(values["IMAGE_REPOSITORY"]): + raise ValueError("invalid repository") + if not _TAG.fullmatch(values["IMAGE_TAG"]): + raise ValueError("invalid tag") + if any(any(character.isspace() for character in value) for value in values.values()): + raise ValueError("managed values cannot contain whitespace") + + +def image_reference(values: Mapping[str, str]) -> str: + values = normalize_values(values) + validate_values(values) + prefix = values.get("IMAGE_REGISTRY", "") + repository = values["IMAGE_REPOSITORY"] + return f"{prefix}/{repository}:{values['IMAGE_TAG']}" if prefix else f"{repository}:{values['IMAGE_TAG']}" + + +def update_env(path: Path, values: Mapping[str, str]) -> bool: + """Atomically update managed entries; return whether the file changed.""" + values = normalize_values(values) + validate_values(values) + if path.is_symlink(): + raise ValueError("refusing to update symlink .env") + old = path.read_text() if path.exists() else "" + parse_env(old) # detect malformed managed syntax before changing anything + lines = old.splitlines() + output: list[str] = [] + written: set[str] = set() + for line in lines: + stripped = line.strip() + export = re.match(r"export[ \t]+", stripped) + candidate = stripped[export.end() :].lstrip() if export else stripped + key = candidate.split("=", 1)[0].strip() if "=" in candidate else "" + if key in MANAGED_KEYS: + if key in written: + continue + output.append(f"{key}={values[key]}") + written.add(key) + else: + output.append(line) + output.extend(f"{key}={values.get(key, '')}" for key in MANAGED_KEYS if key not in written) + new = "\n".join(output) + "\n" + if old == new: + return False + path.parent.mkdir(parents=True, exist_ok=True) + mode = os.stat(path).st_mode & 0o777 if path.exists() else 0o600 + fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + try: + os.fchmod(fd, mode) + with os.fdopen(fd, "w") as handle: + handle.write(new) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + finally: + if os.path.exists(temporary): + os.unlink(temporary) + return True diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/next_image_tag.py b/.agents/skills/greptimedb-development-docker-image/scripts/next_image_tag.py new file mode 100755 index 00000000000..ad261ebd657 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/next_image_tag.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Print a tag with its final numeric component incremented.""" + +import argparse +import re +import sys + + +def increment_tag(tag: str) -> str: + match = re.search(r"(\d+)$", tag) + if match is None: + raise ValueError("tag must end with a numeric version component") + + number = match.group(1) + return f"{tag[:match.start()]}{int(number) + 1:0{len(number)}d}" + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Increment the final numeric component of a Docker image tag." + ) + parser.add_argument("--tag", required=True, help="existing Docker image tag") + args = parser.parse_args() + + try: + print(increment_tag(args.tag)) + except ValueError as error: + print(f"error: {error}: {args.tag!r}", file=sys.stderr) + return 2 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/prepare_context.sh b/.agents/skills/greptimedb-development-docker-image/scripts/prepare_context.sh new file mode 100755 index 00000000000..f0ea66fdf34 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/prepare_context.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash + +set -euo pipefail + +usage() { + printf 'Usage: %s --context PATH --binary PATH --dockerfile PATH --platform PLATFORM\n' "$0" >&2 +} + +context="" +binary="" +dockerfile="" +platform="" + +while [[ $# -gt 0 ]]; do + case "$1" in + --context) + context="$2" + shift 2 + ;; + --binary) + binary="$2" + shift 2 + ;; + --dockerfile) + dockerfile="$2" + shift 2 + ;; + --platform) + platform="$2" + shift 2 + ;; + *) + usage + exit 2 + ;; + esac +done + +if [[ -z "$context" || -z "$binary" || -z "$dockerfile" || -z "$platform" ]]; then + usage + exit 2 +fi + +script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +python3 "$script_dir/binary_platform.py" --binary "$binary" --platform "$platform" + +if [[ ! -f "$dockerfile" ]]; then + printf 'error: Dockerfile template does not exist: %s\n' "$dockerfile" >&2 + exit 1 +fi + +if [[ -e "$context" && ( -L "$context" || -n "$(ls -A "$context")" ) ]]; then + printf 'error: build context must be a new or empty directory: %s\n' "$context" >&2 + exit 1 +fi +mkdir -p "$context" +if [[ -e "$context/greptime" || -e "$context/Dockerfile" || -L "$context/greptime" || -L "$context/Dockerfile" ]]; then + printf 'error: build context already contains output files: %s\n' "$context" >&2 + exit 1 +fi +cp "$binary" "$context/greptime" +cp "$dockerfile" "$context/Dockerfile" +printf 'created Dockerfile: %s\n' "$context/Dockerfile" + +file "$context/greptime" diff --git a/.agents/skills/greptimedb-development-docker-image/scripts/update_image_env.py b/.agents/skills/greptimedb-development-docker-image/scripts/update_image_env.py new file mode 100755 index 00000000000..9bf8cbeea04 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/scripts/update_image_env.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Create or update non-secret GreptimeDB image settings in a .env file.""" + +import argparse +import sys +from pathlib import Path + +from image_config import image_reference, normalize_values, update_env, validate_values + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--env", type=Path) + parser.add_argument("--registry", required=True) + parser.add_argument("--repository", required=True) + parser.add_argument("--tag", required=True) + parser.add_argument("--validate-only", action="store_true") + args = parser.parse_args() + + try: + values = normalize_values({"IMAGE_REGISTRY": args.registry, "IMAGE_REPOSITORY": args.repository, "IMAGE_TAG": args.tag}) + validate_values(values) + except ValueError as error: + print(f"error: {error}", file=sys.stderr) + return 2 + + if args.validate_only: + print(image_reference(values)) + return 0 + if args.env is None: + parser.error("--env is required unless --validate-only is set") + if update_env(args.env, values): + print(f"updated {args.env}") + else: + print(f"unchanged {args.env}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.agents/skills/greptimedb-development-docker-image/tests/test_binary_platform.py b/.agents/skills/greptimedb-development-docker-image/tests/test_binary_platform.py new file mode 100644 index 00000000000..e89bafa0696 --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/tests/test_binary_platform.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib +import struct +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +SCRIPTS = Path(__file__).resolve().parents[1] / "scripts" +sys.path.insert(0, str(SCRIPTS)) + +binary_platform = importlib.import_module("binary_platform") + + +def elf(machine: int, *, bits: int = 2, endian: int = 1, abi: int = 3, elf_type: int = 3) -> bytes: + header = bytearray(64) + header[:4] = b"\x7fELF" + header[4:8] = bytes((bits, endian, 1, abi)) + struct.pack_into("H", header, 18, machine) + struct.pack_into(" Path: + path = Path(self.directory.name) / "greptime" + path.write_bytes(content) + path.chmod(0o755) + return path + + def setUp(self): + self.directory = tempfile.TemporaryDirectory() + + def tearDown(self): + self.directory.cleanup() + + def test_matching_architectures(self): + binary_platform.validate_binary(self.make_binary(elf(0x3E)), "linux/amd64") + binary_platform.validate_binary(self.make_binary(elf(0xB7)), "linux/arm64") + + def test_rejects_bad_executable_elf_and_platform(self): + path = self.make_binary(elf(0x3E)) + path.chmod(0o644) + with self.assertRaises(ValueError): + binary_platform.validate_binary(path, "linux/amd64") + path = self.make_binary(b"not an elf") + with self.assertRaises(ValueError): + binary_platform.validate_binary(path, "linux/amd64") + path = self.make_binary(elf(0x3E, bits=1)) + with self.assertRaises(ValueError): + binary_platform.validate_binary(path, "linux/amd64") + with self.assertRaises(ValueError): + binary_platform.validate_binary(self.make_binary(elf(0x3E, abi=9)), "linux/amd64") + with self.assertRaises(ValueError): + binary_platform.validate_binary(self.make_binary(elf(0x3E, elf_type=1)), "linux/amd64") + with self.assertRaises(ValueError): + binary_platform.validate_binary(self.make_binary(b"\x7fELF"), "linux/amd64") + with self.assertRaises(ValueError): + binary_platform.validate_binary(self.make_binary(elf(0x3E)), "darwin/amd64") + + def test_rejects_wrong_architecture_and_symlink(self): + with self.assertRaises(ValueError): + binary_platform.validate_binary(self.make_binary(elf(0x3E)), "linux/arm64") + target = self.make_binary(elf(0x3E)) + link = Path(self.directory.name) / "link" + link.symlink_to(target) + with self.assertRaises(ValueError): + binary_platform.validate_binary(link, "linux/amd64") + + def test_cli_success_and_failures(self): + path = self.make_binary(elf(0x3E)) + assert binary_platform.__file__ is not None + script = Path(binary_platform.__file__) + ok = subprocess.run([sys.executable, str(script), "--binary", str(path), "--platform", "linux/amd64"], capture_output=True, text=True) + self.assertEqual(ok.returncode, 0, ok.stderr) + bad = subprocess.run([sys.executable, str(script), "--binary", str(path), "--platform", "linux/arm64"], capture_output=True, text=True) + self.assertEqual(bad.returncode, 2) + missing = subprocess.run([sys.executable, str(script), "--binary", str(path) + ".missing", "--platform", "linux/amd64"], capture_output=True, text=True) + self.assertEqual(missing.returncode, 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/.agents/skills/greptimedb-development-docker-image/tests/test_image_config.py b/.agents/skills/greptimedb-development-docker-image/tests/test_image_config.py new file mode 100644 index 00000000000..95dc913ea5d --- /dev/null +++ b/.agents/skills/greptimedb-development-docker-image/tests/test_image_config.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib +import stat +import sys +import tempfile +import unittest +from unittest import mock +from pathlib import Path + +SCRIPTS = Path(__file__).resolve().parents[1] / "scripts" +sys.path.insert(0, str(SCRIPTS)) + +image_config = importlib.import_module("image_config") + + +VALUES = {"IMAGE_REGISTRY": "registry.example.com/team", "IMAGE_REPOSITORY": "greptime-dev", "IMAGE_TAG": "dev-001"} + + +class ImageConfigTest(unittest.TestCase): + def test_parse_supported_bounded_dotenv(self): + text = """\n# comment\n export IMAGE_REGISTRY = \"registry.example.com/team\" # note\nIMAGE_REPOSITORY='greptime-dev'\nIMAGE_TAG=dev-001 # safe comment\nOTHER=$HOME\n""" + self.assertEqual(image_config.parse_env(text), VALUES) + + def test_parser_does_not_validate_values(self): + self.assertEqual(image_config.parse_env("IMAGE_TAG=not valid\n"), {"IMAGE_TAG": "not valid"}) + + def test_parser_rejects_unmatched_and_unsupported_syntax(self): + with self.assertRaises(ValueError): + image_config.parse_env('IMAGE_TAG="unterminated\n') + with self.assertRaises(ValueError): + image_config.parse_env("IMAGE_TAG=$(date)\n") + + def test_validation_and_reference(self): + image_config.validate_values(VALUES) + self.assertEqual(image_config.image_reference(VALUES), "registry.example.com/team/greptime-dev:dev-001") + local = dict(VALUES, IMAGE_REGISTRY="") + self.assertEqual(image_config.image_reference(local), "greptime-dev:dev-001") + image_config.validate_values(dict(VALUES, IMAGE_REGISTRY="localhost")) + image_config.validate_values(dict(VALUES, IMAGE_REGISTRY="example.com")) + with self.assertRaises(ValueError): + image_config.validate_values(dict(VALUES, IMAGE_REGISTRY="registry/team")) + + def test_atomic_update_preserves_mode_and_canonicalizes_duplicates(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / ".env" + path.write_text("# keep\nIMAGE_TAG=old\nIMAGE_TAG=duplicate\nOTHER=value\n") + path.chmod(0o640) + self.assertTrue(image_config.update_env(path, VALUES)) + self.assertEqual(stat.S_IMODE(path.stat().st_mode), 0o640) + self.assertEqual(path.read_text(), "# keep\nIMAGE_TAG=dev-001\nOTHER=value\nIMAGE_REGISTRY=registry.example.com/team\nIMAGE_REPOSITORY=greptime-dev\n") + self.assertFalse(image_config.update_env(path, VALUES)) + + def test_update_rejects_symlink(self): + with tempfile.TemporaryDirectory() as directory: + target = Path(directory) / "target" + target.write_text("IMAGE_TAG=old\n") + link = Path(directory) / ".env" + link.symlink_to(target) + with self.assertRaises(ValueError): + image_config.update_env(link, VALUES) + self.assertEqual(target.read_text(), "IMAGE_TAG=old\n") + + def test_update_clears_omitted_registry_and_uses_private_new_file_mode(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / ".env" + path.write_text("IMAGE_REGISTRY=registry.example.com\nIMAGE_REPOSITORY=old\nIMAGE_TAG=old\n") + values = {"IMAGE_REPOSITORY": "greptime-dev", "IMAGE_TAG": "dev-001"} + image_config.update_env(path, values) + self.assertIn("IMAGE_REGISTRY=\n", path.read_text()) + new_path = Path(directory) / "new.env" + image_config.update_env(new_path, VALUES) + self.assertEqual(stat.S_IMODE(new_path.stat().st_mode), 0o600) + + def test_failed_replace_keeps_original_and_removes_temporary_file(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / ".env" + path.write_text("OTHER=value\n") + with mock.patch("image_config.os.replace", side_effect=OSError("no space")): + with self.assertRaises(OSError): + image_config.update_env(path, VALUES) + self.assertEqual(path.read_text(), "OTHER=value\n") + self.assertEqual(list(Path(directory).glob("..env.*")), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/actions/setup-greptimedb-cluster/with-remote-wal-logical-prune.yaml b/.github/actions/setup-greptimedb-cluster/with-remote-wal-logical-prune.yaml index 60c23def6e2..3b25e80a8f2 100644 --- a/.github/actions/setup-greptimedb-cluster/with-remote-wal-logical-prune.yaml +++ b/.github/actions/setup-greptimedb-cluster/with-remote-wal-logical-prune.yaml @@ -3,6 +3,7 @@ logging: format: "json" filters: - log_store=debug + - rskafka=debug meta: configData: |- [runtime] diff --git a/.github/actions/setup-greptimedb-cluster/with-remote-wal.yaml b/.github/actions/setup-greptimedb-cluster/with-remote-wal.yaml index ec595b9e9b8..6a15b25b422 100644 --- a/.github/actions/setup-greptimedb-cluster/with-remote-wal.yaml +++ b/.github/actions/setup-greptimedb-cluster/with-remote-wal.yaml @@ -3,6 +3,7 @@ logging: format: "json" filters: - log_store=debug + - rskafka=debug meta: configData: |- [runtime] diff --git a/.github/actions/setup-kafka-cluster/action.yml b/.github/actions/setup-kafka-cluster/action.yml index fca420da709..3c30ca285e3 100644 --- a/.github/actions/setup-kafka-cluster/action.yml +++ b/.github/actions/setup-kafka-cluster/action.yml @@ -16,8 +16,8 @@ runs: helm upgrade \ --install kafka oci://registry-1.docker.io/bitnamicharts/kafka \ --set controller.replicaCount=${{ inputs.controller-replicas }} \ - --set controller.resources.requests.cpu=50m \ - --set controller.resources.requests.memory=128Mi \ + --set controller.resources.requests.cpu=500m \ + --set controller.resources.requests.memory=512Mi \ --set controller.resources.limits.cpu=2000m \ --set controller.resources.limits.memory=2Gi \ --set listeners.controller.protocol=PLAINTEXT \ diff --git a/.github/runner-scale-sets/query-regression/Dockerfile b/.github/runner-scale-sets/query-regression/Dockerfile index 1f9f7203b9d..61392919587 100644 --- a/.github/runner-scale-sets/query-regression/Dockerfile +++ b/.github/runner-scale-sets/query-regression/Dockerfile @@ -3,6 +3,19 @@ FROM ghcr.io/actions/actions-runner@sha256:08c30b0a7105f64bddfc485d2487a22aa0393 USER root +ARG SCCACHE_VERSION=0.16.0 +ARG SCCACHE_SHA256=aec995a83ad3dff3d14b6314e08858b7b73d35ca85a5bcf3d3a9ec07dee35588 +ARG RUSTUP_INIT_VERSION=1.29.0 +ARG RUSTUP_INIT_TARGET=x86_64-unknown-linux-gnu +ARG RUSTUP_INIT_SHA256=4acc9acc76d5079515b46346a485974457b5a79893cfb01112423c89aeb5aa10 +ARG RUST_TOOLCHAIN=nightly-2026-03-21 + +ENV RUSTUP_HOME=/opt/rustup \ + CARGO_HOME=/opt/cargo \ + RUSTUP_TOOLCHAIN=${RUST_TOOLCHAIN} \ + RUSTUP_AUTO_INSTALL=0 \ + PATH=/opt/cargo/bin:${PATH} + RUN apt-get update \ && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ build-essential \ @@ -13,6 +26,7 @@ RUN apt-get update \ git \ gzip \ jq \ + libprotobuf-dev \ libssl-dev \ mold \ openssh-client \ @@ -30,4 +44,58 @@ RUN apt-get update \ COPY --from=uv /uv /uvx /usr/local/bin/ +RUN curl --fail --location --silent --show-error \ + --output /tmp/sccache.tar.gz \ + "https://github.com/mozilla/sccache/releases/download/v${SCCACHE_VERSION}/sccache-v${SCCACHE_VERSION}-x86_64-unknown-linux-musl.tar.gz" \ + && echo "${SCCACHE_SHA256} /tmp/sccache.tar.gz" | sha256sum --check --status - \ + && mkdir --parents /tmp/sccache \ + && tar --extract --gzip --file /tmp/sccache.tar.gz \ + --directory /tmp/sccache --strip-components=1 \ + "sccache-v${SCCACHE_VERSION}-x86_64-unknown-linux-musl/sccache" \ + && install --mode=0755 /tmp/sccache/sccache /usr/local/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache + +RUN curl --fail --location --silent --show-error \ + --output /tmp/rustup-init \ + "https://static.rust-lang.org/rustup/archive/${RUSTUP_INIT_VERSION}/${RUSTUP_INIT_TARGET}/rustup-init" \ + && echo "${RUSTUP_INIT_SHA256} /tmp/rustup-init" | sha256sum --check --status - \ + && chmod 0755 /tmp/rustup-init \ + && /tmp/rustup-init -y --profile minimal --default-toolchain "${RUST_TOOLCHAIN}" --no-modify-path \ + && rm -f /tmp/rustup-init \ + && chown -R root:root /opt/rustup /opt/cargo \ + && chmod -R go-w /opt/rustup /opt/cargo \ + && find /opt/rustup /opt/cargo -type d -exec chmod 0755 {} + \ + && find /opt/rustup /opt/cargo -type f -exec chmod a+r {} + + USER runner + +RUN set -eu \ + && test "$(id -u)" = "1001" \ + && temporary_cargo_home="$(mktemp --directory)" \ + && temporary_proto_dir="" \ + && cleanup() { rm -rf "${temporary_cargo_home}" "${temporary_proto_dir}"; } \ + && trap cleanup 0 \ + && export CARGO_HOME="${temporary_cargo_home}" \ + && test "${RUSTUP_AUTO_INSTALL}" = "0" \ + && rustup --version \ + && cargo --version \ + && rustc --version \ + && active_toolchain="$(rustup show active-toolchain)" \ + && printf 'Active toolchain: %s\n' "${active_toolchain}" \ + && case "${active_toolchain}" in "${RUST_TOOLCHAIN}-x86_64-unknown-linux-gnu"|"${RUST_TOOLCHAIN}-x86_64-unknown-linux-gnu "*) ;; *) exit 1;; esac \ + && test ! -w /opt/rustup \ + && test ! -w /opt/cargo/bin \ + && test -r /usr/include/google/protobuf/any.proto \ + && test -r /usr/include/google/protobuf/empty.proto \ + && temporary_proto_dir="$(mktemp --directory)" \ + && printf '%s\n' \ + 'syntax = "proto3";' \ + 'package smoke;' \ + 'import "google/protobuf/any.proto";' \ + 'import "google/protobuf/empty.proto";' \ + 'message Smoke { google.protobuf.Any any = 1; google.protobuf.Empty empty = 2; }' \ + > "${temporary_proto_dir}/smoke.proto" \ + && protoc --proto_path="${temporary_proto_dir}" --proto_path=/usr/include \ + --descriptor_set_out="${temporary_proto_dir}/smoke.pb" \ + "${temporary_proto_dir}/smoke.proto" \ + && test -s "${temporary_proto_dir}/smoke.pb" diff --git a/.github/runner-scale-sets/query-regression/README.md b/.github/runner-scale-sets/query-regression/README.md index eb1a7e12f79..4fe8c0f3afb 100644 --- a/.github/runner-scale-sets/query-regression/README.md +++ b/.github/runner-scale-sets/query-regression/README.md @@ -1,21 +1,14 @@ # Query regression self-hosted runners -The `Query Regression` workflow targets self-hosted GitHub Actions runners via -runner labels or ARC runner scale set names: +The `Query Regression` workflow uses the dedicated ARC runner scale set +`perf-regression-8-cores`. ARC runner Pods run in the target Kubernetes cluster +and connect outbound to GitHub. The live scale set is currently **paused**: +`minRunners=0`, `maxRunners=0`, and no runner Pods. Do not resume it without +explicit approval. -- `perf-regression-8-cores` +## Prerequisites and trust admission -The name intentionally avoids generic labels such as `ubuntu-22.04-8-cores`, -which may already be used by GitHub-hosted larger runners or other runner pools. - -For Kubernetes-based runners, ARC runner pods run inside the target Kubernetes -cluster and connect outbound to GitHub; GitHub can then dispatch jobs whose -`runs-on` value matches the runner scale set name or runner labels. - -## Prerequisites - -Install the ARC scale set controller in the Kubernetes cluster if it is not -already installed: +Install the ARC scale set controller if it is not already installed: ```bash helm upgrade --install arc \ @@ -25,8 +18,8 @@ helm upgrade --install arc \ --version 0.14.2 ``` -Create the GitHub App or PAT secret in the namespace that will host the runner -scale sets. Prefer a GitHub App with access limited to `GreptimeTeam/greptimedb`. +Create the GitHub App secret in the runner namespace. Prefer an App limited to +`GreptimeTeam/greptimedb`: ```bash kubectl -n arc-runners create secret generic greptimedb-arc-github-app \ @@ -35,14 +28,47 @@ kubectl -n arc-runners create secret generic greptimedb-arc-github-app \ --from-file=github_app_private_key= ``` -The values files in this directory reference that secret by name. +The values files here reference that secret by name. -## Build the runner image +A maintainer applying the `query-regression` label is **trust admission for +that exact PR revision**. The admitted job may use this scale set's dedicated, +writable persistent cache. `pull_request: labeled` is the only PR trigger: the +label event snapshots its merge, head, and base SHAs. A queued job fetches that +immutable event merge SHA directly, verifies it is a two-parent merge whose +parents include the snapshotted head exactly once, and uses its other parent as +the actual base build revision. The snapshotted event base is retained for audit +only, so a difference from the merge's non-head parent is not a failure. The job +never follows a newer mutable PR merge ref. An unavailable event merge, or one +that does not contain exactly one snapshotted head parent, fails closed. A later +PR head change does not retarget an already queued run: it may execute only its +previously trusted event revision if that revision remains fetchable. To run the +new revision, the maintainer must review it, remove the label, and re-add +`query-regression`; cancel the old run if it is no longer wanted. An existing +label does not automatically rerun the benchmark. -Use a derived ARC runner image instead of the minimal upstream runner image. The -image keeps the official `/home/runner/run.sh` entrypoint layout and adds the -tools this workflow expects to be present on a normal CI host, including `wget`, -`uv`, C/C++ build tools, OpenSSL headers, and protobuf tooling. +Admission does not relax runner hardening or GitHub permissions. Keep +service-account token mounting disabled; do not mount host paths, the Docker +socket, kubeconfig, or long-lived credentials. The runner and cache initializer +use UID/GID 1001, disallow privilege escalation, drop all capabilities, and use +the RuntimeDefault seccomp profile. Keep GitHub tokens least-privilege and +review workflow changes before admission. Where the CNI supports it, restrict +egress to required GitHub Actions, artifact/cache, Rust/crate/toolchain, DNS, +and image-registry endpoints; block unrelated cluster services, private ranges, +and metadata endpoints unless a case requires them. + +### Network routing prerequisite + +Required split routing is an **external environment-specific prerequisite**. The +responsible network operator must route GitHub Actions, GitHub content, +artifact/cache, crates.io, Rust toolchain, and image-registry traffic through +the approved path rather than the VPN where applicable. Neither this repository +nor Kubernetes configures that route. Verify it with the responsible network +operator before any canary. + +## Runner image and workflow tools + +Build and push the derived runner image; it preserves the official +`/home/runner/run.sh` entrypoint and supplies CI tools needed at runtime: ```bash docker build \ @@ -53,100 +79,232 @@ docker build \ docker push greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner:latest ``` -Deploy the runner image by digest rather than by a mutable tag. Update -`values-8-cores.yaml` after pushing a rebuilt image. +Deploy by digest, not mutable tag, by updating `values-8-cores.yaml` after a +rebuild. If the registry is private, use a dedicated read-only pull secret only +as `imagePullSecrets`; never expose registry credentials to runner containers. +Both digest-pinned init and runner containers use `IfNotPresent`: the immutable +digest makes a cached image safe and avoids adding a registry dependency to every +runner startup. -The workflow still runs setup actions for pinned Rust and `uv` behavior. `mold` -is installed in the image and selected through `CARGO_BUILD_RUSTFLAGS`, so jobs -do not need privileged package installation at runtime. +The runner optionally imports only the non-sensitive `HTTP_PROXY`, +`HTTPS_PROXY`, and `NO_PROXY` variables from the +`query-regression-runner-local-env` ConfigMap. Manage that ConfigMap locally in +the target namespace; private endpoint configuration must not be committed, and +credentials or secrets must never be placed in a ConfigMap. -Decide whether the registry repository is public or private. Public pull access -avoids distributing registry credentials to the runner namespace. If the image -must be private, create a dedicated read-only image-pull secret for this image -and attach it only as `imagePullSecrets`; do not expose registry credentials to -runner containers. +Before builds, the workflow asserts UID/GID 1001 and exact image tool versions: +`libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.30.0`, `Python 3.12.3`, `sccache +0.16.0`, root-owned `rustup 1.29.0`, and the image-baked +`nightly-2026-03-21` Rust toolchain. Rustup, Cargo, and Rustc must resolve from +`/opt/cargo/bin`; the runner cannot write `/opt/rustup` or `/opt/cargo/bin`. +Protobuf well-known includes, including `google/protobuf/any.proto` and +`google/protobuf/empty.proto`, are an image contract and must compile with +`protoc`. +`actions-rust-lang/setup-rust-toolchain@v1` is intentionally removed. The +workflow sets its warning-denying mold `RUSTFLAGS` directly, disables automatic +Rustup installation, and performs no runtime toolchain downloads. -## Install the query-regression scale set +The workflow no longer uses GitHub `rust-cache`, `setup-protoc`, `setup-uv`, or +runtime Rust setup: the image establishes immutable executable state and the PVC +supplies only reusable Cargo data. Do not reintroduce those actions unless the +corresponding cache or image contract changes. -Install the runner scale set. The Helm release name and `runnerScaleSetName` -should match the `runs-on` value used by the workflow. +## Capacity and persistent cache + +`values-8-cores.yaml` is normal operation: `minRunners=0`, `maxRunners=1`. +`values-paused.yaml` is the mandatory pause overlay: `minRunners=0`, +`maxRunners=0`. The job uses group `query-regression-persistent-cache-v1`, +`queue: max`, and `cancel-in-progress: false`; admitted jobs queue rather than +replacing older pending jobs. During maintenance, cancel admitted queued runs as +well as pausing ARC. Runner Pods have `activeDeadlineSeconds=12600`. +The runner requests 6 CPU and limits at 8 CPU to preserve `minipc-3` +allocatable-capacity scheduling headroom; do not reset the request to 8 CPU +without revalidating scheduling capacity. + +The cache claim `query-regression-build-cache` is a nominal 600Gi `local-path` +PVC in `arc-runners`. It is `ReadWriteOnce`; `local-path` uses +WaitForFirstConsumer binding and Delete reclaim behavior, produces a +node-affine local PV, is non-expandable, and the 600Gi request is not a hard +storage quota. The runner's `minipc-3` selector is its only consumer candidate. + +The initializer mounts the PVC root at `/cache`, creates and write-tests these +versioned subpaths as non-root UID/GID 1001, and the runner mounts them as: + +| Persistent state | PVC subpath | Runner mount | +| --- | --- | --- | +| Ephemeral Cargo home | `emptyDir` | `/home/runner/.cargo` | +| Cargo registry data | `cargo-registry-v1` | `/home/runner/.cargo/registry` | +| Cargo Git data | `cargo-git-v1` | `/home/runner/.cargo/git` | +| Cargo target | `query-regression-target-v1` | `/home/runner/query-regression-target` | +| Cache metadata | `meta-v1` | `/home/runner/query-regression-cache-meta` | +| sccache local disk cache | `sccache-v1` | `/home/runner/.cache/sccache` | +| Immutable Rust toolchain | image-owned | `/opt/rustup`, `/opt/cargo/bin` | + +The Pod security context uses UID/GID and `fsGroup` 1001 with +`fsGroupChangePolicy: OnRootMismatch`; no privileged `chown` or raw `hostPath` +is used. `CARGO_HOME` is a per-Pod `emptyDir`; only its nested `registry` and +`git` mounts are persistent. `RUSTUP_HOME=/opt/rustup` and `/opt/cargo/bin` are +image-owned immutable paths, while `CARGO_TARGET_DIR`, cache metadata, and +`SCCACHE_DIR` are persistent absolute paths. The runner sets +`RUSTC_WRAPPER=/usr/local/bin/sccache`, +`SCCACHE_DIR=/home/runner/.cache/sccache`, `SCCACHE_CACHE_SIZE=40G`, and +`CARGO_INCREMENTAL=0`. sccache uses its local PVC disk backend and self-evicts +at 40G; do not add runtime downloads, object storage, or a shared backend. + +The repository's `.cargo/config.toml` remains a trusted per-revision build input. +In contrast, `$CARGO_HOME/config*`, credentials, installed bins, and Cargo +metadata outside the persistent `registry` and `git` data mounts are ephemeral +and cannot survive to another Pod. + +The local disk backend has a one-server constraint. `maxRunners=1` and the +unchanged `query-regression-persistent-cache-v1` workflow concurrency group +serialize runs; do not increase runner capacity or relax that serialization +while this backend is in use. Base and candidate builds share the target; Cargo +fingerprints invalidate source and dependency changes. The workflow records the +sccache version and relevant environment in the target ABI marker, starts and +zeros sccache after cache and toolchain checks, shows initial/base/candidate +statistics, and resets statistics between base and candidate builds. + +### Disk preflight and cleanup contract + +Before applying or unpausing, verify the backing filesystem on `minipc-3` has +at least 900GiB free. The current local-path provisioner source is +`/opt/local-path-provisioner`; measure the filesystem containing it: + +```bash +df -PB1G /opt/local-path-provisioner +``` + +The workflow reports `du`, `df -P`, human-readable free space, and inode +availability before builds and in an always-run report. Its cleanup is narrow +and non-destructive: + +- warn at target size 400GiB; at 450GiB clear only the complete target root; +- warn at Cargo registry-plus-Git data size 60GiB; at 80GiB remove only + `registry/src` and `git/checkouts`, then abort if that persistent data remains + at least 80GiB; +- below 300GiB backing free space, clear the complete target root first, + remeasure, then remove only those Cargo extracted trees and checkouts; abort + if free space is still below 300GiB; +- never automatically remove Cargo registry cache/index, Git database, the + image-owned Cargo bin or Rustup toolchain, cache metadata, the self-evicting + sccache directory, or the PVC. + +The target clear uses fixed absolute roots and removes all entries, including +dotfiles. After migration validation, remove obsolete `cargo-home-v1` and +`rustup-home-v1` only in explicit maintenance while ARC is 0/0 and no runner Pod +exists; they are not mounted by the current configuration. + +## Deploy and pause safely + +First verify the configured external network route and the disk preflight. Apply the PVC; while +the scale set is paused, it is expected to remain `Pending` because +WaitForFirstConsumer has no scheduled runner: + +```bash +kubectl apply --dry-run=server \ + -f .github/runner-scale-sets/query-regression/cache-pvc.yaml +kubectl apply -f .github/runner-scale-sets/query-regression/cache-pvc.yaml +``` + +Render normal and paused configurations. Normal values are always first; the +pause overlay is always last: + +```bash +helm template perf-regression-8-cores \ + oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ + --namespace arc-runners --version 0.14.2 \ + --set controllerServiceAccount.name=arc-gha-rs-controller \ + --set controllerServiceAccount.namespace=arc-systems \ + -f .github/runner-scale-sets/query-regression/values-8-cores.yaml + +helm template perf-regression-8-cores \ + oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ + --namespace arc-runners --version 0.14.2 \ + --set controllerServiceAccount.name=arc-gha-rs-controller \ + --set controllerServiceAccount.namespace=arc-systems \ + -f .github/runner-scale-sets/query-regression/values-8-cores.yaml \ + -f .github/runner-scale-sets/query-regression/values-paused.yaml +``` + +The **first post-merge Helm deployment must reconcile the release in paused +mode**. Keep the pause overlay last: + +```bash +# First post-merge deployment and every return to paused mode: 0/0. +helm upgrade --install perf-regression-8-cores \ + oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ + --namespace arc-runners --create-namespace --version 0.14.2 \ + --reset-values --wait \ + -f .github/runner-scale-sets/query-regression/values-8-cores.yaml \ + -f .github/runner-scale-sets/query-regression/values-paused.yaml + +# Expect 0/0 and no runner resources before considering normal mode. +kubectl -n arc-runners get autoscalingrunnerset perf-regression-8-cores \ + -o jsonpath='{.spec.minRunners}{"/"}{.spec.maxRunners}{"\n"}' +kubectl -n arc-runners get ephemeralrunners,pods \ + -l actions.github.com/scale-set-name=perf-regression-8-cores +``` + +Only after that verification and separate explicit approval, apply normal 0/1 +operation without the pause overlay: ```bash helm upgrade --install perf-regression-8-cores \ oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners \ - --create-namespace \ - --version 0.14.2 \ + --namespace arc-runners --create-namespace --version 0.14.2 \ + --reset-values --wait \ -f .github/runner-scale-sets/query-regression/values-8-cores.yaml ``` -Check registration and pods: +Do not use bare `helm rollback`, `--atomic`, or `--reuse-values`: a stored +revision can restore nonzero runner capacity. Inspect rendered manifests for +capacity, the `minipc-3` selector, cache claim and mounts, initializer, +security context, and resources. After approved normal mode receives its first +canary, the PVC binds to `minipc-3`. -```bash -kubectl -n arc-runners get pods -kubectl -n arc-runners get autoscalingrunnersets -``` +For local-PV node loss, cache recovery is intentionally disposable: return to +0/0, recreate the PVC on a healthy node, cold-fill it, and run a new canary. -The scale sets should also appear under repository Actions runner settings. +## Canary and rollback -## Security notes for fork PRs +With explicit approval, run two identical `workflow_dispatch` canaries on +`perf-regression-8-cores`, using immutable full base and candidate commit SHAs +and `cargo_profile=nightly`. The first is the cold fill; the second verifies warm +reuse. Record the workflow's base/candidate build elapsed logs and cache +ABI-marker output, initial/base/candidate sccache statistics, and cache report. +Confirm the image tool contract (root-owned Rustup/Cargo paths, baked nightly, +and non-writable `/opt` roots) and that the ephemeral Cargo home contains only +the mounted registry/Git data before Cargo creates per-Pod state. +Obtain dependency and tool network byte counters from the configured +environment counter source, filtered to `minipc-3` and the dependency/tool +destinations. -Maintainer-approved fork PRs can run on self-hosted runners. Approval only lets -the workflow execute; it does not make the fork code trusted. +Accept the canary only when all of the following hold: -The `query-regression` label is the explicit trigger for PR runs. Updating a PR -does not rerun the benchmark automatically; remove and re-add the label after -reviewing the updated changes. +- exactly one runner Pod runs on `minipc-3`, and both jobs use the same bound PV; +- UID/GID 1001 cache mounts are writable; the warm run does not invalidate the + target or bulk-redownload crates or toolchains; sccache reports separate base + and candidate build statistics without server or cache-path errors; immutable + Rustup/Cargo roots remain non-writable and only registry/Git data persists; +- warm base build time is at most 50% of cold base build time; +- warm dependency/tool network bytes are at most 10% of cold fill bytes; +- cache sizes remain below soft watermarks, node free space remains at least + 300GiB, and the benchmark is correct without TLS EOFs or timeouts; +- the configured environment counter source confirms this traffic is outside VPN + accounting. -These runners execute PR code and should be treated as untrusted execution -capacity: +Immediately return to 0/0 after either canary unless ongoing normal operation +has been explicitly approved; return immediately on any traffic, cache, disk, +TLS, or correctness failure. To roll back, use the paused Helm upgrade above, +or another explicit `helm upgrade --install` with known-good values followed by +`values-paused.yaml`, `--reset-values`, and `--wait`. Do not delete the PVC +automatically; preserve it for diagnosis unless intentionally discarding cache. -- keep them isolated from sensitive internal services unless explicitly required; -- do not mount host paths, Docker socket, kubeconfig, or long-lived credentials; -- use ephemeral runner pods and no shared work directory with trusted jobs; -- disable service account token mounting in runner pods unless Kubernetes API - access is intentionally required; -- use a runner image whose default user is non-root, disable privilege - escalation, drop Linux capabilities, and use the runtime-default seccomp - profile; -- keep GitHub tokens least-privilege and rely on normal `pull_request` behavior - for fork PRs, where repository secrets are withheld and `GITHUB_TOKEN` is - read-only; -- review fork workflow changes before approving the run. +## Future optional phases -If stronger isolation is required, install a separate runner group/namespace just -for query-regression PR workloads and restrict repository/workflow access to the -`Query Regression` workflow. - -Use namespace or cluster network policy to restrict runner egress where the CNI -supports the needed controls. Query-regression runners need outbound access to -GitHub Actions services, GitHub artifact/cache endpoints, Rust/crate/toolchain -endpoints, DNS, and the configured image registries. Block access to unrelated -cluster services, private network ranges, and cloud metadata endpoints unless a -case explicitly needs them. - -## Build cache - -The workflow builds base and candidate in the same job, the same source path, -and a shared `CARGO_TARGET_DIR`. It checks out the base ref into `src`, builds -and copies the base binary aside, then resets that same `src` checkout to the -candidate ref before building the candidate binary. Keeping the workspace path -stable improves Cargo incremental reuse for local workspace crates compared with -building separate `base-src` and `candidate-src` checkouts. - -The workflow also uses the GitHub Actions Rust cache for restore-only cache -reuse. PR and dispatch runs do not save cache entries. Refresh shared caches from -trusted maintenance workflows only. - -The ARC values in this directory do not configure a cross-run runner-local -compiler cache. If cross-run compile time still dominates the benchmark, -deployers may add one of the following at the runner infrastructure layer: - -- a custom runner image with Rust tooling, `mold`, and `sccache` preinstalled; -- `RUSTC_WRAPPER=sccache` plus a shared `sccache` backend such as object storage, - Redis, or a Kubernetes storage class that is safe for concurrent runner pods; -- separate cache namespaces/buckets for untrusted PR code and trusted branches to - avoid cache poisoning across trust boundaries. - -Do not mount a shared writable host path, Docker socket, kubeconfig, or other -privileged credentials into runners that execute PR code. +The current phase uses a digest-pinned image with sccache 0.16.0 and no shared +cache service. Optional follow-ups are an image additionally seeded with the +exact Rust toolchain and `cargo fetch --locked`; or an internal read/write +sccache backend or Cargo/Git mirror. Evaluate them only if persistent PVC reuse +is insufficient. diff --git a/.github/runner-scale-sets/query-regression/cache-pvc.yaml b/.github/runner-scale-sets/query-regression/cache-pvc.yaml new file mode 100644 index 00000000000..692f3a2d6e2 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/cache-pvc.yaml @@ -0,0 +1,15 @@ +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: query-regression-build-cache + namespace: arc-runners + labels: + app.kubernetes.io/part-of: perf-regression-8-cores +spec: + volumeMode: Filesystem + accessModes: + - ReadWriteOnce + storageClassName: local-path + resources: + requests: + storage: 600Gi diff --git a/.github/runner-scale-sets/query-regression/values-8-cores.yaml b/.github/runner-scale-sets/query-regression/values-8-cores.yaml index 54ade99bad1..057ff55b507 100644 --- a/.github/runner-scale-sets/query-regression/values-8-cores.yaml +++ b/.github/runner-scale-sets/query-regression/values-8-cores.yaml @@ -3,31 +3,144 @@ githubConfigSecret: greptimedb-arc-github-app runnerScaleSetName: "perf-regression-8-cores" minRunners: 0 -maxRunners: 4 +maxRunners: 1 template: spec: automountServiceAccountToken: false - activeDeadlineSeconds: 11400 + activeDeadlineSeconds: 12600 + nodeSelector: + kubernetes.io/hostname: minipc-3 securityContext: + runAsNonRoot: true + runAsUser: 1001 + runAsGroup: 1001 + fsGroup: 1001 + fsGroupChangePolicy: OnRootMismatch seccompProfile: type: RuntimeDefault - containers: - - name: runner - image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:108b7357532e5cf3329853373f561b0c674e954d6307b53d37075df5a189b0b2 - imagePullPolicy: Always - command: ["/home/runner/run.sh"] + volumes: + - name: cargo-home + emptyDir: {} + - name: build-cache + persistentVolumeClaim: + claimName: query-regression-build-cache + initContainers: + - name: initialize-build-cache + image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:2436278bf0756700c718927907cdbedac64d55c593f80b515acb44dec8b02cc1 + imagePullPolicy: IfNotPresent + command: + - /bin/sh + - -ec + - | + umask 0002 + for directory in cargo-registry-v1 cargo-git-v1 query-regression-target-v1 meta-v1 sccache-v1; do + cache_directory="/cache/${directory}" + mkdir -p "${cache_directory}" + test -w "${cache_directory}" + probe_file="${cache_directory}/.write-probe" + : > "${probe_file}" + test -f "${probe_file}" + rm "${probe_file}" + done securityContext: + runAsNonRoot: true + runAsUser: 1001 + runAsGroup: 1001 allowPrivilegeEscalation: false capabilities: drop: - ALL resources: requests: - cpu: "8" + cpu: 100m + memory: 128Mi + ephemeral-storage: 1Gi + limits: + cpu: 500m + memory: 512Mi + ephemeral-storage: 1Gi + volumeMounts: + - name: build-cache + mountPath: /cache + containers: + - name: runner + image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:2436278bf0756700c718927907cdbedac64d55c593f80b515acb44dec8b02cc1 + imagePullPolicy: IfNotPresent + command: ["/home/runner/run.sh"] + env: + - name: CARGO_HOME + value: /home/runner/.cargo + - name: HTTP_PROXY + valueFrom: + configMapKeyRef: + name: query-regression-runner-local-env + key: HTTP_PROXY + optional: true + - name: HTTPS_PROXY + valueFrom: + configMapKeyRef: + name: query-regression-runner-local-env + key: HTTPS_PROXY + optional: true + - name: NO_PROXY + valueFrom: + configMapKeyRef: + name: query-regression-runner-local-env + key: NO_PROXY + optional: true + - name: UV_CACHE_DIR + value: /home/runner/.cargo/uv-cache + - name: RUSTUP_HOME + value: /opt/rustup + - name: RUSTUP_TOOLCHAIN + value: nightly-2026-03-21 + - name: RUSTUP_AUTO_INSTALL + value: "0" + - name: CARGO_TARGET_DIR + value: /home/runner/query-regression-target + - name: QUERY_REGRESSION_CACHE_META + value: /home/runner/query-regression-cache-meta + - name: RUSTC_WRAPPER + value: /usr/local/bin/sccache + - name: SCCACHE_DIR + value: /home/runner/.cache/sccache + - name: SCCACHE_CACHE_SIZE + value: 40G + - name: CARGO_INCREMENTAL + value: "0" + securityContext: + runAsNonRoot: true + runAsUser: 1001 + runAsGroup: 1001 + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + resources: + requests: + cpu: "6" memory: 24Gi ephemeral-storage: 80Gi limits: cpu: "8" memory: 24Gi ephemeral-storage: 80Gi + volumeMounts: + - name: cargo-home + mountPath: /home/runner/.cargo + - name: build-cache + mountPath: /home/runner/.cargo/registry + subPath: cargo-registry-v1 + - name: build-cache + mountPath: /home/runner/.cargo/git + subPath: cargo-git-v1 + - name: build-cache + mountPath: /home/runner/query-regression-target + subPath: query-regression-target-v1 + - name: build-cache + mountPath: /home/runner/query-regression-cache-meta + subPath: meta-v1 + - name: build-cache + mountPath: /home/runner/.cache/sccache + subPath: sccache-v1 diff --git a/.github/runner-scale-sets/query-regression/values-paused.yaml b/.github/runner-scale-sets/query-regression/values-paused.yaml new file mode 100644 index 00000000000..fe734ef7615 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/values-paused.yaml @@ -0,0 +1,2 @@ +minRunners: 0 +maxRunners: 0 diff --git a/.github/scripts/query-regression-comment.cjs b/.github/scripts/query-regression-comment.cjs index 5a98e0b8161..6f4386fc3be 100644 --- a/.github/scripts/query-regression-comment.cjs +++ b/.github/scripts/query-regression-comment.cjs @@ -241,8 +241,10 @@ module.exports = async function validateQueryRegressionComment({ github, context '> Rendered by a trusted workflow from JSON artifacts produced by the query-regression run. Results from untrusted PR code are advisory until reviewed.', '', `- **Workflow run:** ${serverUrl}/${context.repo.owner}/${context.repo.repo}/actions/runs/${expectedRunId}`, - `- **Base SHA:** \`${text(metadata.base_sha)}\``, + `- **Built base SHA:** \`${text(metadata.built_base_sha)}\``, + `- **Event base SHA:** \`${text(metadata.event_base_sha)}\``, `- **Head SHA:** \`${text(metadata.head_sha)}\``, + `- **Candidate merge SHA:** \`${text(metadata.candidate_sha)}\``, '', ].join('\n'); diff --git a/.github/scripts/query-regression-pr-metadata.py b/.github/scripts/query-regression-pr-metadata.py index 2d8ed3f2e70..ba7ab113b6c 100644 --- a/.github/scripts/query-regression-pr-metadata.py +++ b/.github/scripts/query-regression-pr-metadata.py @@ -31,7 +31,10 @@ def main() -> int: metadata = { "pr_number": int(os.environ["PR_NUMBER"]), "head_sha": os.environ["HEAD_SHA"], - "base_sha": os.environ["BASE_SHA"], + "base_sha": os.environ["EVENT_BASE_SHA"], + "event_base_sha": os.environ["EVENT_BASE_SHA"], + "built_base_sha": os.environ["VERIFIED_BASE_SHA"], + "candidate_sha": os.environ["VERIFIED_CANDIDATE_SHA"], "head_repo": os.environ["HEAD_REPO"], "base_repo": os.environ["BASE_REPO"], "run_id": int(os.environ["RUN_ID"]), diff --git a/.github/scripts/query-regression-run.py b/.github/scripts/query-regression-run.py index b89857be81f..62d11f44203 100644 --- a/.github/scripts/query-regression-run.py +++ b/.github/scripts/query-regression-run.py @@ -26,6 +26,10 @@ from pathlib import Path DEFAULT_CASES = [ "tests/perf/query_cases/smoke_direct_sst/case.toml", + "tests/perf/query_cases/prom_remote_write_seeded_random/case.toml", + "tests/perf/query_cases/prom_remote_write_run_heavy/case.toml", + "tests/perf/query_cases/prom_remote_write_mixed_every/case.toml", + "tests/perf/query_cases/prom_remote_write_integer_counter/case.toml", "tests/perf/query_cases/promql_pushdown_7913/case.toml", "tests/perf/query_cases/sql_topk_order_by/case.toml", "tests/perf/query_cases/sql_aggregate_order_by/case.toml", @@ -106,8 +110,6 @@ def run_case(args: argparse.Namespace, case_path: Path, work_dir: Path) -> int: str(candidate_bin), "--fixture-generator", str(fixture_generator), - "--remote-write-generator", - str(args.remote_write_generator), "--work-dir", str(work_dir), "--http-timeout", @@ -160,11 +162,6 @@ def main() -> int: parser.add_argument("--work-dir", default=Path("query-regression-work"), type=Path) parser.add_argument("--http-timeout", default=os.environ.get("HTTP_TIMEOUT", "300")) parser.add_argument("--allow-large-fixture", default=os.environ.get("ALLOW_LARGE_FIXTURE", "false")) - parser.add_argument( - "--remote-write-generator", - type=Path, - default=configured_path(os.environ.get("REMOTE_WRITE_GENERATOR")), - ) parser.add_argument( "--summary-script", type=Path, @@ -177,9 +174,6 @@ def main() -> int: parser.add_argument("--candidate-ref", default=os.environ.get("CANDIDATE_REF", "")) parser.add_argument("--github-output", default=os.environ.get("GITHUB_OUTPUT")) args = parser.parse_args() - if args.remote_write_generator is None: - args.remote_write_generator = args.candidate_src / "target" / profile_dir(args.cargo_profile) / "prom_remote_write_fixture" - try: cases = split_cases(args.cases or [os.environ.get("CASE_PATHS", "all")]) except ValueError as err: diff --git a/.github/scripts/wait-kafka-ready.sh b/.github/scripts/wait-kafka-ready.sh new file mode 100755 index 00000000000..7617bdfa1f2 --- /dev/null +++ b/.github/scripts/wait-kafka-ready.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash + +set -euo pipefail + +GT_KAFKA_NAMESPACE="${GT_KAFKA_NAMESPACE:-kafka-cluster}" +GT_KAFKA_CLIENT_NAMESPACE="${GT_KAFKA_CLIENT_NAMESPACE:-my-greptimedb}" +GT_KAFKA_BOOTSTRAP_SERVERS="${GT_KAFKA_BOOTSTRAP_SERVERS:-kafka.kafka-cluster.svc.cluster.local:9092}" +GT_KAFKA_IMAGE="${GT_KAFKA_IMAGE:-docker.io/greptime/kafka:3.9.0-debian-12-r1}" +GT_KAFKA_READY_TIMEOUT_SECS="${GT_KAFKA_READY_TIMEOUT_SECS:-120}" +GT_KAFKA_READY_INTERVAL_SECS="${GT_KAFKA_READY_INTERVAL_SECS:-5}" +GT_KAFKA_POD_LABEL="${GT_KAFKA_POD_LABEL:-app.kubernetes.io/instance=kafka}" + +log() { + printf '[wait-kafka-ready] %s\n' "$*" +} + +run_client_probe() { + local pod_name="kafka-readiness-${RANDOM}-${RANDOM}" + + kubectl run "${pod_name}" \ + --namespace "${GT_KAFKA_CLIENT_NAMESPACE}" \ + --image "${GT_KAFKA_IMAGE}" \ + --restart=Never \ + --rm \ + --attach \ + --quiet \ + --env "GT_KAFKA_BOOTSTRAP_SERVERS=${GT_KAFKA_BOOTSTRAP_SERVERS}" \ + --command -- bash -euo pipefail -c ' + if command -v kafka-broker-api-versions.sh >/dev/null 2>&1; then + kafka_tool="kafka-broker-api-versions.sh" + elif [ -x /opt/bitnami/kafka/bin/kafka-broker-api-versions.sh ]; then + kafka_tool="/opt/bitnami/kafka/bin/kafka-broker-api-versions.sh" + else + echo "kafka-broker-api-versions.sh not found" >&2 + exit 1 + fi + + "${kafka_tool}" --bootstrap-server "${GT_KAFKA_BOOTSTRAP_SERVERS}" >/tmp/kafka-ready.out + cat /tmp/kafka-ready.out + ' +} + +log "wait for Kafka pods in namespace ${GT_KAFKA_NAMESPACE}" +kubectl wait \ + --for=condition=Ready \ + pod -l "${GT_KAFKA_POD_LABEL}" \ + --timeout="${GT_KAFKA_READY_TIMEOUT_SECS}s" \ + -n "${GT_KAFKA_NAMESPACE}" + +deadline=$((SECONDS + GT_KAFKA_READY_TIMEOUT_SECS)) +attempt=1 + +while [ "${SECONDS}" -le "${deadline}" ]; do + log "probe Kafka from namespace ${GT_KAFKA_CLIENT_NAMESPACE}, attempt ${attempt}" + if run_client_probe; then + log "Kafka is ready at ${GT_KAFKA_BOOTSTRAP_SERVERS}" + exit 0 + fi + + log "Kafka probe failed; retrying in ${GT_KAFKA_READY_INTERVAL_SECS}s" + sleep "${GT_KAFKA_READY_INTERVAL_SECS}" + attempt=$((attempt + 1)) +done + +log "timed out waiting for Kafka bootstrap ${GT_KAFKA_BOOTSTRAP_SERVERS}" +kubectl get pods -o wide -n "${GT_KAFKA_NAMESPACE}" || true +kubectl get pods -o wide -n "${GT_KAFKA_CLIENT_NAMESPACE}" || true +exit 1 diff --git a/.github/workflows/develop.yml b/.github/workflows/develop.yml index 08c44a826b2..0d9eceacb8e 100644 --- a/.github/workflows/develop.yml +++ b/.github/workflows/develop.yml @@ -457,6 +457,11 @@ jobs: name: Port forward Kafka WAL helper run: | kubectl port-forward service/kafka-wal-helper 8080:8080 -n kafka-cluster& + - if: matrix.mode.kafka + name: Check Kafka readiness + shell: bash + run: | + .github/scripts/wait-kafka-ready.sh - name: Fuzz Test uses: ./.github/actions/fuzz-test env: @@ -648,6 +653,11 @@ jobs: - name: Port forward (mysql) run: | kubectl port-forward service/my-greptimedb-frontend 4002:4002 -n my-greptimedb& + - if: matrix.mode.kafka + name: Check Kafka readiness + shell: bash + run: | + .github/scripts/wait-kafka-ready.sh - name: Fuzz Test uses: ./.github/actions/fuzz-test env: diff --git a/.github/workflows/nightly-jsonbench.yaml b/.github/workflows/nightly-jsonbench.yaml index da13a428e38..d5db6f9ba4a 100644 --- a/.github/workflows/nightly-jsonbench.yaml +++ b/.github/workflows/nightly-jsonbench.yaml @@ -10,14 +10,19 @@ on: description: The nightly build workflow run id to download GreptimeDB artifacts from required: true type: string + download_dataset: + description: Download the JSONBench dataset before running the benchmark + required: false + default: false + type: boolean permissions: actions: read contents: read concurrency: - group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} - cancel-in-progress: true + group: ${{ github.workflow }} + cancel-in-progress: false jobs: resolve-artifact: @@ -115,6 +120,34 @@ jobs: security-group-id: ${{ vars.EC2_RUNNER_SECURITY_GROUP_ID }} subnet-id: ${{ vars.EC2_RUNNER_SUBNET_ID }} + - name: Attach JSONBench data volume + if: ${{ steps.start-linux-runner.outputs.ec2-instance-id != '' && inputs.download_dataset != true }} + shell: bash + env: + INSTANCE_ID: ${{ steps.start-linux-runner.outputs.ec2-instance-id }} + VOLUME_ID: ${{ secrets.JSONBENCH_DATA_VOLUME_ID }} + run: | + set -euo pipefail + + if [[ -z "${VOLUME_ID}" ]]; then + echo "JSONBENCH_DATA_VOLUME_ID is not configured" + exit 1 + fi + + echo "Waiting for JSONBench data volume to become available" + aws ec2 wait volume-available --volume-ids "${VOLUME_ID}" + + echo "Attaching JSONBench data volume to EC2 instance ${INSTANCE_ID} as /dev/sdf" + aws ec2 attach-volume \ + --instance-id "${INSTANCE_ID}" \ + --volume-id "${VOLUME_ID}" \ + --device /dev/sdf \ + > /dev/null + + echo "Waiting for JSONBench data volume attachment to complete" + aws ec2 wait volume-in-use --volume-ids "${VOLUME_ID}" + echo "JSONBench data volume is attached to EC2 instance ${INSTANCE_ID}" + jsonbench: name: Run JSONBench if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success') }} @@ -123,7 +156,37 @@ jobs: timeout-minutes: 120 env: JSONBENCH_OUTPUT_PREFIX: _linux-${{ needs.resolve-artifact.outputs.artifact-arch }} + JSONBENCH_DATASET: 3 + JSONBENCH_DATA_DIR: /root/data/bluesky steps: + - name: Mount JSONBench data volume + if: ${{ needs.allocate-runner.outputs.linux-ec2-runner-instance-id != '' && inputs.download_dataset != true }} + shell: bash + env: + VOLUME_ID: ${{ secrets.JSONBENCH_DATA_VOLUME_ID }} + run: | + set -euo pipefail + + volume_serial="${VOLUME_ID//-/}" + device="" + echo "Waiting for the JSONBench data volume block device" + for _ in {1..30}; do + device=$(lsblk -dpno NAME,SERIAL | awk -v serial="${volume_serial}" '$2 == serial { print $1; exit }') + [[ -n "${device}" ]] && break + sleep 2 + done + + if [[ -z "${device}" ]]; then + echo "Cannot find the JSONBench data volume block device" + exit 1 + fi + + echo "Resolved JSONBench data volume to ${device}" + mkdir -p "${JSONBENCH_DATA_DIR}" + echo "Mounting ${device} read-only at ${JSONBENCH_DATA_DIR}" + mount -o ro "${device}" "${JSONBENCH_DATA_DIR}" + echo "JSONBench data volume is mounted at ${JSONBENCH_DATA_DIR}" + - name: Download GreptimeDB nightly artifact uses: actions/download-artifact@v4 with: @@ -142,26 +205,33 @@ jobs: chmod +x ./greptime rm -rf greptimedb-artifact "${{ needs.resolve-artifact.outputs.artifact-name }}" - - name: Run JSONBench - env: - # TODO(LFC): Change to "3" (100m) when JSON2 ingestion performance is optimized. - JSONBENCH_DATASET: 2 + - name: Clone JSONBench shell: bash run: | set -euo pipefail - export JSONBENCH_DATA_DIR="/root/data/bluesky" - echo "Use JSONBench data directory ${JSONBENCH_DATA_DIR}" - echo "${JSONBENCH_DATASET}" > jsonbench.dataset - - echo "Cloning JSONBench" git clone --branch greptimedb-new-json --depth 1 https://github.com/GreptimeTeam/JSONBench.git JSONBench + - name: Download JSONBench dataset + if: ${{ inputs.download_dataset == true }} + shell: bash + run: | + set -euo pipefail + echo "Downloading JSONBench dataset choice ${JSONBENCH_DATASET} to ${JSONBENCH_DATA_DIR}" mkdir -p "${JSONBENCH_DATA_DIR}" printf "${JSONBENCH_DATASET}\n" | ./JSONBench/download_data.sh + + - name: Run JSONBench + shell: bash + run: | + set -euo pipefail + + echo "Use JSONBench data directory ${JSONBENCH_DATA_DIR}" + echo "${JSONBENCH_DATASET}" > jsonbench.dataset + downloaded_files=$(find "${JSONBENCH_DATA_DIR}" -type f | wc -l) - echo "Downloaded JSONBench dataset files: ${downloaded_files}" + echo "Available JSONBench dataset files: ${downloaded_files}" export GREPTIMEDB_STANDALONE__WAL__DIR=greptimedb_data/wal export GREPTIMEDB_STANDALONE__STORAGE__DATA_HOME=greptimedb_data diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index a898be6ad29..3485bb32434 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -45,10 +45,6 @@ on: permissions: contents: read -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.run_id }} - cancel-in-progress: true - jobs: query-regression: if: >- @@ -59,12 +55,30 @@ jobs: (github.event.action != 'labeled' || github.event.label.name == 'query-regression')) }} runs-on: ${{ github.event_name == 'workflow_dispatch' && inputs.runner || 'perf-regression-8-cores' }} timeout-minutes: 180 + concurrency: + group: query-regression-persistent-cache-v1 + queue: max + cancel-in-progress: false env: CARGO_PROFILE: ${{ github.event_name == 'pull_request' && 'nightly' || inputs.cargo_profile }} - CARGO_TARGET_DIR: ${{ github.workspace }}/query-regression-target - CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" + CARGO_HOME: /home/runner/.cargo + UV_CACHE_DIR: /home/runner/.cargo/uv-cache + RUSTUP_HOME: /opt/rustup + RUSTUP_TOOLCHAIN: nightly-2026-03-21 + RUSTUP_AUTO_INSTALL: "0" + CARGO_TARGET_DIR: /home/runner/query-regression-target + QUERY_REGRESSION_CACHE_META: /home/runner/query-regression-cache-meta + RUSTC_WRAPPER: /usr/local/bin/sccache + SCCACHE_DIR: /home/runner/.cache/sccache + SCCACHE_CACHE_SIZE: 40G + CARGO_INCREMENTAL: "0" + RUSTFLAGS: -D warnings -C link-arg=-fuse-ld=mold + QUERY_REGRESSION_CACHE_EPOCH: "1" BASE_REF: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || inputs.base_ref }} - CANDIDATE_REF: ${{ github.event_name == 'pull_request' && format('refs/pull/{0}/merge', github.event.pull_request.number) || inputs.candidate_ref || github.ref }} + CANDIDATE_REF: ${{ github.event_name == 'workflow_dispatch' && (inputs.candidate_ref || github.ref) || '' }} + EVENT_MERGE_SHA: ${{ github.event_name == 'pull_request' && github.sha || '' }} + EVENT_HEAD_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || '' }} + EVENT_BASE_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || '' }} steps: - name: Checkout base source uses: actions/checkout@v4 @@ -82,28 +96,382 @@ jobs: cp src/.github/scripts/query-regression-summary.py \ query-regression-trusted-scripts/query-regression-summary.py - - uses: arduino/setup-protoc@v3 - with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + - name: Resolve immutable candidate and build base + working-directory: src + shell: bash + run: | + set -euo pipefail - - uses: actions-rust-lang/setup-rust-toolchain@v1 - with: - cache: false + fail_closed() { + printf '%s\n' \ + "PR changed or merge result unavailable; maintainer must remove and re-add query-regression after reviewing current revision. $1" >&2 + exit 1 + } - - uses: astral-sh/setup-uv@v5 + is_full_sha() { + [[ "$1" =~ ^[[:xdigit:]]{40}$ ]] + } - - name: Rust cache - uses: Swatinem/rust-cache@v2 - with: - workspaces: src - shared-key: query-regression - cache-all-crates: "true" - save-if: false + fetch_with_retry() { + local ref="$1" + local depth="$2" + local attempt=1 + local max_attempts=5 + local delay + + while (( attempt <= max_attempts )); do + if git fetch --no-tags --prune --depth="${depth}" origin "${ref}"; then + return 0 + fi + printf 'Git fetch failed for %s (attempt %d/%d).\n' \ + "${ref}" "${attempt}" "${max_attempts}" >&2 + if (( attempt < max_attempts )); then + delay=$(( 2 ** (attempt - 1) )) + printf 'Retrying git fetch for %s in %d seconds.\n' "${ref}" "${delay}" >&2 + sleep "${delay}" + fi + ((attempt += 1)) + done + return 1 + } + + initially_checked_out_base_sha="$(git rev-parse --verify HEAD)" + is_full_sha "${initially_checked_out_base_sha}" || { + printf 'Could not resolve initially checked-out BASE_REF commit.\n' >&2 + exit 1 + } + + if [[ "${GITHUB_EVENT_NAME}" == "pull_request" ]]; then + for identity in EVENT_MERGE_SHA EVENT_HEAD_SHA EVENT_BASE_SHA; do + value="${!identity-}" + is_full_sha "${value}" || fail_closed "Missing or invalid ${identity}." + done + + if ! fetch_with_retry "${EVENT_MERGE_SHA}" 2; then + fail_closed "Could not fetch event merge SHA ${EVENT_MERGE_SHA}." + fi + fetched_sha="$(git rev-parse --verify FETCH_HEAD 2>/dev/null)" \ + || fail_closed "Fetched event merge SHA could not be resolved." + [[ "${fetched_sha,,}" == "${EVENT_MERGE_SHA,,}" ]] \ + || fail_closed "FETCH_HEAD ${fetched_sha} does not match event merge SHA ${EVENT_MERGE_SHA}." + + read -r -a parents <<< "$(git show -s --format=%P FETCH_HEAD)" + (( ${#parents[@]} == 2 )) \ + || fail_closed "Event merge SHA ${EVENT_MERGE_SHA} is not a two-parent merge commit." + head_parent_count=0 + for parent in "${parents[@]}"; do + if [[ "${parent,,}" == "${EVENT_HEAD_SHA,,}" ]]; then + ((head_parent_count += 1)) + else + VERIFIED_BASE_SHA="${parent,,}" + fi + done + (( head_parent_count == 1 )) \ + || fail_closed "Event merge SHA must have exactly one EVENT_HEAD_SHA parent." + is_full_sha "${VERIFIED_BASE_SHA-}" \ + || fail_closed "Event merge SHA has no non-head parent for the base build." + VERIFIED_CANDIDATE_SHA="${fetched_sha,,}" + else + if ! fetch_with_retry "${CANDIDATE_REF}" 1; then + printf 'Could not fetch candidate ref %s.\n' "${CANDIDATE_REF}" >&2 + exit 1 + fi + VERIFIED_CANDIDATE_SHA="$(git rev-parse --verify FETCH_HEAD)" + VERIFIED_BASE_SHA="${initially_checked_out_base_sha}" + fi + + is_full_sha "${VERIFIED_CANDIDATE_SHA}" || { + printf 'Could not resolve candidate SHA.\n' >&2 + exit 1 + } + is_full_sha "${VERIFIED_BASE_SHA}" || { + printf 'Could not resolve build base SHA.\n' >&2 + exit 1 + } + printf 'VERIFIED_CANDIDATE_SHA=%s\n' "${VERIFIED_CANDIDATE_SHA,,}" >> "${GITHUB_ENV}" + printf 'VERIFIED_BASE_SHA=%s\n' "${VERIFIED_BASE_SHA,,}" >> "${GITHUB_ENV}" + printf 'Verified candidate SHA: %s\n' "${VERIFIED_CANDIDATE_SHA,,}" + printf 'Verified build base SHA: %s\n' "${VERIFIED_BASE_SHA,,}" + printf -- '- Verified candidate SHA: `%s`; built base SHA: `%s`\n' \ + "${VERIFIED_CANDIDATE_SHA,,}" "${VERIFIED_BASE_SHA,,}" >> "${GITHUB_STEP_SUMMARY}" + git reset --hard "${VERIFIED_BASE_SHA}" + git clean -ffdx + + - name: Verify runner image tools + shell: bash + run: | + set -euo pipefail + [[ "$(id -u)" == "1001" ]] + [[ "$(id -g)" == "1001" ]] + [[ "${UV_CACHE_DIR}" == "/home/runner/.cargo/uv-cache" ]] + [[ "$(protoc --version)" == "libprotoc 3.21.12" ]] + temporary_proto_dir="$(mktemp --directory)" + trap 'rm -rf "${temporary_proto_dir}"' EXIT + test -r /usr/include/google/protobuf/any.proto + test -r /usr/include/google/protobuf/empty.proto + printf '%s\n' \ + 'syntax = "proto3";' \ + 'package smoke;' \ + 'import "google/protobuf/any.proto";' \ + 'import "google/protobuf/empty.proto";' \ + 'message Smoke { google.protobuf.Any any = 1; google.protobuf.Empty empty = 2; }' \ + > "${temporary_proto_dir}/smoke.proto" + protoc --proto_path="${temporary_proto_dir}" --proto_path=/usr/include \ + --descriptor_set_out="${temporary_proto_dir}/smoke.pb" \ + "${temporary_proto_dir}/smoke.proto" + test -s "${temporary_proto_dir}/smoke.pb" + [[ "$(uv --version)" =~ ^uv[[:space:]]0\.11\.26([[:space:]]|$) ]] + mold_version="$(mold --version)" + [[ "${mold_version}" =~ ^mold[[:space:]]2\.30\.0([[:space:]]|$) ]] + [[ "$(python3 --version)" == "Python 3.12.3" ]] + sccache_version="$(sccache --version)" + [[ "${sccache_version}" =~ ^sccache[[:space:]]0\.16\.0([[:space:]]|$) ]] + [[ "$(command -v rustup)" == "/opt/cargo/bin/rustup" ]] + [[ "$(command -v cargo)" == "/opt/cargo/bin/cargo" ]] + [[ "$(command -v rustc)" == "/opt/cargo/bin/rustc" ]] + [[ "$(rustup --version)" =~ ^rustup[[:space:]]1\.29\.0([[:space:]]|$) ]] + cargo_version="$(cargo --version)" + [[ "${cargo_version}" =~ ^cargo[[:space:]]1\.96\.0-nightly[[:space:]]\(cbb9bb8bd[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$ ]] + rustc_version="$(rustc --version)" + [[ "${rustc_version}" =~ ^rustc[[:space:]]1\.96\.0-nightly[[:space:]]\(ac7f9ec7d[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$ ]] + active_toolchain="$(rustup show active-toolchain)" + [[ "${active_toolchain}" =~ ^nightly-2026-03-21-x86_64-unknown-linux-gnu([[:space:]]|$) ]] + [[ "${RUSTUP_HOME}" == "/opt/rustup" ]] + [[ "${RUSTUP_TOOLCHAIN}" == "nightly-2026-03-21" ]] + [[ "${RUSTUP_AUTO_INSTALL}" == "0" ]] + test -r /opt/rustup && test -x /opt/rustup + test ! -w /opt/rustup + test ! -w /opt/cargo/bin + for entry in config config.toml credentials credentials.toml bin .crates.toml .crates2.json .global-cache .package-cache; do + test ! -e "${CARGO_HOME}/${entry}" + done + mkdir -p "${CARGO_HOME}/registry" "${CARGO_HOME}/git" + + - name: Prepare persistent query regression cache + shell: bash + working-directory: src + run: | + set -euo pipefail + + readonly EXPECTED_CARGO_HOME="/home/runner/.cargo" + readonly EXPECTED_CARGO_REGISTRY="/home/runner/.cargo/registry" + readonly EXPECTED_CARGO_GIT="/home/runner/.cargo/git" + readonly EXPECTED_UV_CACHE_DIR="/home/runner/.cargo/uv-cache" + readonly EXPECTED_RUSTUP_HOME="/opt/rustup" + readonly EXPECTED_TARGET_DIR="/home/runner/query-regression-target" + readonly EXPECTED_CACHE_META="/home/runner/query-regression-cache-meta" + readonly EXPECTED_SCCACHE_DIR="/home/runner/.cache/sccache" + readonly EXPECTED_RUSTC_WRAPPER="/usr/local/bin/sccache" + readonly RUNNER_IMAGE_DIGEST="sha256:2436278bf0756700c718927907cdbedac64d55c593f80b515acb44dec8b02cc1" + readonly RUNNER_IMAGE_EPOCH="4" + + require_expected_root() { + local name="$1" + local actual="$2" + local expected="$3" + [[ "${actual}" == /* && "${actual}" == "${expected}" ]] || { + printf 'Refusing unexpected %s root: %s\n' "${name}" "${actual}" >&2 + exit 1 + } + } + + clear_directory() { + local root="$1" + [[ "${root}" == /* && "${root}" != "/" ]] || { + printf 'Refusing to clear non-absolute or root path: %s\n' "${root}" >&2 + exit 1 + } + case "${root}" in + "${EXPECTED_TARGET_DIR}"|"${EXPECTED_TARGET_DIR}/debug"|"${EXPECTED_TARGET_DIR}/release"|"${EXPECTED_TARGET_DIR}/nightly") ;; + *) + printf 'Refusing to clear unexpected cache path: %s\n' "${root}" >&2 + exit 1 + ;; + esac + find "${root}" -mindepth 1 -maxdepth 1 -exec rm -rf -- {} + + } + + clear_cargo_extracted_trees() { + require_expected_root CARGO_REGISTRY "${CARGO_HOME}/registry" "${EXPECTED_CARGO_REGISTRY}" + require_expected_root CARGO_GIT "${CARGO_HOME}/git" "${EXPECTED_CARGO_GIT}" + rm -rf -- "${CARGO_HOME}/registry/src" "${CARGO_HOME}/git/checkouts" + } + + cargo_size_kib() { + du -sk -- "${CARGO_HOME}/registry" "${CARGO_HOME}/git" | awk '{ total += $1 } END { print total }' + } + + target_size_kib() { + du -sk -- "${CARGO_TARGET_DIR}" | cut -f1 + } + + free_kib() { + df -Pk "${CARGO_TARGET_DIR}" | awk 'NR == 2 { print $4 }' + } + + report_cache_usage() { + du -sh -- "${CARGO_HOME}" "${CARGO_HOME}/registry" "${CARGO_HOME}/git" \ + "${RUSTUP_HOME}" "${CARGO_TARGET_DIR}" "${QUERY_REGRESSION_CACHE_META}" "${SCCACHE_DIR}" + df -P "${CARGO_TARGET_DIR}" + df -hP "${CARGO_TARGET_DIR}" + df -Pi "${CARGO_TARGET_DIR}" + } + + require_expected_root CARGO_HOME "${CARGO_HOME}" "${EXPECTED_CARGO_HOME}" + require_expected_root CARGO_REGISTRY "${CARGO_HOME}/registry" "${EXPECTED_CARGO_REGISTRY}" + require_expected_root CARGO_GIT "${CARGO_HOME}/git" "${EXPECTED_CARGO_GIT}" + require_expected_root UV_CACHE_DIR "${UV_CACHE_DIR}" "${EXPECTED_UV_CACHE_DIR}" + require_expected_root RUSTUP_HOME "${RUSTUP_HOME}" "${EXPECTED_RUSTUP_HOME}" + require_expected_root CARGO_TARGET_DIR "${CARGO_TARGET_DIR}" "${EXPECTED_TARGET_DIR}" + require_expected_root QUERY_REGRESSION_CACHE_META "${QUERY_REGRESSION_CACHE_META}" "${EXPECTED_CACHE_META}" + require_expected_root SCCACHE_DIR "${SCCACHE_DIR}" "${EXPECTED_SCCACHE_DIR}" + [[ "${RUSTC_WRAPPER}" == "${EXPECTED_RUSTC_WRAPPER}" ]] || { + printf 'Refusing unexpected RUSTC_WRAPPER: %s\n' "${RUSTC_WRAPPER}" >&2 + exit 1 + } + [[ "${SCCACHE_CACHE_SIZE}" == "40G" ]] || { + printf 'Refusing unexpected SCCACHE_CACHE_SIZE: %s\n' "${SCCACHE_CACHE_SIZE}" >&2 + exit 1 + } + [[ "${CARGO_INCREMENTAL}" == "0" ]] || { + printf 'Refusing unexpected CARGO_INCREMENTAL: %s\n' "${CARGO_INCREMENTAL}" >&2 + exit 1 + } + test -r "${RUSTUP_HOME}" + test -x "${RUSTUP_HOME}" + test ! -w "${RUSTUP_HOME}" + mkdir -p "${CARGO_HOME}" "${CARGO_HOME}/registry" "${CARGO_HOME}/git" "${UV_CACHE_DIR}" \ + "${CARGO_TARGET_DIR}" "${QUERY_REGRESSION_CACHE_META}" "${SCCACHE_DIR}" + for root in "${CARGO_HOME}" "${CARGO_HOME}/registry" "${CARGO_HOME}/git" \ + "${UV_CACHE_DIR}" "${CARGO_TARGET_DIR}" "${QUERY_REGRESSION_CACHE_META}" "${SCCACHE_DIR}"; do + test -w "${root}" + touch "${root}/.query-regression-write-test" + rm -f "${root}/.query-regression-write-test" + done + + case "${CARGO_PROFILE}" in + dev) target_profile="debug" ;; + release|nightly) target_profile="${CARGO_PROFILE}" ;; + *) printf 'Unsupported CARGO_PROFILE: %s\n' "${CARGO_PROFILE}" >&2; exit 1 ;; + esac + + rustc_vv="$(rustc -vV)" + if cargo -Vv >/dev/null 2>&1; then + cargo_vv="$(cargo -Vv)" + else + cargo_vv="$(cargo -V)" + fi + host_triple="$(printf '%s\n' "${rustc_vv}" | grep '^host: ' | cut -d' ' -f2-)" + cc_version="$(cc --version)" + mold_version="$(mold --version | grep -m1 .)" + protoc_version="$(protoc --version)" + sccache_version="$(sccache --version)" + rustup_version="$(rustup --version)" + active_toolchain="$(rustup show active-toolchain)" + rustup_path="$(command -v rustup)" + cargo_path="$(command -v cargo)" + rustc_path="$(command -v rustc)" + profile_target_dir="${CARGO_TARGET_DIR}/${target_profile}" + marker="${QUERY_REGRESSION_CACHE_META}/target-${target_profile}.abi" + new_marker="$(mktemp "${QUERY_REGRESSION_CACHE_META}/.${target_profile}.abi.XXXXXX")" + { + printf 'cache_epoch=%s\n' "${QUERY_REGRESSION_CACHE_EPOCH}" + printf 'runner_image_digest=%s\n' "${RUNNER_IMAGE_DIGEST}" + printf 'runner_image_epoch=%s\n' "${RUNNER_IMAGE_EPOCH}" + printf 'host_triple=%s\n' "${host_triple}" + printf 'architecture=%s\n' "$(uname -m)" + printf 'cargo_profile=%s\n' "${CARGO_PROFILE}" + printf 'target_profile=%s\n' "${target_profile}" + printf 'RUSTFLAGS=%s\n' "${RUSTFLAGS-}" + printf 'CARGO_ENCODED_RUSTFLAGS=%s\n' "${CARGO_ENCODED_RUSTFLAGS-}" + printf 'CARGO_INCREMENTAL=%s\n' "${CARGO_INCREMENTAL-}" + printf 'RUSTC_WRAPPER=%s\n' "${RUSTC_WRAPPER-}" + printf 'SCCACHE_DIR=%s\n' "${SCCACHE_DIR-}" + printf 'SCCACHE_CACHE_SIZE=%s\n' "${SCCACHE_CACHE_SIZE-}" + printf 'RUSTUP_HOME=%s\n' "${RUSTUP_HOME-}" + printf 'RUSTUP_TOOLCHAIN=%s\n' "${RUSTUP_TOOLCHAIN-}" + printf 'RUSTUP_AUTO_INSTALL=%s\n' "${RUSTUP_AUTO_INSTALL-}" + printf 'PATH=%s\n' "${PATH-}" + printf '%s\n' "rustc_vV=${rustc_vv}" + printf '%s\n' "cargo_vV=${cargo_vv}" + printf '%s\n' "cc_version=${cc_version}" + printf '%s\n' "mold_version=${mold_version}" + printf '%s\n' "protoc_version=${protoc_version}" + printf '%s\n' "sccache_version=${sccache_version}" + printf '%s\n' "rustup_version=${rustup_version}" + printf '%s\n' "active_toolchain=${active_toolchain}" + printf '%s\n' "rustup_path=${rustup_path}" + printf '%s\n' "cargo_path=${cargo_path}" + printf '%s\n' "rustc_path=${rustc_path}" + } > "${new_marker}" + + if ! cmp -s "${marker}" "${new_marker}"; then + printf 'Cache ABI marker mismatch; invalidating target profile: %s\n' "${target_profile}" + mkdir -p "${profile_target_dir}" + clear_directory "${profile_target_dir}" + else + printf 'Cache ABI marker match; reusing target profile: %s\n' "${target_profile}" + fi + mv -f "${new_marker}" "${marker}" + + report_cache_usage + + target_size="$(target_size_kib)" + if (( target_size >= 400 * 1024 * 1024 )); then + printf 'Warning: target cache is at least 400 GiB (%s KiB)\n' "${target_size}" >&2 + fi + if (( target_size >= 450 * 1024 * 1024 )); then + printf 'Target cache is at least 450 GiB; clearing complete target root\n' >&2 + clear_directory "${CARGO_TARGET_DIR}" + fi + + cargo_size="$(cargo_size_kib)" + if (( cargo_size >= 60 * 1024 * 1024 )); then + printf 'Warning: Cargo cache is at least 60 GiB (%s KiB)\n' "${cargo_size}" >&2 + fi + if (( cargo_size >= 80 * 1024 * 1024 )); then + printf 'Cargo cache is at least 80 GiB; removing extracted sources and checkouts\n' >&2 + clear_cargo_extracted_trees + cargo_size="$(cargo_size_kib)" + if (( cargo_size >= 80 * 1024 * 1024 )); then + printf 'Cargo cache remains at least 80 GiB after cleanup (%s KiB)\n' "${cargo_size}" >&2 + exit 1 + fi + fi + + free_space="$(free_kib)" + if (( free_space < 300 * 1024 * 1024 )); then + printf 'Free space is below 300 GiB; clearing complete target root\n' >&2 + clear_directory "${CARGO_TARGET_DIR}" + free_space="$(free_kib)" + if (( free_space < 300 * 1024 * 1024 )); then + printf 'Free space remains below 300 GiB; removing Cargo extracted sources and checkouts\n' >&2 + clear_cargo_extracted_trees + free_space="$(free_kib)" + if (( free_space < 300 * 1024 * 1024 )); then + printf 'Free space remains below 300 GiB after cleanup (%s KiB)\n' "${free_space}" >&2 + exit 1 + fi + fi + fi + + report_cache_usage + + sccache --start-server + sccache --zero-stats + sccache --show-stats - name: Build base greptime working-directory: src run: | + set -euo pipefail + SECONDS=0 cargo build --profile "${CARGO_PROFILE}" -p cmd --bin greptime + base_build_elapsed="${SECONDS}" + printf 'Base greptime cargo build elapsed: %s seconds\n' "${base_build_elapsed}" + printf -- '- Base greptime cargo build: %s seconds\n' "${base_build_elapsed}" >> "${GITHUB_STEP_SUMMARY}" + sccache --show-stats + sccache --zero-stats target_dir="${CARGO_PROFILE}" if [[ "${CARGO_PROFILE}" == "dev" ]]; then target_dir="debug" @@ -114,15 +482,35 @@ jobs: - name: Switch source to candidate working-directory: src + shell: bash run: | - git fetch --no-tags --prune --depth=1 origin "${CANDIDATE_REF}" - git reset --hard FETCH_HEAD + set -euo pipefail + [[ "${VERIFIED_CANDIDATE_SHA}" =~ ^[[:xdigit:]]{40}$ ]] || { + printf 'Verified candidate SHA is missing or invalid.\n' >&2 + exit 1 + } + git cat-file -e "${VERIFIED_CANDIDATE_SHA}^{commit}" || { + printf 'Verified candidate SHA is unavailable in the local repository.\n' >&2 + exit 1 + } + git reset --hard "${VERIFIED_CANDIDATE_SHA}" git clean -ffdx + - name: Test query regression runner lifecycle + working-directory: src + run: uv run --no-project python tests/perf/test_query_regression_runner_compaction_toctou.py + - name: Build candidate greptime and fixture generators working-directory: src run: | - cargo build --profile "${CARGO_PROFILE}" -p cmd --bin greptime --bin query_perf_fixture --bin prom_remote_write_fixture + set -euo pipefail + SECONDS=0 + cargo build --profile "${CARGO_PROFILE}" -p cmd --bin greptime + cargo build --profile "${CARGO_PROFILE}" -p cmd --bin query_perf_fixture --features dev-tools + candidate_build_elapsed="${SECONDS}" + printf 'Candidate greptime and fixture generator cargo builds elapsed: %s seconds\n' "${candidate_build_elapsed}" + printf -- '- Candidate greptime and fixture generator cargo builds: %s seconds\n' "${candidate_build_elapsed}" >> "${GITHUB_STEP_SUMMARY}" + sccache --show-stats target_dir="${CARGO_PROFILE}" if [[ "${CARGO_PROFILE}" == "dev" ]]; then target_dir="debug" @@ -132,8 +520,6 @@ jobs: "${GITHUB_WORKSPACE}/query-regression-bins/candidate/greptime" cp "${CARGO_TARGET_DIR}/${target_dir}/query_perf_fixture" \ "${GITHUB_WORKSPACE}/query-regression-bins/candidate/query_perf_fixture" - cp "${CARGO_TARGET_DIR}/${target_dir}/prom_remote_write_fixture" \ - "${GITHUB_WORKSPACE}/query-regression-bins/candidate/prom_remote_write_fixture" - name: Run query regression id: run @@ -146,7 +532,6 @@ jobs: BASE_BIN: ${{ github.workspace }}/query-regression-bins/base/greptime CANDIDATE_BIN: ${{ github.workspace }}/query-regression-bins/candidate/greptime FIXTURE_GENERATOR: ${{ github.workspace }}/query-regression-bins/candidate/query_perf_fixture - REMOTE_WRITE_GENERATOR: ${{ github.workspace }}/query-regression-bins/candidate/prom_remote_write_fixture run: >- uv run --no-project python src/.github/scripts/query-regression-run.py --base-src src @@ -171,9 +556,13 @@ jobs: with: name: query-regression-report path: | - query-regression-work/** + query-regression-work/**/query-regression-report.json + query-regression-work/**/base/report.json + query-regression-work/**/candidate/report.json + query-regression-work/**/logs/** query-regression-summary.md if-no-files-found: warn + retention-days: 7 - name: Upload trusted comment artifact if: ${{ always() && github.event_name == 'pull_request' }} @@ -186,6 +575,46 @@ jobs: if-no-files-found: warn retention-days: 7 + - name: Report persistent cache usage + if: ${{ always() }} + shell: bash + run: | + set -euo pipefail + for root in \ + /home/runner/.cargo \ + /home/runner/.cargo/registry \ + /home/runner/.cargo/git \ + /opt/rustup \ + /home/runner/query-regression-target \ + /home/runner/query-regression-cache-meta \ + /home/runner/.cache/sccache; do + if [[ -e "${root}" ]]; then + du -sh -- "${root}" + else + printf 'Missing cache path (report only): %s\n' "${root}" + fi + done + if [[ -e /opt/rustup ]]; then + if [[ -w /opt/rustup ]]; then + printf 'Immutable Rustup path is unexpectedly writable: /opt/rustup\n' >&2 + else + printf 'Immutable Rustup path is read-only to runner: /opt/rustup\n' + fi + fi + report_root=/home/runner/query-regression-target + if [[ ! -e "${report_root}" ]]; then + report_root=/home/runner + printf 'Target cache path is missing; reporting its filesystem through %s\n' "${report_root}" + fi + df -P "${report_root}" + df -hP "${report_root}" + df -Pi "${report_root}" + if command -v sccache >/dev/null 2>&1; then + sccache --show-stats || printf 'Unable to show sccache statistics (report only)\n' >&2 + else + printf 'sccache is unavailable (report only)\n' >&2 + fi + - name: Fail on regression failure if: ${{ steps.run.outputs.status != '0' }} run: exit 1 diff --git a/AGENTS.md b/AGENTS.md index ce25569aff1..fd2ec44aba1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -77,8 +77,10 @@ change-coupling points, and gotchas: 2. `make clippy` 3. `make test` (or `cargo nextest run`) 4. `make check-udeps` (run `make fix-udeps` if it reports unused dependencies). -5. If you changed sample config under `config/`, run `make config-docs` (needs - Docker) and commit the regenerated `config/config.md`. +5. If you added or changed a public configuration option, update the applicable + example TOMLs, configuration-loading and serialized-config snapshot tests, + and related user-facing documentation. Run `make config-docs` (needs Docker) + and commit the regenerated `config/config.md`. 6. If you changed a persisted or wire format, add a compatibility test case (see `.agents/architecture-invariants.md`). 7. Use a conventional-commit title, sign off commits (`git commit -s`), and sign diff --git a/Cargo.lock b/Cargo.lock index e00e30e28a6..87827905fe6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -948,11 +948,14 @@ dependencies = [ "common-test-util", "digest 0.10.7", "hex", + "hmac 0.12.1", "pbkdf2", + "rand 0.9.4", "sha1 0.10.6", "sha2 0.10.9", "snafu 0.8.6", "sql", + "stringprep", "subtle", "tokio", ] @@ -5938,7 +5941,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=4a69c9e75240a76afeb0796798085813669288f8#4a69c9e75240a76afeb0796798085813669288f8" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=c6f92ee9e12ee7bd5ee3c0956b1c9caa39be3f76#c6f92ee9e12ee7bd5ee3c0956b1c9caa39be3f76" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", @@ -8516,6 +8519,7 @@ dependencies = [ "toml 0.8.23", "tracing", "uuid", + "zstd", ] [[package]] @@ -11020,6 +11024,7 @@ dependencies = [ "bytemuck", "common-error", "common-macro", + "common-query", "common-recordbatch", "common-telemetry", "criterion 0.7.0", @@ -11038,9 +11043,9 @@ dependencies = [ [[package]] name = "promql-parser" -version = "0.7.3" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f514ab521bb2b4ce26aca29bf133498b2a2651b785d03a262bbf5772329c279f" +checksum = "9e8c90c956be4237a93be6f2298302b56974ef862ab11e0f98869d1f27c69f88" dependencies = [ "cfgrammar", "chrono", @@ -11183,7 +11188,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.11.0", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -11196,7 +11201,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" dependencies = [ "anyhow", - "itertools 0.11.0", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", diff --git a/Cargo.toml b/Cargo.toml index 5ee60c450e8..12a1d2d346c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -158,7 +158,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 = "4a69c9e75240a76afeb0796798085813669288f8" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "c6f92ee9e12ee7bd5ee3c0956b1c9caa39be3f76" } hex = "0.4" http = "1" humantime = "2.1" @@ -198,7 +198,7 @@ paste = "1.0" pin-project = "1.0" pretty_assertions = "1.4.0" prometheus = { version = "0.14", features = ["process"] } -promql-parser = { version = "0.7.3", features = ["ser"] } +promql-parser = { version = "0.10.0", features = ["ser"] } prost = { version = "0.14", features = ["no-recursion-limit"] } prost-types = "0.14" raft-engine = { version = "0.4.1", default-features = false } @@ -242,6 +242,7 @@ sqlx = { version = "0.8", default-features = false, features = [ "tls-rustls-aws-lc-rs", "rust_decimal", ] } +stringprep = "0.1" strum = { version = "0.27", features = ["derive"] } sysinfo = "0.33" tempfile = "3" diff --git a/README.md b/README.md index 3ba52bc5266..99c78814609 100644 --- a/README.md +++ b/README.md @@ -209,7 +209,13 @@ Read the [v1.0 highlights](https://greptime.com/blogs/2025-11-05-greptimedb-v1-h If GreptimeDB is useful to you, please star the repo. -[![Star History Chart](https://api.star-history.com/svg?repos=GreptimeTeam/GreptimeDB&type=Date)](https://www.star-history.com/#GreptimeTeam/GreptimeDB&Date) + + + + + Star History Chart + + Known Users diff --git a/config/config.md b/config/config.md index 4ffd851967c..13aeccaead7 100644 --- a/config/config.md +++ b/config/config.md @@ -15,7 +15,7 @@ | `default_timezone` | String | Unset | The default timezone of the server. | | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index and value columns. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | -| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:"
"pbkdf2_sha256" protects passwords at rest but is not compatible with mysql_native_password. | +| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's
native password handshake; a MySQL client must send the password in cleartext for such users.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | | `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | | `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | | `init_regions_in_background` | Bool | `false` | Initialize all regions in the background during the startup.
By default, it provides services after all regions have been initialized. | @@ -157,7 +157,8 @@ | `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. | | `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. | | `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size`. | -| `region_engine.mito.sst_meta_cache_size` | String | Auto | Cache size for SST metadata. Setting it to 0 to disable the cache.
If not set, it's default to 1/32 of OS memory with a max limitation of 128MB. | +| `region_engine.mito.default_region_write_buffer_size` | String | `0` | Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. | +| `region_engine.mito.sst_meta_cache_size` | String | Auto | Cache size for SST metadata. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory with a max limitation of 512MB. | | `region_engine.mito.vector_cache_size` | String | Auto | Cache size for vectors and arrow arrays. Setting it to 0 to disable the cache.
If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. | | `region_engine.mito.page_cache_size` | String | Auto | Cache size for pages of SST row groups. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory. | | `region_engine.mito.selector_result_cache_size` | String | Auto | Cache size for time series selector (e.g. `last_value()`). Setting it to 0 to disable the cache.
If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. | @@ -205,7 +206,6 @@ | `region_engine.mito.bloom_filter_index.mem_threshold_on_create` | String | `auto` | Memory threshold for bloom filter creation.
- `auto`: automatically determine the threshold based on the system memory size (default)
- `unlimited`: no memory limit
- `[size]` e.g. `64MB`: fixed memory threshold | | `region_engine.file` | -- | -- | Enable the file engine. | | `region_engine.metric` | -- | -- | Metric engine options. | -| `region_engine.metric.sparse_primary_key_encoding` | Bool | `true` | Whether to use sparse primary key encoding. | | `region_engine.metric.experimental_enable_metric_value_split` | Bool | `false` | Whether to store exact integral metric values in internal Int64 columns. | | `logging` | -- | -- | The logging options. | | `logging.dir` | String | `./greptimedb_data/logs` | The directory to store the log files. If set to empty, logs will not be written to files. | @@ -240,7 +240,7 @@ | `default_timezone` | String | Unset | The default timezone of the server. | | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index and value columns. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | -| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:"
"pbkdf2_sha256" protects passwords at rest but is not compatible with mysql_native_password. | +| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's
native password handshake; a MySQL client must send the password in cleartext for such users.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | | `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | | `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | | `runtime` | -- | -- | The runtime options. | @@ -438,6 +438,9 @@ | `gc` | -- | -- | -- | | `gc.enable` | Bool | `false` | Whether GC is enabled. Default to false. Need to be the same with datanode's `mito.gc.enable`
If set to false, no GC will be performed | | `gc.gc_cooldown_period` | String | `5m` | Cooldown period between GC operations on the same region. | +| `gc.experimental_soft_drop` | -- | -- | -- | +| `gc.experimental_soft_drop.enable` | Bool | `false` | Reserved experimental option. Currently ignored. | +| `gc.experimental_soft_drop.retention` | String | `7d` | Reserved retention duration. Currently ignored. | | `logging` | -- | -- | The logging options. | | `logging.dir` | String | `./greptimedb_data/logs` | The directory to store the log files. If set to empty, logs will not be written to files. | | `logging.level` | String | Unset | The log level. Can be `info`/`debug`/`warn`/`error`. | @@ -515,7 +518,7 @@ | `wal.timeout` | String | `3s` | The timeout for kafka client.
**It's only used when the provider is `kafka`**. | | `wal.max_batch_bytes` | String | `1MB` | The max size of a single producer batch.
Warning: Kafka has a default limit of 1MB per message in a topic.
**It's only used when the provider is `kafka`**. | | `wal.consumer_wait_timeout` | String | `100ms` | The consumer wait timeout.
**It's only used when the provider is `kafka`**. | -| `wal.create_index` | Bool | `true` | Whether to enable WAL index creation.
**It's only used when the provider is `kafka`**. | +| `wal.create_index` | Bool | `false` | Whether to enable WAL index creation.
**It's only used when the provider is `kafka`**. | | `wal.dump_index_interval` | String | `60s` | The interval for dumping WAL indexes.
**It's only used when the provider is `kafka`**. | | `wal.overwrite_entry_start_id` | Bool | `false` | Ignore missing entries during read WAL.
**It's only used when the provider is `kafka`**.

This option ensures that when Kafka messages are deleted, the system
can still successfully replay memtable data without throwing an
out-of-range error.
However, enabling this option might lead to unexpected data loss,
as the system will skip over missing entries instead of treating
them as critical errors. | | `query` | -- | -- | The query engine options. | @@ -561,7 +564,8 @@ | `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. | | `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. | | `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size` | -| `region_engine.mito.sst_meta_cache_size` | String | Auto | Cache size for SST metadata. Setting it to 0 to disable the cache.
If not set, it's default to 1/32 of OS memory with a max limitation of 128MB. | +| `region_engine.mito.default_region_write_buffer_size` | String | `0` | Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. | +| `region_engine.mito.sst_meta_cache_size` | String | Auto | Cache size for SST metadata. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory with a max limitation of 512MB. | | `region_engine.mito.vector_cache_size` | String | Auto | Cache size for vectors and arrow arrays. Setting it to 0 to disable the cache.
If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. | | `region_engine.mito.page_cache_size` | String | Auto | Cache size for pages of SST row groups. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory. | | `region_engine.mito.selector_result_cache_size` | String | Auto | Cache size for time series selector (e.g. `last_value()`). Setting it to 0 to disable the cache.
If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. | @@ -613,7 +617,6 @@ | `region_engine.mito.gc.unknown_file_lingering_time` | String | `1d` | Lingering time before deleting unknown files (files with undetermined expel time).
Only applies during full file listing GC.
This uses the object's last-modified timestamp as a heuristic (strict less-than comparison);
do not configure this value too small in production to avoid deleting pre-manifest files
from in-progress compaction or flush.
For active/open regions, an unknown file is deleted only if its object last-modified time exceeds this TTL.
If the object store does not provide a last-modified timestamp, the file is conservatively kept.
For dropped regions, unknown files are deleted immediately. | | `region_engine.file` | -- | -- | Enable the file engine. | | `region_engine.metric` | -- | -- | Metric engine options. | -| `region_engine.metric.sparse_primary_key_encoding` | Bool | `true` | Whether to use sparse primary key encoding. | | `region_engine.metric.experimental_enable_metric_value_split` | Bool | `false` | Whether to store exact integral metric values in internal Int64 columns. | | `logging` | -- | -- | The logging options. | | `logging.dir` | String | `./greptimedb_data/logs` | The directory to store the log files. If set to empty, logs will not be written to files. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index f85505a0fa9..16e64b76ff8 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -185,7 +185,7 @@ consumer_wait_timeout = "100ms" ## Whether to enable WAL index creation. ## **It's only used when the provider is `kafka`**. -create_index = true +create_index = false ## The interval for dumping WAL indexes. ## **It's only used when the provider is `kafka`**. @@ -468,10 +468,13 @@ auto_flush_interval = "1h" ## @toml2docs:none-default="Auto" #+ global_write_buffer_reject_size = "2GB" +## Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. +#+ default_region_write_buffer_size = "0" + ## Cache size for SST metadata. Setting it to 0 to disable the cache. -## If not set, it's default to 1/32 of OS memory with a max limitation of 128MB. +## If not set, it's default to 1/8 of OS memory with a max limitation of 512MB. ## @toml2docs:none-default="Auto" -#+ sst_meta_cache_size = "128MB" +#+ sst_meta_cache_size = "512MB" ## Cache size for vectors and arrow arrays. Setting it to 0 to disable the cache. ## If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. @@ -697,8 +700,6 @@ unknown_file_lingering_time = "1d" [[region_engine]] ## Metric engine options. [region_engine.metric] -## Whether to use sparse primary key encoding. -sparse_primary_key_encoding = true ## Whether to store exact integral metric values in internal Int64 columns. experimental_enable_metric_value_split = false diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 9fc65ece8af..73efe8c1376 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -13,8 +13,16 @@ default_column_prefix = "greptime" ## The user provider for authentication. ## Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", -## "mysql_native_password:" -## "pbkdf2_sha256" protects passwords at rest but is not compatible with mysql_native_password. +## "mysql_native_password:", +## "pg_scram_sha256::::" +## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's +## native password handshake; a MySQL client must send the password in cleartext for such users. +## "mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all. +## PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or +## "mysql_native_password", PostgreSQL falls back to cleartext password auth for every user. +## For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are +## observable in the SCRAM server-first message, and non-default values weaken resistance to username +## enumeration. ## @toml2docs:none-default #+ user_provider = "static_user_provider:file:/path/to/users" diff --git a/config/metasrv.example.toml b/config/metasrv.example.toml index 28d831d4f2a..55e6497078b 100644 --- a/config/metasrv.example.toml +++ b/config/metasrv.example.toml @@ -326,6 +326,12 @@ enable = false ## Cooldown period between GC operations on the same region. gc_cooldown_period = "5m" +[gc.experimental_soft_drop] +## Reserved experimental option. Currently ignored. +enable = false +## Reserved retention duration. Currently ignored. +retention = "7d" + ## The logging options. [logging] ## The directory to store the log files. If set to empty, logs will not be written to files. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index d3f166e6968..01e0ea2a576 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -13,8 +13,16 @@ default_column_prefix = "greptime" ## The user provider for authentication. ## Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", -## "mysql_native_password:" -## "pbkdf2_sha256" protects passwords at rest but is not compatible with mysql_native_password. +## "mysql_native_password:", +## "pg_scram_sha256::::" +## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's +## native password handshake; a MySQL client must send the password in cleartext for such users. +## "mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all. +## PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or +## "mysql_native_password", PostgreSQL falls back to cleartext password auth for every user. +## For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are +## observable in the SCRAM server-first message, and non-default values weaken resistance to username +## enumeration. ## @toml2docs:none-default #+ user_provider = "static_user_provider:file:/path/to/users" @@ -597,10 +605,13 @@ auto_flush_interval = "1h" ## @toml2docs:none-default="Auto" #+ global_write_buffer_reject_size = "2GB" +## Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. +#+ default_region_write_buffer_size = "0" + ## Cache size for SST metadata. Setting it to 0 to disable the cache. -## If not set, it's default to 1/32 of OS memory with a max limitation of 128MB. +## If not set, it's default to 1/8 of OS memory with a max limitation of 512MB. ## @toml2docs:none-default="Auto" -#+ sst_meta_cache_size = "128MB" +#+ sst_meta_cache_size = "512MB" ## Cache size for vectors and arrow arrays. Setting it to 0 to disable the cache. ## If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. @@ -807,8 +818,6 @@ mem_threshold_on_create = "auto" [[region_engine]] ## Metric engine options. [region_engine.metric] -## Whether to use sparse primary key encoding. -sparse_primary_key_encoding = true ## Whether to store exact integral metric values in internal Int64 columns. experimental_enable_metric_value_split = false diff --git a/src/api/src/helper.rs b/src/api/src/helper.rs index ed8d64b57a3..1ff437f127a 100644 --- a/src/api/src/helper.rs +++ b/src/api/src/helper.rs @@ -20,7 +20,7 @@ use common_decimal::decimal128::{DECIMAL128_DEFAULT_SCALE, DECIMAL128_MAX_PRECIS use common_time::time::Time; use common_time::timestamp::TimeUnit; use common_time::{Date, IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth, Timestamp}; -use datatypes::json::value::{JsonNumber, JsonValue, JsonValueRef, JsonVariant}; +use datatypes::json::value::{JsonNumber, JsonValue, JsonValueRef, JsonVariant, JsonVariantRef}; use datatypes::prelude::{ConcreteDataType, ValueRef}; use datatypes::types::json_type::JsonNativeType; use datatypes::types::{ @@ -931,33 +931,59 @@ pub fn encode_json_value(value: JsonValue) -> v1::JsonValue { } fn decode_json_value(value: &v1::JsonValue) -> JsonValueRef<'_> { + let (variant, json_type) = decode_json_value_parts(value); + JsonValueRef::new_with_type(variant, json_type) +} + +fn decode_json_value_parts(value: &v1::JsonValue) -> (JsonVariantRef<'_>, JsonNativeType) { let Some(value) = &value.value else { - return JsonValueRef::null(); + return (JsonVariantRef::Null, JsonNativeType::Null); }; match value { - json_value::Value::Boolean(x) => (*x).into(), - json_value::Value::Int(x) => (*x).into(), - json_value::Value::Uint(x) => (*x).into(), - json_value::Value::Float(x) => (*x).into(), - json_value::Value::Str(x) => (x.as_str()).into(), - json_value::Value::Array(array) => array - .items - .iter() - .map(|x| decode_json_value(x).into_variant()) - .collect::>() - .into(), - json_value::Value::Object(x) => x - .entries - .iter() - .filter_map(|entry| { - entry - .value - .as_ref() - .map(|v| (entry.key.as_str(), decode_json_value(v).into_variant())) - }) - .collect::>() - .into(), - json_value::Value::Variant(x) => x.as_slice().into(), + json_value::Value::Boolean(x) => (JsonVariantRef::Bool(*x), JsonNativeType::Bool), + json_value::Value::Int(x) => ((*x).into(), JsonNativeType::i64()), + json_value::Value::Uint(x) => ((*x).into(), JsonNativeType::u64()), + json_value::Value::Float(x) => ((*x).into(), JsonNativeType::f64()), + json_value::Value::Str(x) => (x.as_str().into(), JsonNativeType::String), + json_value::Value::Array(array) => { + let mut variants = Vec::with_capacity(array.items.len()); + let mut item_type = JsonNativeType::Null; + for item in &array.items { + let (variant, ty) = decode_json_value_parts(item); + variants.push(variant); + if !matches!(item_type, JsonNativeType::Variant) { + item_type.merge(&ty); + } + } + ( + JsonVariantRef::Array(variants), + JsonNativeType::Array(Box::new(item_type)), + ) + } + json_value::Value::Object(object) => { + let mut variants = Vec::with_capacity(object.entries.len()); + let mut fields = Vec::with_capacity(object.entries.len()); + for entry in &object.entries { + let Some(value) = &entry.value else { + continue; + }; + let (variant, json_type) = decode_json_value_parts(value); + variants.push((entry.key.as_str(), variant)); + fields.push((entry.key.clone(), json_type)); + } + let variants = variants.into_iter().collect::>(); + let fields = fields.into_iter().collect::>(); + let json_type = if fields.is_empty() { + JsonNativeType::Null + } else { + JsonNativeType::Object(fields) + }; + (JsonVariantRef::Object(variants), json_type) + } + json_value::Value::Variant(x) => ( + JsonVariantRef::Variant(x.as_slice()), + JsonNativeType::Variant, + ), } } @@ -1854,15 +1880,40 @@ mod tests { ] })) ); + assert_eq!( + JsonNativeType::Array(Box::new(JsonNativeType::i64())), + decode_json_value_parts(&proto).1 + ); let value = decode_json_value(&proto); assert_eq!(json.as_ref(), value); + let proto = v1::JsonValue { + value: Some(json_value::Value::Array(JsonList { + items: vec![ + v1::JsonValue { + value: Some(json_value::Value::Int(1)), + }, + v1::JsonValue { + value: Some(json_value::Value::Float(2.0)), + }, + ], + })), + }; + assert_eq!( + JsonNativeType::Array(Box::new(JsonNativeType::f64())), + decode_json_value_parts(&proto).1 + ); + let json: JsonValue = [(); 0].into(); let proto = encode_json_value(json.clone()); assert_eq!( proto.value, Some(json_value::Value::Array(JsonList { items: vec![] })) ); + assert_eq!( + JsonNativeType::Array(Box::new(JsonNativeType::Null)), + decode_json_value_parts(&proto).1 + ); let value = decode_json_value(&proto); assert_eq!(json.as_ref(), value); @@ -1893,6 +1944,14 @@ mod tests { ] })) ); + assert_eq!( + JsonNativeType::Object(JsonObjectType::from([ + ("k1".to_string(), JsonNativeType::i64()), + ("k2".to_string(), JsonNativeType::i64()), + ("k3".to_string(), JsonNativeType::i64()), + ])), + decode_json_value_parts(&proto).1 + ); let value = decode_json_value(&proto); assert_eq!(json.as_ref(), value); @@ -1902,6 +1961,7 @@ mod tests { proto.value, Some(json_value::Value::Object(JsonObject { entries: vec![] })) ); + assert_eq!(JsonNativeType::Null, decode_json_value_parts(&proto).1); let value = decode_json_value(&proto); assert_eq!(json.as_ref(), value); diff --git a/src/auth/Cargo.toml b/src/auth/Cargo.toml index 77a676ea25b..959eb25078e 100644 --- a/src/auth/Cargo.toml +++ b/src/auth/Cargo.toml @@ -21,11 +21,14 @@ common-macro.workspace = true common-telemetry.workspace = true digest = "0.10" hex.workspace = true +hmac = "0.12" pbkdf2 = "0.12" +rand.workspace = true sha1 = "0.10" sha2 = "0.10" snafu.workspace = true sql.workspace = true +stringprep.workspace = true subtle = "2.6" tokio.workspace = true diff --git a/src/auth/src/common.rs b/src/auth/src/common.rs index 3122822a520..630995b2afc 100644 --- a/src/auth/src/common.rs +++ b/src/auth/src/common.rs @@ -12,14 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use common_base::secrets::SecretString; use digest::Digest; +use hmac::{Hmac, Mac}; use pbkdf2::pbkdf2_hmac; use sha1::Sha1; use sha2::Sha256; use snafu::{OptionExt, ensure}; +use subtle::ConstantTimeEq; use crate::error::{IllegalParamSnafu, InvalidConfigSnafu, Result, UserPasswordMismatchSnafu}; use crate::user_info::DefaultUserInfo; @@ -31,9 +33,17 @@ use crate::{UserInfoRef, UserProviderRef}; pub(crate) const DEFAULT_USERNAME: &str = "greptime"; pub const DEFAULT_PBKDF2_SHA256_ITERATIONS: u32 = 4096; +pub const DEFAULT_PBKDF2_SHA256_SALT_LEN: usize = 16; pub const PBKDF2_SHA256_HASH_LEN: usize = 32; pub const MAX_PBKDF2_SHA256_ITERATIONS: u32 = 1_000_000; pub const MAX_PBKDF2_SHA256_SALT_LEN: usize = 1024; +pub const PG_SCRAM_SHA256_KEY_LEN: usize = 32; + +type HmacSha256 = Hmac; + +/// Process-wide secret used to derive a stable mock salt for unknown users, so +/// the SCRAM `server-first-message` can't be used to enumerate usernames. +static PG_SCRAM_MOCK_SECRET: LazyLock<[u8; PG_SCRAM_SHA256_KEY_LEN]> = LazyLock::new(rand::random); /// construct a [`UserInfo`](crate::user_info::UserInfo) impl with name /// use default username `greptime` if None is provided @@ -200,6 +210,212 @@ pub fn format_pbkdf2_sha256_password_verifier( )) } +#[derive(Clone, PartialEq, Eq)] +pub struct PgScramSha256Verifier { + iterations: u32, + salt: Vec, + stored_key: Vec, + server_key: Vec, +} + +impl std::fmt::Debug for PgScramSha256Verifier { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PgScramSha256Verifier") + .field("iterations", &self.iterations) + .field("salt", &"") + .field("stored_key", &"") + .field("server_key", &"") + .finish() + } +} + +impl PgScramSha256Verifier { + pub fn new( + iterations: u32, + salt: Vec, + stored_key: Vec, + server_key: Vec, + ) -> Result { + ensure!( + iterations > 0 && iterations <= MAX_PBKDF2_SHA256_ITERATIONS, + IllegalParamSnafu { + msg: format!( + "pg_scram_sha256 iterations must be in 1..={}", + MAX_PBKDF2_SHA256_ITERATIONS + ) + } + ); + ensure!( + !salt.is_empty() && salt.len() <= MAX_PBKDF2_SHA256_SALT_LEN, + IllegalParamSnafu { + msg: format!( + "pg_scram_sha256 salt length must be in 1..={}", + MAX_PBKDF2_SHA256_SALT_LEN + ) + } + ); + ensure!( + stored_key.len() == PG_SCRAM_SHA256_KEY_LEN, + IllegalParamSnafu { + msg: "pg_scram_sha256 stored key must be 32 bytes" + } + ); + ensure!( + server_key.len() == PG_SCRAM_SHA256_KEY_LEN, + IllegalParamSnafu { + msg: "pg_scram_sha256 server key must be 32 bytes" + } + ); + + Ok(Self { + iterations, + salt, + stored_key, + server_key, + }) + } + + pub fn from_password(password: &[u8], salt: &[u8], iterations: u32) -> Result { + let salted_password = pg_scram_sha256_salted_password(password, salt, iterations)?; + Self::from_salted_password(salted_password, salt.to_vec(), iterations) + } + + pub fn from_salted_password( + salted_password: Vec, + salt: Vec, + iterations: u32, + ) -> Result { + ensure!( + salted_password.len() == PG_SCRAM_SHA256_KEY_LEN, + IllegalParamSnafu { + msg: "pg_scram_sha256 salted password must be 32 bytes" + } + ); + let client_key = hmac_sha256(&salted_password, b"Client Key"); + let stored_key = sha256(&client_key); + let server_key = hmac_sha256(&salted_password, b"Server Key"); + Self::new(iterations, salt, stored_key, server_key) + } + + /// Builds a throwaway verifier for an unknown user. The salt is derived + /// deterministically from the username and a process-wide secret, so the + /// SCRAM `server-first-message` (salt + iterations) stays stable per + /// username and indistinguishable from a real user across reconnects. No + /// PBKDF2 is run (avoids a CPU-exhaustion DoS keyed on unknown usernames), + /// and the random keys guarantee the client proof never matches. + pub fn mock_for_unknown_user(username: &[u8]) -> Self { + let mock_salt = hmac_sha256(PG_SCRAM_MOCK_SECRET.as_slice(), username); + Self { + iterations: DEFAULT_PBKDF2_SHA256_ITERATIONS, + salt: mock_salt[..DEFAULT_PBKDF2_SHA256_SALT_LEN].to_vec(), + stored_key: rand::random::<[u8; PG_SCRAM_SHA256_KEY_LEN]>().to_vec(), + server_key: rand::random::<[u8; PG_SCRAM_SHA256_KEY_LEN]>().to_vec(), + } + } + + pub fn iterations(&self) -> u32 { + self.iterations + } + + pub fn salt(&self) -> &[u8] { + &self.salt + } + + pub fn verify_plain_password(&self, password: &[u8]) -> Result { + let salted_password = + pg_scram_sha256_salted_password(password, &self.salt, self.iterations)?; + let client_key = hmac_sha256(&salted_password, b"Client Key"); + let stored_key = sha256(&client_key); + Ok(self.stored_key.ct_eq(&stored_key).into()) + } + + pub fn verify_client_proof( + &self, + auth_message: &[u8], + client_proof: &[u8], + ) -> Result>> { + if client_proof.len() != PG_SCRAM_SHA256_KEY_LEN { + return Ok(None); + } + + let client_signature = hmac_sha256(&self.stored_key, auth_message); + let client_key = xor(client_proof, &client_signature); + let stored_key = sha256(&client_key); + if self.stored_key.ct_eq(&stored_key).into() { + Ok(Some(hmac_sha256(&self.server_key, auth_message))) + } else { + Ok(None) + } + } +} + +pub fn format_pg_scram_sha256_password_verifier( + password: &[u8], + salt: &[u8], + iterations: u32, +) -> Result { + let verifier = PgScramSha256Verifier::from_password(password, salt, iterations)?; + Ok(format!( + "pg_scram_sha256:{iterations}:{}:{}:{}", + hex::encode(verifier.salt), + hex::encode(verifier.stored_key), + hex::encode(verifier.server_key) + )) +} + +fn pg_scram_sha256_salted_password( + password: &[u8], + salt: &[u8], + iterations: u32, +) -> Result> { + ensure!( + iterations > 0 && iterations <= MAX_PBKDF2_SHA256_ITERATIONS, + IllegalParamSnafu { + msg: format!( + "pg_scram_sha256 iterations must be in 1..={}", + MAX_PBKDF2_SHA256_ITERATIONS + ) + } + ); + ensure!( + !salt.is_empty() && salt.len() <= MAX_PBKDF2_SHA256_SALT_LEN, + IllegalParamSnafu { + msg: format!( + "pg_scram_sha256 salt length must be in 1..={}", + MAX_PBKDF2_SHA256_SALT_LEN + ) + } + ); + + let prepared_password = std::str::from_utf8(password) + .ok() + .and_then(|password| stringprep::saslprep(password).ok()); + let password = prepared_password + .as_deref() + .map(str::as_bytes) + .unwrap_or(password); + + let mut salted_password = [0u8; PG_SCRAM_SHA256_KEY_LEN]; + pbkdf2_hmac::(password, salt, iterations, &mut salted_password); + Ok(salted_password.to_vec()) +} + +fn hmac_sha256(key: &[u8], msg: &[u8]) -> Vec { + let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key length"); + mac.update(msg); + mac.finalize().into_bytes().to_vec() +} + +fn sha256(data: &[u8]) -> Vec { + let mut hasher = Sha256::new(); + hasher.update(data); + hasher.finalize().to_vec() +} + +fn xor(lhs: &[u8], rhs: &[u8]) -> Vec { + lhs.iter().zip(rhs).map(|(l, r)| l ^ r).collect() +} + fn sha1_two(input_1: &[u8], input_2: &[u8]) -> Vec { let mut hasher = Sha1::new(); hasher.update(input_1); @@ -274,4 +490,62 @@ mod tests { .is_err() ); } + + #[test] + fn test_format_pg_scram_sha256_password_verifier() { + let verifier = + format_pg_scram_sha256_password_verifier("password".as_bytes(), b"salt", 4096).unwrap(); + assert_eq!( + "pg_scram_sha256:4096:73616c74:945e1c466fc9932efadc23781edc5d1e78d5e10f005933652af1a6105154f084:b9bf0e811b1fb6793671c0cc3adedf7c75cd72291191092ad65878c5a02aad2c", + verifier + ); + } + + #[test] + fn test_pg_scram_sha256_applies_saslprep() { + let normalized = PgScramSha256Verifier::from_password(b"pass word", b"salt", 4096).unwrap(); + let non_breaking_space = + PgScramSha256Verifier::from_password("pass\u{00a0}word".as_bytes(), b"salt", 4096) + .unwrap(); + + assert_eq!(normalized, non_breaking_space); + assert!( + non_breaking_space + .verify_plain_password(b"pass word") + .unwrap() + ); + } + + #[test] + fn test_pg_scram_sha256_uses_original_bytes_when_saslprep_is_not_possible() { + let invalid_utf8 = b"password\xff"; + let prohibited = b"password\x07"; + + for password in [invalid_utf8.as_slice(), prohibited.as_slice()] { + let verifier = PgScramSha256Verifier::from_password(password, b"salt", 4096).unwrap(); + assert!(verifier.verify_plain_password(password).unwrap()); + } + } + + #[test] + fn test_mock_verifier_is_stable_per_username() { + let alice = PgScramSha256Verifier::mock_for_unknown_user(b"alice"); + let alice_again = PgScramSha256Verifier::mock_for_unknown_user(b"alice"); + let bob = PgScramSha256Verifier::mock_for_unknown_user(b"bob"); + + // Same username must yield the same salt and iterations across reconnects, + // otherwise the SCRAM server-first message leaks user (non-)existence. + assert_eq!(alice.salt(), alice_again.salt()); + assert_ne!(alice.salt(), bob.salt()); + assert_eq!(alice.iterations(), DEFAULT_PBKDF2_SHA256_ITERATIONS); + assert_eq!(alice.salt().len(), DEFAULT_PBKDF2_SHA256_SALT_LEN); + + // The mock verifier must never accept a client proof. + assert!( + alice + .verify_client_proof(b"auth-message", &[0u8; PG_SCRAM_SHA256_KEY_LEN]) + .unwrap() + .is_none() + ); + } } diff --git a/src/auth/src/lib.rs b/src/auth/src/lib.rs index 180042eefd8..f1dab981452 100644 --- a/src/auth/src/lib.rs +++ b/src/auth/src/lib.rs @@ -22,16 +22,17 @@ mod user_provider; pub mod tests; pub use common::{ - DEFAULT_PBKDF2_SHA256_ITERATIONS, HashedPassword, Identity, MAX_PBKDF2_SHA256_ITERATIONS, - MAX_PBKDF2_SHA256_SALT_LEN, PBKDF2_SHA256_HASH_LEN, Password, auth_mysql, + DEFAULT_PBKDF2_SHA256_ITERATIONS, DEFAULT_PBKDF2_SHA256_SALT_LEN, HashedPassword, Identity, + MAX_PBKDF2_SHA256_ITERATIONS, MAX_PBKDF2_SHA256_SALT_LEN, PBKDF2_SHA256_HASH_LEN, + PG_SCRAM_SHA256_KEY_LEN, Password, PgScramSha256Verifier, auth_mysql, format_mysql_native_password_verifier, format_pbkdf2_sha256_password_verifier, - mysql_native_password_hash, static_user_provider_from_option, user_provider_from_option, - userinfo_by_name, + format_pg_scram_sha256_password_verifier, mysql_native_password_hash, + static_user_provider_from_option, user_provider_from_option, userinfo_by_name, }; pub use permission::{DefaultPermissionChecker, PermissionChecker, PermissionReq, PermissionResp}; pub use user_info::UserInfo; -pub use user_provider::UserProvider; pub use user_provider::static_user_provider::StaticUserProvider; +pub use user_provider::{PgAuthInfo, UserProvider}; /// pub type alias pub type UserInfoRef = std::sync::Arc; diff --git a/src/auth/src/permission.rs b/src/auth/src/permission.rs index 8914635290f..cff9b414777 100644 --- a/src/auth/src/permission.rs +++ b/src/auth/src/permission.rs @@ -36,7 +36,11 @@ pub enum PermissionReq<'a> { PromStoreRead, Otlp, LogWrite, - BulkInsert, + BulkInsert { + catalog: &'a str, + schema: &'a str, + table: &'a str, + }, } impl<'a> PermissionReq<'a> { @@ -57,7 +61,7 @@ impl<'a> PermissionReq<'a> { | PermissionReq::PromStoreWrite | PermissionReq::Otlp | PermissionReq::LogWrite - | PermissionReq::BulkInsert => false, + | PermissionReq::BulkInsert { .. } => false, } } @@ -79,6 +83,15 @@ pub trait PermissionChecker: Send + Sync { user_info: UserInfoRef, req: PermissionReq, ) -> Result; + + fn check_permission_with_context( + &self, + user_info: UserInfoRef, + req: PermissionReq, + _current_schema: Option<&str>, + ) -> Result { + self.check_permission(user_info, req) + } } impl PermissionChecker for Option<&PermissionCheckerRef> { @@ -86,13 +99,24 @@ impl PermissionChecker for Option<&PermissionCheckerRef> { &self, user_info: UserInfoRef, req: PermissionReq, + ) -> Result { + self.check_permission_with_context(user_info, req, None) + } + + fn check_permission_with_context( + &self, + user_info: UserInfoRef, + req: PermissionReq, + current_schema: Option<&str>, ) -> Result { match self { - Some(checker) => match checker.check_permission(user_info, req) { - Ok(PermissionResp::Reject) => PermissionDeniedSnafu.fail(), - Ok(PermissionResp::Allow) => Ok(PermissionResp::Allow), - Err(e) => Err(e), - }, + Some(checker) => { + match checker.check_permission_with_context(user_info, req, current_schema) { + Ok(PermissionResp::Reject) => PermissionDeniedSnafu.fail(), + Ok(PermissionResp::Allow) => Ok(PermissionResp::Allow), + Err(e) => Err(e), + } + } None => Ok(PermissionResp::Allow), } } @@ -209,4 +233,15 @@ mod tests { assert!(req.is_write()); } + + #[test] + fn test_bulk_insert_is_write_request() { + let req = PermissionReq::BulkInsert { + catalog: "greptime", + schema: "public", + table: "metrics", + }; + + assert!(req.is_write()); + } } diff --git a/src/auth/src/user_provider.rs b/src/auth/src/user_provider.rs index 2ff61bbe9bc..344e6669015 100644 --- a/src/auth/src/user_provider.rs +++ b/src/auth/src/user_provider.rs @@ -28,8 +28,9 @@ use snafu::{OptionExt, ResultExt, ensure}; use subtle::ConstantTimeEq; use crate::common::{ - Identity, MAX_PBKDF2_SHA256_ITERATIONS, MAX_PBKDF2_SHA256_SALT_LEN, PBKDF2_SHA256_HASH_LEN, - Password, auth_mysql_with_hash_stage_2, + DEFAULT_PBKDF2_SHA256_SALT_LEN, Identity, MAX_PBKDF2_SHA256_ITERATIONS, + MAX_PBKDF2_SHA256_SALT_LEN, PBKDF2_SHA256_HASH_LEN, PG_SCRAM_SHA256_KEY_LEN, Password, + PgScramSha256Verifier, auth_mysql_with_hash_stage_2, }; use crate::error::{ IllegalParamSnafu, InvalidConfigSnafu, IoSnafu, Result, UnsupportedPasswordTypeSnafu, @@ -64,15 +65,35 @@ pub trait UserProvider: Send + Sync { Ok(user_info) } + async fn postgres_auth_info(&self, _id: Identity<'_>) -> Result { + Ok(PgAuthInfo::Cleartext) + } + /// Returns whether this user provider implementation is backed by an external system. fn external(&self) -> bool { false } } -#[derive(Clone, PartialEq, Eq)] +pub enum PgAuthInfo { + ScramSha256 { + verifier: PgScramSha256Verifier, + user_info: Option, + }, + Cleartext, +} + +#[derive(Clone)] pub(crate) enum PasswordVerifier { - PlainText(String), + PlainText { + password: String, + /// SCRAM verifier derived once at load time. Precomputing it keeps the + /// Postgres SCRAM `server-first-message` (stable salt, fixed iterations) + /// and per-connection cost indistinguishable from stored-hash and + /// unknown users, instead of running PBKDF2 with a fresh salt on every + /// connection. + scram: PgScramSha256Verifier, + }, Pbkdf2Sha256 { iterations: u32, salt: Vec, @@ -81,12 +102,46 @@ pub(crate) enum PasswordVerifier { MysqlNativePassword { hash_stage_2: Vec, }, + PgScramSha256(PgScramSha256Verifier), } +impl PartialEq for PasswordVerifier { + fn eq(&self, other: &Self) -> bool { + // The precomputed SCRAM verifier is a cache derived from the password, so + // two plaintext verifiers are equal iff their passwords match. + match (self, other) { + ( + PasswordVerifier::PlainText { password: a, .. }, + PasswordVerifier::PlainText { password: b, .. }, + ) => a == b, + ( + PasswordVerifier::Pbkdf2Sha256 { + iterations: i1, + salt: s1, + hash: h1, + }, + PasswordVerifier::Pbkdf2Sha256 { + iterations: i2, + salt: s2, + hash: h2, + }, + ) => i1 == i2 && s1 == s2 && h1 == h2, + ( + PasswordVerifier::MysqlNativePassword { hash_stage_2: a }, + PasswordVerifier::MysqlNativePassword { hash_stage_2: b }, + ) => a == b, + (PasswordVerifier::PgScramSha256(a), PasswordVerifier::PgScramSha256(b)) => a == b, + _ => false, + } + } +} + +impl Eq for PasswordVerifier {} + impl fmt::Debug for PasswordVerifier { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { - PasswordVerifier::PlainText(_) => { + PasswordVerifier::PlainText { .. } => { f.debug_tuple("PlainText").field(&"").finish() } PasswordVerifier::Pbkdf2Sha256 { iterations, .. } => f @@ -99,14 +154,33 @@ impl fmt::Debug for PasswordVerifier { .debug_struct("MysqlNativePassword") .field("hash_stage_2", &"") .finish(), + PasswordVerifier::PgScramSha256(_) => f + .debug_struct("PgScramSha256") + .field("verifier", &"") + .finish(), } } } impl PasswordVerifier { + /// Builds a plaintext verifier, precomputing a stable SCRAM verifier so the + /// Postgres SCRAM handshake for this user carries a stable salt and runs no + /// per-connection PBKDF2. Uses [`DEFAULT_PBKDF2_SHA256_ITERATIONS`] to match + /// the mock verifier handed to unknown users. + fn plain_text(password: String) -> Option { + let salt = rand::random::<[u8; DEFAULT_PBKDF2_SHA256_SALT_LEN]>(); + let scram = PgScramSha256Verifier::from_password( + password.as_bytes(), + &salt, + crate::DEFAULT_PBKDF2_SHA256_ITERATIONS, + ) + .ok()?; + Some(Self::PlainText { password, scram }) + } + fn parse(input: &str) -> Option { if let Some(password) = input.strip_prefix("plain:") { - return Some(Self::PlainText(password.to_string())); + return Self::plain_text(password.to_string()); } if let Some(verifier) = input.strip_prefix("pbkdf2_sha256:") { @@ -140,14 +214,50 @@ impl PasswordVerifier { return Some(Self::MysqlNativePassword { hash_stage_2 }); } - Some(Self::PlainText(input.to_string())) + if let Some(verifier) = input.strip_prefix("pg_scram_sha256:") { + let mut parts = verifier.split(':'); + let iterations = parts.next()?.parse::().ok()?; + let salt = hex::decode(parts.next()?).ok()?; + let stored_key = hex::decode(parts.next()?).ok()?; + let server_key = hex::decode(parts.next()?).ok()?; + if parts.next().is_some() + || stored_key.len() != PG_SCRAM_SHA256_KEY_LEN + || server_key.len() != PG_SCRAM_SHA256_KEY_LEN + { + return None; + } + + return PgScramSha256Verifier::new(iterations, salt, stored_key, server_key) + .ok() + .map(Self::PgScramSha256); + } + + Self::plain_text(input.to_string()) + } + + fn supports_pg_scram_sha256(&self) -> bool { + matches!( + self, + PasswordVerifier::PlainText { .. } | PasswordVerifier::PgScramSha256(_) + ) + } + + fn to_pg_scram_sha256_verifier(&self) -> Option { + match self { + PasswordVerifier::PlainText { scram, .. } => Some(scram.clone()), + // Legacy PBKDF2 verifiers were derived without SASLprep, and the + // original password is unavailable to normalize them safely. + PasswordVerifier::Pbkdf2Sha256 { .. } => None, + PasswordVerifier::PgScramSha256(verifier) => Some(verifier.clone()), + PasswordVerifier::MysqlNativePassword { .. } => None, + } } fn verify_plain_text(&self, password: &str) -> bool { match self { - PasswordVerifier::PlainText(expected) => { - expected.as_bytes().ct_eq(password.as_bytes()).into() - } + PasswordVerifier::PlainText { + password: expected, .. + } => expected.as_bytes().ct_eq(password.as_bytes()).into(), PasswordVerifier::Pbkdf2Sha256 { iterations, salt, @@ -161,6 +271,9 @@ impl PasswordVerifier { hash.as_slice().ct_eq(&actual[..]).into() } PasswordVerifier::MysqlNativePassword { .. } => false, + PasswordVerifier::PgScramSha256(verifier) => verifier + .verify_plain_password(password.as_bytes()) + .unwrap_or(false), } } @@ -171,7 +284,7 @@ impl PasswordVerifier { username: &str, ) -> Result<()> { match self { - PasswordVerifier::PlainText(password) => { + PasswordVerifier::PlainText { password, .. } => { auth_mysql(auth_data, salt, username, password.as_bytes()) } PasswordVerifier::MysqlNativePassword { hash_stage_2 } => { @@ -181,6 +294,10 @@ impl PasswordVerifier { password_type: "mysql_native_password_with_pbkdf2_sha256_verifier", } .fail(), + PasswordVerifier::PgScramSha256(_) => UnsupportedPasswordTypeSnafu { + password_type: "mysql_native_password_with_pg_scram_sha256_verifier", + } + .fail(), } } } @@ -243,6 +360,7 @@ fn load_credential_from_file(filepath: &str) -> Result { /// - `plain:` /// - `pbkdf2_sha256:::` /// - `mysql_native_password:` +/// - `pg_scram_sha256::::` pub(crate) fn parse_credential_line( line: &str, ) -> Option<(String, (PasswordVerifier, PermissionMode))> { @@ -263,6 +381,53 @@ pub(crate) fn parse_credential_line( Some((username.to_string(), (verifier, permission_mode))) } +pub(crate) fn postgres_auth_info_with_credential( + users: &UserInfoMap, + input_id: Identity<'_>, +) -> Result { + match input_id { + Identity::UserId(username, _) => { + ensure!( + !username.is_empty(), + IllegalParamSnafu { + msg: "blank username" + } + ); + + if !users + .values() + .all(|(verifier, _)| verifier.supports_pg_scram_sha256()) + { + // PostgreSQL chooses one auth method during startup. Selecting it + // per username would expose user or verifier existence. + return Ok(PgAuthInfo::Cleartext); + } + + if let Some((verifier, permission_mode)) = users.get(username) { + if let Some(verifier) = verifier.to_pg_scram_sha256_verifier() { + return Ok(PgAuthInfo::ScramSha256 { + verifier, + user_info: Some(DefaultUserInfo::with_name_and_permission( + username, + *permission_mode, + )), + }); + } + + return Ok(PgAuthInfo::Cleartext); + } + + // Unknown user: hand back a deterministic mock verifier so the SCRAM + // handshake is indistinguishable from a real user, without running + // PBKDF2 or leaking existence through an unstable salt. + Ok(PgAuthInfo::ScramSha256 { + verifier: PgScramSha256Verifier::mock_for_unknown_user(username.as_bytes()), + user_info: None, + }) + } + } +} + fn authenticate_with_credential( users: &UserInfoMap, input_id: Identity<'_>, @@ -317,10 +482,10 @@ mod tests { use sha1::Sha1; use super::*; - use crate::common::mysql_native_password_hash; + use crate::common::{format_pg_scram_sha256_password_verifier, mysql_native_password_hash}; fn plain(password: &str) -> PasswordVerifier { - PasswordVerifier::PlainText(password.to_string()) + PasswordVerifier::plain_text(password.to_string()).unwrap() } fn sha1_one(data: &[u8]) -> Vec { @@ -484,6 +649,23 @@ mod tests { let result = parse_credential_line("user=mysql_native_password:abcd"); assert_eq!(result, None); + let verifier = + format_pg_scram_sha256_password_verifier(b"password", b"salt", 4096).unwrap(); + let result = parse_credential_line(&format!("user={verifier}")); + assert!(matches!( + result, + Some(( + _, + ( + PasswordVerifier::PgScramSha256(PgScramSha256Verifier { .. }), + PermissionMode::ReadWrite + ) + )) + )); + + let result = parse_credential_line("user=pg_scram_sha256:4096:73616c74:abcd:abcd"); + assert_eq!(result, None); + // Invalid format - no equals sign let result = parse_credential_line("invalid_line"); assert_eq!(result, None); @@ -541,7 +723,7 @@ mod tests { let users = HashMap::from([( "user".to_string(), ( - PasswordVerifier::PlainText(password.to_string()), + PasswordVerifier::plain_text(password.to_string()).unwrap(), PermissionMode::default(), ), )]); @@ -585,9 +767,148 @@ mod tests { assert!(result.is_err()); } + #[test] + fn test_authenticate_with_pg_scram_sha256_verifier_plain_text() { + let verifier = + format_pg_scram_sha256_password_verifier(b"password", b"salt", 4096).unwrap(); + let (_, user) = parse_credential_line(&format!("user={verifier}")).unwrap(); + let users = HashMap::from([("user".to_string(), user)]); + + let result = authenticate_with_credential( + &users, + Identity::UserId("user", None), + Password::PlainText("password".to_string().into()), + ); + assert!(result.is_ok()); + + let result = authenticate_with_credential( + &users, + Identity::UserId("user", None), + Password::PlainText("wrong".to_string().into()), + ); + assert!(result.is_err()); + } + + #[test] + fn test_plain_text_scram_verifier_is_stable() { + let verifier = PasswordVerifier::plain_text("password".to_string()).unwrap(); + let first = verifier.to_pg_scram_sha256_verifier().unwrap(); + let second = verifier.to_pg_scram_sha256_verifier().unwrap(); + + // A plaintext-backed user must present a stable salt and fixed iterations + // across connections without re-running PBKDF2, otherwise the SCRAM + // server-first message (and its timing) leaks that the user exists. + assert_eq!(first.salt(), second.salt()); + assert_eq!(first.salt().len(), DEFAULT_PBKDF2_SHA256_SALT_LEN); + assert_eq!(first.iterations(), crate::DEFAULT_PBKDF2_SHA256_ITERATIONS); + + // The derived verifier must still accept the real password. + assert!(first.verify_plain_password(b"password").unwrap()); + assert!(!first.verify_plain_password(b"wrong").unwrap()); + } + + #[test] + fn test_postgres_auth_info_uses_scram_for_unknown_user() { + let verifier = + format_pg_scram_sha256_password_verifier(b"password", b"salt", 4096).unwrap(); + let (_, user) = parse_credential_line(&format!("user={verifier}")).unwrap(); + let users = HashMap::from([("user".to_string(), user)]); + + let auth_info = + postgres_auth_info_with_credential(&users, Identity::UserId("unknown", None)).unwrap(); + assert!(matches!( + auth_info, + PgAuthInfo::ScramSha256 { + user_info: None, + .. + } + )); + } + + #[test] + fn test_postgres_auth_info_falls_back_to_cleartext() { + let hash_stage_2 = mysql_native_password_hash("password".as_bytes()); + let users = HashMap::from([( + "user".to_string(), + ( + PasswordVerifier::MysqlNativePassword { hash_stage_2 }, + PermissionMode::default(), + ), + )]); + + let auth_info = + postgres_auth_info_with_credential(&users, Identity::UserId("user", None)).unwrap(); + assert!(matches!(auth_info, PgAuthInfo::Cleartext)); + + let auth_info = + postgres_auth_info_with_credential(&users, Identity::UserId("unknown", None)).unwrap(); + assert!(matches!(auth_info, PgAuthInfo::Cleartext)); + } + + #[test] + fn test_postgres_auth_info_with_pbkdf2_falls_back_to_cleartext() { + let iterations = 4096; + let salt = b"salt"; + let password = "pass\u{00a0}word"; + let mut hash = [0u8; PBKDF2_SHA256_HASH_LEN]; + pbkdf2_hmac::(password.as_bytes(), salt, iterations, &mut hash); + let users = HashMap::from([( + "user".to_string(), + ( + PasswordVerifier::Pbkdf2Sha256 { + iterations, + salt: salt.to_vec(), + hash: hash.to_vec(), + }, + PermissionMode::default(), + ), + )]); + + let auth_info = + postgres_auth_info_with_credential(&users, Identity::UserId("user", None)).unwrap(); + assert!(matches!(auth_info, PgAuthInfo::Cleartext)); + + assert!( + authenticate_with_credential( + &users, + Identity::UserId("user", None), + Password::PlainText(password.to_string().into()), + ) + .is_ok() + ); + } + + #[test] + fn test_postgres_auth_info_mixed_verifiers_fall_back_to_cleartext() { + let verifier = + format_pg_scram_sha256_password_verifier(b"password", b"salt", 4096).unwrap(); + let (_, scram_user) = parse_credential_line(&format!("scram={verifier}")).unwrap(); + let mysql_user = ( + PasswordVerifier::MysqlNativePassword { + hash_stage_2: mysql_native_password_hash("password".as_bytes()), + }, + PermissionMode::default(), + ); + let users = HashMap::from([ + ("scram".to_string(), scram_user), + ("mysql".to_string(), mysql_user), + ]); + + let auth_info = + postgres_auth_info_with_credential(&users, Identity::UserId("scram", None)).unwrap(); + assert!(matches!(auth_info, PgAuthInfo::Cleartext)); + + let auth_info = + postgres_auth_info_with_credential(&users, Identity::UserId("unknown", None)).unwrap(); + assert!(matches!(auth_info, PgAuthInfo::Cleartext)); + } + #[test] fn test_password_verifier_debug_redacts_secrets() { - let debug = format!("{:?}", PasswordVerifier::PlainText("secret".to_string())); + let debug = format!( + "{:?}", + PasswordVerifier::plain_text("secret".to_string()).unwrap() + ); assert!(debug.contains("")); assert!(!debug.contains("secret")); diff --git a/src/auth/src/user_provider/static_user_provider.rs b/src/auth/src/user_provider/static_user_provider.rs index 689b9627261..4ea7efb4345 100644 --- a/src/auth/src/user_provider/static_user_provider.rs +++ b/src/auth/src/user_provider/static_user_provider.rs @@ -17,7 +17,8 @@ use snafu::OptionExt; use crate::error::{InvalidConfigSnafu, Result}; use crate::user_provider::{ - UserInfoMap, authenticate_with_credential, load_credential_from_file, parse_credential_line, + PgAuthInfo, UserInfoMap, authenticate_with_credential, load_credential_from_file, + parse_credential_line, postgres_auth_info_with_credential, }; use crate::{Identity, Password, UserInfoRef, UserProvider}; @@ -67,6 +68,10 @@ impl UserProvider for StaticUserProvider { authenticate_with_credential(&self.users, id, pwd) } + async fn postgres_auth_info(&self, id: Identity<'_>) -> Result { + postgres_auth_info_with_credential(&self.users, id) + } + async fn authorize( &self, _catalog: &str, diff --git a/src/auth/src/user_provider/watch_file_user_provider.rs b/src/auth/src/user_provider/watch_file_user_provider.rs index 451efd5cc4d..4d39ae3e4bc 100644 --- a/src/auth/src/user_provider/watch_file_user_provider.rs +++ b/src/auth/src/user_provider/watch_file_user_provider.rs @@ -20,7 +20,10 @@ use common_telemetry::{info, warn}; use snafu::ResultExt; use crate::error::{FileWatchSnafu, Result}; -use crate::user_provider::{UserInfoMap, authenticate_with_credential, load_credential_from_file}; +use crate::user_provider::{ + PgAuthInfo, UserInfoMap, authenticate_with_credential, load_credential_from_file, + postgres_auth_info_with_credential, +}; use crate::{Identity, Password, UserInfoRef, UserProvider}; pub(crate) const WATCH_FILE_USER_PROVIDER: &str = "watch_file_user_provider"; @@ -83,6 +86,11 @@ impl UserProvider for WatchFileUserProvider { authenticate_with_credential(&users, id, password) } + async fn postgres_auth_info(&self, id: Identity<'_>) -> Result { + let users = self.users.lock().expect("users credential must be valid"); + postgres_auth_info_with_credential(&users, id) + } + async fn authorize(&self, _: &str, _: &str, _: &UserInfoRef) -> Result<()> { // default allow all Ok(()) diff --git a/src/catalog/src/information_extension.rs b/src/catalog/src/information_extension.rs index 08d81ac1f03..c41c816c6d7 100644 --- a/src/catalog/src/information_extension.rs +++ b/src/catalog/src/information_extension.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + use api::v1::meta::ProcedureStatus; use common_error::ext::BoxedError; use common_meta::cluster::{ClusterInfo, NodeInfo, Role}; @@ -22,8 +24,16 @@ use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor}; use common_meta::rpc::procedure; use common_procedure::{ProcedureInfo, ProcedureState}; use common_query::request::QueryRequest; -use common_recordbatch::SendableRecordBatchStream; +use common_recordbatch::adapter::{AsyncRecordBatchStreamAdapter, DfRecordBatchStreamAdapter}; use common_recordbatch::util::{ChainedRecordBatchStream, LimitedRecordBatchStream}; +use common_recordbatch::{DfSendableRecordBatchStream, SendableRecordBatchStream}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; +use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef; +use datatypes::schema::SchemaRef; +use futures_util::future::try_join_all; use meta_client::MetaClientRef; use snafu::ResultExt; use store_api::storage::RegionId; @@ -43,6 +53,168 @@ impl DistributedInformationExtension { datanode_manager, } } + + async fn inspect_datanode_stream( + meta_client: MetaClientRef, + datanode_manager: DatanodeManagerRef, + request: DatanodeInspectRequest, + ) -> std::result::Result { + let nodes = meta_client + .list_nodes(Some(Role::Datanode)) + .await + .map_err(BoxedError::new) + .context(crate::error::ListNodesSnafu)?; + + let limit = request.scan.limit; + let plan = request + .build_plan() + .context(crate::error::DatafusionSnafu)?; + + let streams = try_join_all(nodes.into_iter().map(|node| { + let datanode_manager = datanode_manager.clone(); + let plan = plan.clone(); + async move { + let client = datanode_manager.datanode(&node.peer).await; + client + .handle_query(QueryRequest { + plan, + region_id: RegionId::default(), + header: None, + }) + .await + .context(crate::error::HandleQuerySnafu) + } + })) + .await?; + + let chained = + ChainedRecordBatchStream::new(streams).context(crate::error::CreateRecordBatchSnafu)?; + match limit { + Some(limit) => Ok(Box::pin(LimitedRecordBatchStream::new( + Box::pin(chained), + limit, + ))), + None => Ok(Box::pin(chained)), + } + } +} + +struct DistributedInspectExec { + request: DatanodeInspectRequest, + meta_client: MetaClientRef, + datanode_manager: DatanodeManagerRef, + schema: SchemaRef, + arrow_schema: ArrowSchemaRef, + properties: Arc, +} + +impl std::fmt::Debug for DistributedInspectExec { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DistributedInspectExec") + .field("request", &self.request) + .field("schema", &self.schema) + .finish_non_exhaustive() + } +} + +impl DistributedInspectExec { + fn new( + request: DatanodeInspectRequest, + schema: SchemaRef, + meta_client: MetaClientRef, + datanode_manager: DatanodeManagerRef, + ) -> Self { + let arrow_schema = schema.arrow_schema().clone(); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(arrow_schema.clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + request, + meta_client, + datanode_manager, + schema, + arrow_schema, + properties, + } + } +} + +impl DisplayAs for DistributedInspectExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "DistributedInspectExec: kind={:?}, scan={}, schema={:?}", + self.request.kind, self.request.scan, self.arrow_schema + ) + } +} + +impl ExecutionPlan for DistributedInspectExec { + fn name(&self) -> &str { + "DistributedInspectExec" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn schema(&self) -> ArrowSchemaRef { + self.arrow_schema.clone() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> datafusion::error::Result> { + if children.is_empty() { + Ok(self) + } else { + Err(datafusion::error::DataFusionError::Internal( + "DistributedInspectExec is a leaf execution plan and cannot accept children" + .to_string(), + )) + } + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> datafusion::error::Result { + if partition != 0 { + return Err(datafusion::error::DataFusionError::Execution(format!( + "DistributedInspectExec invalid partition. Expected 0, got {partition}" + ))); + } + + let meta_client = self.meta_client.clone(); + let datanode_manager = self.datanode_manager.clone(); + let request = self.request.clone(); + let schema = self.schema.clone(); + let future = async move { + DistributedInformationExtension::inspect_datanode_stream( + meta_client, + datanode_manager, + request, + ) + .await + .map_err(BoxedError::new) + .context(common_recordbatch::error::ExternalSnafu) + }; + let stream = AsyncRecordBatchStreamAdapter::new(schema, Box::pin(future)); + Ok(Box::pin(DfRecordBatchStreamAdapter::new(Box::pin(stream)))) + } } #[async_trait::async_trait] @@ -112,41 +284,24 @@ impl InformationExtension for DistributedInformationExtension { &self, request: DatanodeInspectRequest, ) -> std::result::Result { - // Aggregate results from all datanodes - let nodes = self - .meta_client - .list_nodes(Some(Role::Datanode)) - .await - .map_err(BoxedError::new) - .context(crate::error::ListNodesSnafu)?; + Self::inspect_datanode_stream( + self.meta_client.clone(), + self.datanode_manager.clone(), + request, + ) + .await + } - let limit = request.scan.limit; - let plan = request - .build_plan() - .context(crate::error::DatafusionSnafu)?; - - let mut streams = Vec::with_capacity(nodes.len()); - for node in nodes { - let client = self.datanode_manager.datanode(&node.peer).await; - let stream = client - .handle_query(QueryRequest { - plan: plan.clone(), - region_id: RegionId::default(), - header: None, - }) - .await - .context(crate::error::HandleQuerySnafu)?; - streams.push(stream); - } - - let chained = - ChainedRecordBatchStream::new(streams).context(crate::error::CreateRecordBatchSnafu)?; - match limit { - Some(limit) => Ok(Box::pin(LimitedRecordBatchStream::new( - Box::pin(chained), - limit, - ))), - None => Ok(Box::pin(chained)), - } + fn inspect_datanode_plan( + &self, + request: DatanodeInspectRequest, + schema: SchemaRef, + ) -> std::result::Result>, Self::Error> { + Ok(Some(Arc::new(DistributedInspectExec::new( + request, + schema, + self.meta_client.clone(), + self.datanode_manager.clone(), + )))) } } diff --git a/src/catalog/src/system_schema.rs b/src/catalog/src/system_schema.rs index 6c6a5fc2028..fe926d893c9 100644 --- a/src/catalog/src/system_schema.rs +++ b/src/catalog/src/system_schema.rs @@ -25,6 +25,7 @@ use std::sync::Arc; use common_error::ext::BoxedError; use common_recordbatch::{RecordBatchStreamWrapper, SendableRecordBatchStream}; use common_telemetry::tracing::Span; +use datafusion::physical_plan::ExecutionPlan; use datatypes::schema::SchemaRef; use futures_util::StreamExt; use snafu::ResultExt; @@ -107,6 +108,10 @@ pub trait SystemTable { fn to_stream(&self, request: ScanRequest) -> Result; + fn scan_plan(&self, _request: ScanRequest) -> Result>> { + Ok(None) + } + fn table_type(&self) -> TableType { TableType::Temporary } @@ -173,4 +178,15 @@ impl DataSource for SystemTableDataSource { Ok(Box::pin(stream)) } + + fn get_physical_plan( + &self, + request: ScanRequest, + ) -> std::result::Result>, BoxedError> { + self.table + .scan_plan(request) + .map_err(BoxedError::new) + .context(TablesRecordBatchSnafu) + .map_err(BoxedError::new) + } } diff --git a/src/catalog/src/system_schema/information_schema.rs b/src/catalog/src/system_schema/information_schema.rs index 01e366da886..0a0864aed85 100644 --- a/src/catalog/src/system_schema/information_schema.rs +++ b/src/catalog/src/system_schema/information_schema.rs @@ -46,6 +46,7 @@ use common_procedure::ProcedureInfo; use common_recordbatch::SendableRecordBatchStream; use datafusion::error::DataFusionError; use datafusion::logical_expr::LogicalPlan; +use datafusion::physical_plan::ExecutionPlan; use datatypes::schema::SchemaRef; use lazy_static::lazy_static; use paste::paste; @@ -420,6 +421,10 @@ pub trait InformationTable { fn to_stream(&self, request: ScanRequest) -> Result; + fn scan_plan(&self, _request: ScanRequest) -> Result>> { + Ok(None) + } + fn table_type(&self) -> TableType { TableType::Temporary } @@ -449,6 +454,10 @@ where fn to_stream(&self, request: ScanRequest) -> Result { InformationTable::to_stream(self, request) } + + fn scan_plan(&self, request: ScanRequest) -> Result>> { + InformationTable::scan_plan(self, request) + } } pub type InformationExtensionRef = Arc + Send + Sync>; @@ -475,6 +484,16 @@ pub trait InformationExtension { &self, request: DatanodeInspectRequest, ) -> std::result::Result; + + /// Builds a physical plan for datanode inspect if the extension can expose + /// the distributed fan-in semantics to DataFusion. + fn inspect_datanode_plan( + &self, + _request: DatanodeInspectRequest, + _schema: SchemaRef, + ) -> std::result::Result>, Self::Error> { + Ok(None) + } } /// The request to inspect the datanode. diff --git a/src/catalog/src/system_schema/information_schema/region_info.rs b/src/catalog/src/system_schema/information_schema/region_info.rs index ffc9dfc7ae7..56dc04b85c8 100644 --- a/src/catalog/src/system_schema/information_schema/region_info.rs +++ b/src/catalog/src/system_schema/information_schema/region_info.rs @@ -18,6 +18,7 @@ use common_catalog::consts::INFORMATION_SCHEMA_REGION_INFO_TABLE_ID; use common_error::ext::BoxedError; use common_recordbatch::SendableRecordBatchStream; use common_recordbatch::adapter::AsyncRecordBatchStreamAdapter; +use datafusion::physical_plan::ExecutionPlan; use datatypes::schema::SchemaRef; use snafu::ResultExt; use store_api::region_info::RegionInfoEntry; @@ -83,4 +84,20 @@ impl InformationTable for InformationSchemaRegionInfo { Box::pin(future), ))) } + + fn scan_plan(&self, request: ScanRequest) -> Result>> { + let schema = if let Some(p) = request.projection_indices() { + Arc::new(self.schema.try_project(p).context(ProjectSchemaSnafu)?) + } else { + self.schema.clone() + }; + + let info_ext = utils::information_extension(&self.catalog_manager)?; + let req = DatanodeInspectRequest { + kind: DatanodeInspectKind::RegionInfo, + scan: request, + }; + + info_ext.inspect_datanode_plan(req, schema) + } } diff --git a/src/catalog/src/system_schema/information_schema/ssts.rs b/src/catalog/src/system_schema/information_schema/ssts.rs index 520fc6d485d..19c44aee343 100644 --- a/src/catalog/src/system_schema/information_schema/ssts.rs +++ b/src/catalog/src/system_schema/information_schema/ssts.rs @@ -21,6 +21,7 @@ use common_catalog::consts::{ use common_error::ext::BoxedError; use common_recordbatch::SendableRecordBatchStream; use common_recordbatch::adapter::AsyncRecordBatchStreamAdapter; +use datafusion::physical_plan::ExecutionPlan; use datatypes::schema::SchemaRef; use snafu::ResultExt; use store_api::sst_entry::{ManifestSstEntry, PuffinIndexMetaEntry, StorageSstEntry}; @@ -34,6 +35,30 @@ use crate::information_schema::{ }; use crate::system_schema::utils; +fn projected_schema(schema: &SchemaRef, request: &ScanRequest) -> Result { + if let Some(p) = request.projection_indices() { + Ok(Arc::new(schema.try_project(p).context(ProjectSchemaSnafu)?)) + } else { + Ok(schema.clone()) + } +} + +fn inspect_plan( + catalog_manager: &Weak, + kind: DatanodeInspectKind, + schema: &SchemaRef, + request: ScanRequest, +) -> Result>> { + let projected_schema = projected_schema(schema, &request)?; + let info_ext = utils::information_extension(catalog_manager)?; + let req = DatanodeInspectRequest { + kind, + scan: request, + }; + + info_ext.inspect_datanode_plan(req, projected_schema) +} + /// Information schema table for sst manifest. pub struct InformationSchemaSstsManifest { schema: SchemaRef, @@ -63,11 +88,7 @@ impl InformationTable for InformationSchemaSstsManifest { } fn to_stream(&self, request: ScanRequest) -> Result { - let schema = if let Some(p) = request.projection_indices() { - Arc::new(self.schema.try_project(p).context(ProjectSchemaSnafu)?) - } else { - self.schema.clone() - }; + let schema = projected_schema(&self.schema, &request)?; let info_ext = utils::information_extension(&self.catalog_manager)?; let req = DatanodeInspectRequest { kind: DatanodeInspectKind::SstManifest, @@ -86,6 +107,15 @@ impl InformationTable for InformationSchemaSstsManifest { Box::pin(future), ))) } + + fn scan_plan(&self, request: ScanRequest) -> Result>> { + inspect_plan( + &self.catalog_manager, + DatanodeInspectKind::SstManifest, + &self.schema, + request, + ) + } } /// Information schema table for sst storage. @@ -117,11 +147,7 @@ impl InformationTable for InformationSchemaSstsStorage { } fn to_stream(&self, request: ScanRequest) -> Result { - let schema = if let Some(p) = request.projection_indices() { - Arc::new(self.schema.try_project(p).context(ProjectSchemaSnafu)?) - } else { - self.schema.clone() - }; + let schema = projected_schema(&self.schema, &request)?; let info_ext = utils::information_extension(&self.catalog_manager)?; let req = DatanodeInspectRequest { @@ -141,6 +167,15 @@ impl InformationTable for InformationSchemaSstsStorage { Box::pin(future), ))) } + + fn scan_plan(&self, request: ScanRequest) -> Result>> { + inspect_plan( + &self.catalog_manager, + DatanodeInspectKind::SstStorage, + &self.schema, + request, + ) + } } /// Information schema table for index metadata. @@ -172,11 +207,7 @@ impl InformationTable for InformationSchemaSstsIndexMeta { } fn to_stream(&self, request: ScanRequest) -> Result { - let schema = if let Some(p) = request.projection_indices() { - Arc::new(self.schema.try_project(p).context(ProjectSchemaSnafu)?) - } else { - self.schema.clone() - }; + let schema = projected_schema(&self.schema, &request)?; let info_ext = utils::information_extension(&self.catalog_manager)?; let req = DatanodeInspectRequest { @@ -196,4 +227,116 @@ impl InformationTable for InformationSchemaSstsIndexMeta { Box::pin(future), ))) } + + fn scan_plan(&self, request: ScanRequest) -> Result>> { + inspect_plan( + &self.catalog_manager, + DatanodeInspectKind::SstIndexMeta, + &self.schema, + request, + ) + } +} + +#[cfg(test)] +mod tests { + use cache::{build_fundamental_cache_registry, with_default_composite_cache_registry}; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, INFORMATION_SCHEMA_NAME}; + use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder}; + use common_meta::cluster::NodeInfo; + use common_meta::datanode::RegionStat; + use common_meta::key::flow::flow_state::FlowStat; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_procedure::ProcedureInfo; + use common_recordbatch::{RecordBatches, SendableRecordBatchStream}; + use datafusion::physical_plan::empty::EmptyExec; + + use super::*; + use crate::error::Error; + use crate::information_schema::InformationExtension; + use crate::kvbackend::KvBackendCatalogManagerBuilder; + + struct PlanInformationExtension; + + #[async_trait::async_trait] + impl InformationExtension for PlanInformationExtension { + type Error = Error; + + async fn nodes(&self) -> std::result::Result, Self::Error> { + Ok(vec![]) + } + + async fn procedures( + &self, + ) -> std::result::Result, Self::Error> { + Ok(vec![]) + } + + async fn region_stats(&self) -> std::result::Result, Self::Error> { + Ok(vec![]) + } + + async fn flow_stats(&self) -> std::result::Result, Self::Error> { + Ok(None) + } + + async fn inspect_datanode( + &self, + _request: DatanodeInspectRequest, + ) -> std::result::Result { + Ok(RecordBatches::empty().as_stream()) + } + + fn inspect_datanode_plan( + &self, + _request: DatanodeInspectRequest, + schema: SchemaRef, + ) -> std::result::Result>, Self::Error> { + Ok(Some(Arc::new(EmptyExec::new( + schema.arrow_schema().clone(), + )))) + } + } + + #[tokio::test] + async fn ssts_manifest_uses_information_extension_scan_plan() { + let backend = Arc::new(MemoryKvBackend::default()); + let layered_cache_builder = LayeredCacheRegistryBuilder::default() + .add_cache_registry(CacheRegistryBuilder::default().build()); + let fundamental_cache_registry = build_fundamental_cache_registry(backend.clone()); + let layered_cache_registry = Arc::new( + with_default_composite_cache_registry( + layered_cache_builder.add_cache_registry(fundamental_cache_registry), + ) + .unwrap() + .build(), + ); + + let catalog_manager = KvBackendCatalogManagerBuilder::new( + Arc::new(PlanInformationExtension), + backend, + layered_cache_registry, + ) + .build(); + + let table = catalog_manager + .table( + DEFAULT_CATALOG_NAME, + INFORMATION_SCHEMA_NAME, + SSTS_MANIFEST, + None, + ) + .await + .unwrap() + .unwrap(); + + let request = ScanRequest { + projection_input: Some(vec![0].into()), + ..Default::default() + }; + let plan = table.scan_to_plan(request).unwrap().unwrap(); + + assert!(plan.as_any().is::()); + assert_eq!(1, plan.schema().fields().len()); + } } diff --git a/src/cli/src/bench/metadata.rs b/src/cli/src/bench/metadata.rs index c21db072b1f..18c25c24ae9 100644 --- a/src/cli/src/bench/metadata.rs +++ b/src/cli/src/bench/metadata.rs @@ -119,6 +119,7 @@ impl TableMetadataBencher { &table_info.table_name(), &table_route, ®ion_wal_options, + None, ) .await; start.elapsed() diff --git a/src/cli/src/metadata/control/del/table.rs b/src/cli/src/metadata/control/del/table.rs index c92e5732248..59ebb2459ed 100644 --- a/src/cli/src/metadata/control/del/table.rs +++ b/src/cli/src/metadata/control/del/table.rs @@ -118,7 +118,13 @@ impl TableMetadataDeleter { .map_err(BoxedError::new)?; self.table_metadata_manager - .delete_table_metadata(table_id, &table_name, &table_route, ®ion_wal_options) + .delete_table_metadata( + table_id, + &table_name, + &table_route, + ®ion_wal_options, + None, + ) .await .map_err(BoxedError::new)?; Ok(()) diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index a8adfa01a48..b06da6e517a 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -12,10 +12,7 @@ path = "src/bin/greptime.rs" [[bin]] name = "query_perf_fixture" path = "src/bin/query_perf_fixture.rs" - -[[bin]] -name = "prom_remote_write_fixture" -path = "src/bin/prom_remote_write_fixture.rs" +required-features = ["dev-tools"] [features] default = [ @@ -25,7 +22,7 @@ default = [ "meta-srv/mysql_kvbackend", ] enterprise = ["common-meta/enterprise", "frontend/enterprise", "meta-srv/enterprise"] -# Developer-only tools (benchmarks, diagnostics) gating the `datanode parquetbench` subcommand. +# Developer-only helper binary gate for `query_perf_fixture`. # Kept out of `default` so normal/release builds don't compile them. dev-tools = [] tokio-console = ["common-telemetry/tokio-console"] diff --git a/src/cmd/src/bin/greptime.rs b/src/cmd/src/bin/greptime.rs index e0f326908a2..957e227355a 100644 --- a/src/cmd/src/bin/greptime.rs +++ b/src/cmd/src/bin/greptime.rs @@ -117,7 +117,6 @@ async fn start(cli: Command) -> Result<()> { } datanode::SubCommand::Objbench(ref bench) => bench.run().await, datanode::SubCommand::Scanbench(ref bench) => bench.run().await, - #[cfg(feature = "dev-tools")] datanode::SubCommand::Parquetbench(ref bench) => bench.run().await, }, SubCommand::Flownode(cmd) => { diff --git a/src/cmd/src/bin/prom_remote_write_fixture.rs b/src/cmd/src/bin/prom_remote_write_fixture.rs deleted file mode 100644 index 0bfdf954823..00000000000 --- a/src/cmd/src/bin/prom_remote_write_fixture.rs +++ /dev/null @@ -1,144 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#![allow(clippy::print_stdout)] - -use std::time::{Duration, Instant}; - -use api::prom_store::remote::{Label, Sample, TimeSeries, WriteRequest}; -use clap::Parser; -use prost::Message; -use serde_json::json; -use servers::prom_store::snappy_compress; - -#[derive(Debug, Parser)] -#[command(about = "Generate a deterministic Prometheus remote-write fixture")] -struct Args { - #[arg(long, default_value = "http://127.0.0.1:4000/v1/prometheus/write")] - endpoint: String, - #[arg(long, default_value = "public")] - database: String, - #[arg(long)] - metric: String, - #[arg(long, default_value = "greptime_physical_table")] - physical_table: String, - #[arg(long, default_value_t = 8)] - series_count: u64, - #[arg(long, default_value_t = 30)] - samples_per_series: u64, - #[arg(long, default_value_t = 1_704_067_200_000)] - start_unix_millis: i64, - #[arg(long, default_value_t = 15_000)] - step_millis: i64, - #[arg(long, alias = "batch-size", default_value_t = 8)] - chunk_series_count: u64, - #[arg(long, default_value_t = 60)] - timeout_seconds: u64, -} - -#[tokio::main] -async fn main() -> Result<(), Box> { - let args = Args::parse(); - let started = Instant::now(); - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(args.timeout_seconds)) - .build()?; - let mut batches = 0_u64; - let mut rows = 0_u64; - let mut http_statuses = Vec::new(); - - let chunk = args.chunk_series_count.max(1); - for first in (0..args.series_count).step_by(chunk as usize) { - let last = (first + chunk).min(args.series_count); - let timeseries = (first..last) - .map(|series_idx| make_series(&args, series_idx)) - .collect::>(); - rows += timeseries - .iter() - .map(|ts| ts.samples.len() as u64) - .sum::(); - let body = snappy_compress( - &WriteRequest { - timeseries, - ..Default::default() - } - .encode_to_vec(), - )?; - let response = client - .post(&args.endpoint) - .header("Content-Encoding", "snappy") - .body(body) - .send() - .await?; - let status = response.status(); - http_statuses.push(status.as_u16()); - if !status.is_success() { - let body_text = response.text().await.unwrap_or_default(); - return Err(format!("remote-write failed with status {status}: {body_text}").into()); - } - batches += 1; - } - - println!( - "{}", - json!({ - "status": "ok", - "endpoint": args.endpoint, - "database": args.database, - "metric": args.metric, - "physical_table": args.physical_table, - "series_count": args.series_count, - "samples_per_series": args.samples_per_series, - "rows": rows, - "samples_written": rows, - "batches": batches, - "elapsed_seconds": started.elapsed().as_secs_f64(), - "http_statuses": http_statuses, - }) - ); - Ok(()) -} - -fn make_series(args: &Args, series_idx: u64) -> TimeSeries { - let mut labels = vec![ - label("__name__", &args.metric), - label("x_greptime_database", &args.database), - label("x_greptime_physical_table", &args.physical_table), - label("host", &format!("host{:04}", series_idx % 1024)), - label("instance", &format!("instance{:06}", series_idx)), - ]; - labels.sort_unstable_by(|left, right| left.name.cmp(&right.name)); - let samples = (0..args.samples_per_series) - .map(|sample_idx| Sample { - value: deterministic_value(series_idx, sample_idx), - timestamp: args.start_unix_millis + (sample_idx as i64 * args.step_millis), - }) - .collect(); - TimeSeries { - labels, - samples, - ..Default::default() - } -} - -fn label(name: &str, value: &str) -> Label { - Label { - name: name.to_string(), - value: value.to_string(), - } -} - -fn deterministic_value(series_idx: u64, sample_idx: u64) -> f64 { - ((series_idx % 97) as f64) + (sample_idx as f64 * 0.125) -} diff --git a/src/cmd/src/bin/query_perf_fixture.rs b/src/cmd/src/bin/query_perf_fixture.rs index dfc3a664a03..d0e185038fa 100644 --- a/src/cmd/src/bin/query_perf_fixture.rs +++ b/src/cmd/src/bin/query_perf_fixture.rs @@ -14,656 +14,10 @@ #![allow(clippy::print_stderr, clippy::print_stdout)] -//! Generic direct-readable-SST fixture generator for query performance labs. -//! -//! This is a bounded MVP: it reads one table/one region case TOML, builds -//! synthetic region metadata, writes real Mito readable SST parquet files, and -//! emits a manifest checkpoint plus validation metadata. The first case uses -//! `sst_format = "flat"`; primary-key output is also accepted for early lab use -//! because the writer can convert the same flat input batch to primary-key SSTs. -//! -//! TODO(query-perf): add multi-region generation and seed-manifest/catalog DB -//! integration so the generated fixture can be mounted directly by a full -//! GreptimeDB open/query path without manual metadata wiring. - -use std::collections::HashMap; -use std::fs; -use std::io::{BufWriter, Write}; -use std::num::{NonZeroU64, NonZeroUsize}; -use std::path::{Path, PathBuf}; -use std::sync::Arc; - -use api::v1::{OpType, SemanticType}; -use async_trait::async_trait; -use clap::Parser; -use datatypes::arrow::array::{ - ArrayRef, BinaryDictionaryBuilder, Float64Array, RecordBatch, StringDictionaryBuilder, - TimestampMillisecondArray, TimestampNanosecondArray, UInt8Array, UInt64Array, -}; -use datatypes::arrow::datatypes::UInt32Type; -use datatypes::prelude::ConcreteDataType; -use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; -use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; -use mito2::access_layer::{FilePathProvider, Metrics, WriteType}; -use mito2::config::IndexConfig; -use mito2::manifest::action::{RegionCheckpoint, RegionManifest, RemovedFilesRecord}; -use mito2::read::FlatSource; -use mito2::sst::file::{FileMeta, RegionFileId}; -use mito2::sst::index::{Indexer, IndexerBuilder}; -use mito2::sst::parquet::writer::ParquetWriter; -use mito2::sst::parquet::{SstInfo, WriteOptions}; -use mito2::sst::{ - DEFAULT_WRITE_BUFFER_SIZE, FlatSchemaOptions, FormatType, to_flat_sst_arrow_schema, -}; -use object_store::ObjectStore; -use object_store::services::Fs as FsBuilder; -use serde::Deserialize; -use store_api::codec::PrimaryKeyEncoding; -use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataRef}; -use store_api::path_utils::region_name; -use store_api::region_request::PathType; -use store_api::storage::{FileId, RegionId}; - -#[derive(Parser, Debug)] -#[command(name = "query_perf_fixture")] -#[command(about = "Generate direct-readable SST fixtures for query performance cases")] -struct Args { - /// TOML case file, for example tests/perf/query_cases/.../case.toml. - #[arg(long, value_name = "PATH")] - case: PathBuf, - - /// Output directory (creates object-store/, manifest/, files.jsonl, summary.json). - #[arg(long, value_name = "DIR")] - out_dir: PathBuf, - - /// Region ID to synthesize. Defaults to table id 1024, region 0. - #[arg(long, default_value = "4398046511104")] - region_id: u64, - - /// Table directory relative to object-store root. Defaults to data/{database}/{table}/. - #[arg(long, value_name = "DIR")] - table_dir: Option, - - /// Table name to generate when the case contains multiple [[tables]]. - #[arg(long, value_name = "NAME")] - table: Option, - - /// Manifest/checkpoint version. - #[arg(long, default_value = "1000000")] - checkpoint_version: u64, - - /// Safety flag when sst_count exceeds 1000. - #[arg(long)] - allow_large: bool, - - /// Print plan only. - #[arg(long)] - dry_run: bool, -} - -#[derive(Debug, Deserialize)] -struct CaseFile { - scenario: Scenario, -} - -#[derive(Debug, Deserialize)] -#[serde(tag = "kind")] -enum Scenario { - #[serde(rename = "direct_readable_sst")] - DirectReadableSst(DirectReadableSstScenario), -} - -#[derive(Debug, Deserialize)] -struct DirectReadableSstScenario { - #[serde(default)] - seed: Option, - tables: Vec, - layout: LayoutConfig, -} - -impl Scenario { - fn kind(&self) -> &'static str { - match self { - Scenario::DirectReadableSst(_) => "direct_readable_sst", - } - } - - fn direct_readable_sst(&self) -> &DirectReadableSstScenario { - match self { - Scenario::DirectReadableSst(scenario) => scenario, - } - } -} - -#[derive(Debug, Deserialize)] -struct TableConfig { - database: String, - name: String, - engine: String, - #[serde(default)] - append_mode: Option, - #[serde(default)] - sst_format: Option, - primary_key: Vec, - time_index: String, - columns: Vec, -} - -#[derive(Debug, Deserialize)] -struct ColumnConfig { - name: String, - #[serde(rename = "type")] - ty: String, - semantic: String, - distribution: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(tag = "kind")] -enum Distribution { - #[serde(rename = "cardinality")] - Cardinality { - values: NonZeroUsize, - prefix: String, - }, - #[serde(rename = "deterministic_wave")] - DeterministicWave { min: f64, max: f64 }, -} - -#[derive(Debug, Deserialize)] -struct LayoutConfig { - regions: usize, - sst_count: usize, - rows_per_sst: usize, - row_group_size: usize, - series_count: NonZeroUsize, - start_unix_nanos: i64, - step_nanos: i64, - time_range_layout: String, - series_layout: String, -} - -struct NoopIndexBuilder; - -#[async_trait] -impl IndexerBuilder for NoopIndexBuilder { - async fn build(&self, _file_id: FileId, _index_version: u64) -> Indexer { - Indexer::default() - } -} - -#[derive(Clone)] -struct FixedPathProvider { - table_dir: String, -} - -impl FilePathProvider for FixedPathProvider { - fn build_index_file_path(&self, file_id: RegionFileId) -> String { - mito2::sst::location::index_file_path_legacy(&self.table_dir, file_id, PathType::Bare) - } - fn build_index_file_path_with_version( - &self, - index_id: mito2::sst::file::RegionIndexId, - ) -> String { - mito2::sst::location::index_file_path(&self.table_dir, index_id, PathType::Bare) - } - fn build_sst_file_path(&self, file_id: RegionFileId) -> String { - mito2::sst::location::sst_file_path(&self.table_dir, file_id, PathType::Bare) - } -} - -fn semantic_type(s: &str) -> SemanticType { - match s.to_ascii_lowercase().as_str() { - "tag" => SemanticType::Tag, - "field" => SemanticType::Field, - "timestamp" => SemanticType::Timestamp, - other => panic!("unsupported semantic {other}"), - } -} - -fn concrete_type(ty: &str) -> ConcreteDataType { - match ty.to_ascii_uppercase().as_str() { - "STRING" => ConcreteDataType::string_datatype(), - "DOUBLE" => ConcreteDataType::float64_datatype(), - "UINT64" => ConcreteDataType::uint64_datatype(), - "TIMESTAMP(9)" => ConcreteDataType::timestamp_nanosecond_datatype(), - "TIMESTAMP(3)" => ConcreteDataType::timestamp_millisecond_datatype(), - other => panic!("unsupported column type {other}"), - } -} - -fn build_region_metadata(table: &TableConfig, region_id: RegionId) -> RegionMetadata { - let mut builder = store_api::metadata::RegionMetadataBuilder::new(region_id); - for (idx, col) in table.columns.iter().enumerate() { - let mut schema = ColumnSchema::new( - &col.name, - concrete_type(&col.ty), - col.semantic != "timestamp", - ); - if col.semantic == "tag" { - schema = schema.with_inverted_index(true); - } - if col.semantic == "tag" || col.semantic == "field" { - schema = schema - .with_skipping_options(SkippingIndexOptions { - granularity: 1, - ..Default::default() - }) - .expect("valid skipping index options for query perf fixture"); - } - if col.name == table.time_index { - schema = schema.with_time_index(true); - } - builder.push_column_metadata(ColumnMetadata { - column_schema: schema, - semantic_type: semantic_type(&col.semantic), - column_id: idx as u32, - }); - } - let pk_ids = table - .primary_key - .iter() - .map(|name| { - table - .columns - .iter() - .position(|c| c.name == *name) - .unwrap_or_else(|| panic!("primary key column {name} not found")) as u32 - }) - .collect(); - builder.primary_key(pk_ids); - builder.primary_key_encoding(PrimaryKeyEncoding::Dense); - builder - .build() - .expect("region metadata should be valid for fixture table config") -} - -fn encode_dense_primary_key(table: &TableConfig, values: &HashMap) -> Vec { - let fields = table - .primary_key - .iter() - .enumerate() - .map(|(idx, _)| { - ( - idx as u32, - SortField::new(ConcreteDataType::string_datatype()), - ) - }) - .collect(); - let converter = DensePrimaryKeyCodec::with_fields(fields); - converter - .encode( - table - .primary_key - .iter() - .map(|name| datatypes::value::ValueRef::String(values[name].as_str())), - ) - .expect("dense primary key encoding should match fixture primary key fields") -} - -fn tag_value(col: &ColumnConfig, series: usize) -> String { - match col - .distribution - .as_ref() - .expect("tag distribution is required") - { - Distribution::Cardinality { values, prefix } => { - format!("{}{}", prefix, series % values.get()) - } - _ => panic!("tag column {} requires cardinality distribution", col.name), - } -} - -fn series_for_row(layout: &LayoutConfig, sst_idx: usize, base_row: usize, row: usize) -> usize { - match layout.series_layout.as_str() { - "round_robin" | "timestamp_major" => (base_row + row) % layout.series_count.get(), - "per_sst" => sst_idx % layout.series_count.get(), - other => panic!("unsupported series_layout {other}"), - } -} - -fn timestamp_for_row(layout: &LayoutConfig, base_row: usize, row: usize) -> i64 { - let logical_row = base_row + row; - let timestamp_step = if layout.series_layout == "timestamp_major" { - logical_row / layout.series_count.get() - } else { - logical_row - }; - - layout.start_unix_nanos + (timestamp_step as i64 * layout.step_nanos) -} - -fn wave_value(min: f64, max: f64, row: usize) -> f64 { - let span = max - min; - let phase = (row % 1024) as f64 / 1023.0; - min + span * (0.5 - 0.5 * (std::f64::consts::TAU * phase).cos()) -} - -fn generate_record_batch( - table: &TableConfig, - metadata: &RegionMetadataRef, - layout: &LayoutConfig, - sst_idx: usize, - sequence: u64, -) -> RecordBatch { - let flat_schema = to_flat_sst_arrow_schema(metadata, &FlatSchemaOptions::default()); - let rows = layout.rows_per_sst; - let mut columns = Vec::new(); - let base_row = sst_idx * rows; - - for col in &table.columns { - match (col.semantic.as_str(), col.ty.to_ascii_uppercase().as_str()) { - ("tag", "STRING") => { - let mut b = StringDictionaryBuilder::::new(); - for row in 0..rows { - let series = series_for_row(layout, sst_idx, base_row, row); - let value = tag_value(col, series); - b.append_value(value); - } - columns.push(Arc::new(b.finish()) as ArrayRef); - } - ("field", "DOUBLE") => { - let (min, max) = match col.distribution.as_ref() { - Some(Distribution::DeterministicWave { min, max }) => (*min, *max), - _ => (0.0, 1.0), - }; - columns.push(Arc::new(Float64Array::from( - (0..rows) - .map(|r| wave_value(min, max, base_row + r)) - .collect::>(), - )) as ArrayRef); - } - ("field", "UINT64") => columns.push(Arc::new(UInt64Array::from( - (0..rows).map(|r| (base_row + r) as u64).collect::>(), - )) as ArrayRef), - ("timestamp", "TIMESTAMP(9)") => columns.push(Arc::new(TimestampNanosecondArray::from( - (0..rows) - .map(|r| timestamp_for_row(layout, base_row, r)) - .collect::>(), - )) as ArrayRef), - ("timestamp", "TIMESTAMP(3)") => { - columns.push(Arc::new(TimestampMillisecondArray::from( - (0..rows) - .map(|r| timestamp_for_row(layout, base_row, r) / 1_000_000) - .collect::>(), - )) as ArrayRef) - } - other => panic!("unsupported column combination {other:?}"), - } - } - - let mut pk_builder = BinaryDictionaryBuilder::::new(); - for row in 0..rows { - let series = series_for_row(layout, sst_idx, base_row, row); - let values = table - .columns - .iter() - .filter(|c| c.semantic == "tag") - .map(|c| (c.name.clone(), tag_value(c, series))) - .collect(); - pk_builder - .append(encode_dense_primary_key(table, &values)) - .expect("append generated dense primary key to Arrow dictionary builder"); - } - columns.push(Arc::new(pk_builder.finish())); - columns.push(Arc::new(UInt64Array::from_value(sequence, rows))); - columns.push(Arc::new(UInt8Array::from_value(OpType::Put as u8, rows))); - RecordBatch::try_new(flat_schema, columns) - .expect("generated fixture columns should match flat SST Arrow schema") -} - -fn file_meta_from_sst_info( - info: &SstInfo, - region_id: RegionId, - file_id: FileId, - sequence: u64, -) -> FileMeta { - FileMeta { - region_id, - file_id, - time_range: info.time_range, - level: 0, - file_size: info.file_size, - max_row_group_uncompressed_size: info.max_row_group_uncompressed_size, - available_indexes: Default::default(), - indexes: Default::default(), - index_file_size: 0, - index_version: 0, - num_rows: info.num_rows as u64, - num_row_groups: info.num_row_groups, - sequence: NonZeroU64::new(sequence), - partition_expr: None, - num_series: info.num_series, - primary_key_min: None, - primary_key_max: None, - } -} - -fn deterministic_file_id(seed: u64, index: usize) -> FileId { - let seed_hi = ((seed >> 16) & 0xffff) as u16; - let seed_lo = (seed & 0xffff) as u16; - FileId::parse_str(&format!( - "00000000-0000-{seed_hi:04x}-{seed_lo:04x}-{index:012x}" - )) - .expect("deterministic file id must be a valid UUID") -} - -fn remap_sst_file( - object_store_dir: &Path, - table_dir: &str, - region_id: RegionId, - old_file_id: FileId, - new_file_id: FileId, -) { - if old_file_id == new_file_id { - return; - } - let old_path = object_store_dir.join(mito2::sst::location::sst_file_path( - table_dir, - RegionFileId::new(region_id, old_file_id), - PathType::Bare, - )); - let new_path = object_store_dir.join(mito2::sst::location::sst_file_path( - table_dir, - RegionFileId::new(region_id, new_file_id), - PathType::Bare, - )); - if let Some(parent) = new_path.parent() { - fs::create_dir_all(parent).expect("failed to create remapped SST parent directory"); - } - fs::rename(&old_path, &new_path).unwrap_or_else(|err| { - panic!( - "failed to remap SST file {} -> {}: {err}", - old_path.display(), - new_path.display() - ) - }); -} - -fn case_name_from_path(path: &Path) -> String { - path.parent() - .and_then(Path::file_name) - .and_then(|name| name.to_str()) - .unwrap_or("query_perf_case") - .to_string() -} +#[path = "query_perf_fixture/mod.rs"] +mod query_perf_fixture; #[tokio::main] async fn main() { - let args = Args::parse(); - let case_text = fs::read_to_string(&args.case).expect("failed to read query perf case TOML"); - let case: CaseFile = toml::from_str(&case_text).expect("failed to parse query perf case TOML"); - let case_name = case_name_from_path(&args.case); - let scenario = case.scenario.direct_readable_sst(); - if scenario.tables.is_empty() || scenario.layout.regions != 1 { - panic!("fixture generator supports one or more tables and exactly one region per table"); - } - if scenario.layout.time_range_layout != "non_overlapping_per_sst" { - panic!("MVP supports time_range_layout=non_overlapping_per_sst"); - } - if scenario.layout.series_layout == "timestamp_major" - && !scenario - .layout - .rows_per_sst - .is_multiple_of(scenario.layout.series_count.get()) - { - panic!( - "series_layout=timestamp_major requires rows_per_sst to be divisible by series_count" - ); - } - if scenario.layout.sst_count > 1000 && !args.allow_large { - panic!("sst_count exceeds 1000; pass --allow-large"); - } - let seed = scenario.seed.unwrap_or(0); - let table_index = match args.table.as_deref() { - Some(name) => scenario - .tables - .iter() - .position(|table| table.name == name) - .unwrap_or_else(|| panic!("--table {name} was not found in case {case_name}")), - None if scenario.tables.len() == 1 => 0, - None => panic!( - "case {} contains {} tables; pass --table to choose one", - case_name, - scenario.tables.len() - ), - }; - let table = &scenario.tables[table_index]; - if table.engine != "mito" { - panic!("MVP supports engine=mito"); - } - let region_id = RegionId::from(args.region_id); - let table_dir = args - .table_dir - .unwrap_or_else(|| format!("data/{}/{}/", table.database, table.name)); - if Path::new(&table_dir).is_absolute() { - panic!("--table-dir must be relative"); - } - if table_dir.trim_end_matches('/').rsplit('/').next() - == Some(region_name(region_id.table_id(), region_id.region_sequence()).as_str()) - { - panic!("--table-dir must be table-level directory"); - } - let region_dir = - mito2::sst::location::region_dir_from_table_dir(&table_dir, region_id, PathType::Bare); - let format = match table.sst_format.as_deref().unwrap_or("flat") { - "flat" => FormatType::Flat, - "primary_key" | "primary-key" => FormatType::PrimaryKey, - other => panic!("unsupported sst_format {other}"), - }; - println!( - "query_perf_fixture case={} scenario={} table={}.{} ssts={} rows_per_sst={} format={format:?}", - case_name, - case.scenario.kind(), - table.database, - table.name, - scenario.layout.sst_count, - scenario.layout.rows_per_sst - ); - println!( - "out_dir={} table_dir={} region_dir={}", - args.out_dir.display(), - table_dir, - region_dir - ); - if args.dry_run { - return; - } - - let obj_store_dir = args.out_dir.join("object-store"); - let manifest_dir = args.out_dir.join("manifest"); - fs::create_dir_all(&obj_store_dir).expect("failed to create fixture object-store directory"); - fs::create_dir_all(&manifest_dir).expect("failed to create fixture manifest directory"); - let ostorage = ObjectStore::new(FsBuilder::default().root(&obj_store_dir.to_string_lossy())) - .expect("failed to create filesystem object store for fixture output") - .finish(); - let metadata: RegionMetadataRef = Arc::new(build_region_metadata(table, region_id)); - let mut files = HashMap::with_capacity(scenario.layout.sst_count); - let mut next_file_index = 1; - for i in 0..scenario.layout.sst_count { - let sequence = 1000 + i as u64; - let batch = generate_record_batch(table, &metadata, &scenario.layout, i, sequence); - let source = - FlatSource::new_iter(batch.schema(), Box::new(vec![batch].into_iter().map(Ok))); - let mut metrics = Metrics::new(WriteType::Flush); - let mut writer = ParquetWriter::new_with_object_store( - ostorage.clone(), - metadata.clone(), - IndexConfig::default(), - NoopIndexBuilder, - FixedPathProvider { - table_dir: table_dir.clone(), - }, - &mut metrics, - ) - .await; - let opts = WriteOptions { - write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, - row_group_size: scenario.layout.row_group_size, - max_file_size: None, - }; - let infos = match format { - FormatType::Flat => writer.write_all_flat(source, Some(sequence), &opts).await, - FormatType::PrimaryKey => { - writer - .write_all_flat_as_primary_key(source, Some(sequence), &opts) - .await - } - } - .expect("failed to write fixture SST through Mito parquet writer"); - for info in infos { - let file_id = deterministic_file_id( - seed.wrapping_add((table_index as u64) << 16), - next_file_index, - ); - next_file_index += 1; - remap_sst_file(&obj_store_dir, &table_dir, region_id, info.file_id, file_id); - files.insert( - file_id, - file_meta_from_sst_info(&info, region_id, file_id, sequence), - ); - } - } - let overall_seq = files - .values() - .filter_map(|m| m.sequence) - .map(|s| s.get()) - .max() - .unwrap_or(0); - let manifest = RegionManifest { - metadata, - files, - removed_files: RemovedFilesRecord::default(), - flushed_entry_id: overall_seq, - flushed_sequence: overall_seq, - committed_sequence: Some(overall_seq), - manifest_version: args.checkpoint_version, - truncated_entry_id: None, - compaction_time_window: None, - sst_format: format, - append_mode: table.append_mode, - }; - let checkpoint = RegionCheckpoint { - last_version: args.checkpoint_version, - compacted_actions: manifest.files.len(), - checkpoint: Some(manifest.clone()), - }; - let checkpoint_bytes = checkpoint - .encode() - .expect("failed to encode fixture region checkpoint"); - let checkpoint_path = manifest_dir.join(format!("{:020}.checkpoint", args.checkpoint_version)); - fs::write(&checkpoint_path, &checkpoint_bytes).expect("failed to write checkpoint file"); - fs::write(manifest_dir.join("_last_checkpoint"), serde_json::to_vec_pretty(&serde_json::json!({ "size": checkpoint_bytes.len(), "version": args.checkpoint_version, "checksum": null, "extend_metadata": {} })).expect("failed to serialize _last_checkpoint metadata")).expect("failed to write _last_checkpoint file"); - let files_jsonl_path = args.out_dir.join("files.jsonl"); - let mut jsonl = BufWriter::new( - fs::File::create(&files_jsonl_path).expect("failed to create fixture files.jsonl"), - ); - let mut entries: Vec<_> = manifest.files.iter().collect(); - entries.sort_by_key(|(id, _)| id.to_string()); - for (file_id, meta) in entries { - writeln!(jsonl, "{}", serde_json::to_string(&serde_json::json!({ "file_id": file_id.to_string(), "region_id": meta.region_id.as_u64(), "object_path": mito2::sst::location::sst_file_path(&table_dir, RegionFileId::new(meta.region_id, *file_id), PathType::Bare), "time_range_start": meta.time_range.0.value(), "time_range_end": meta.time_range.1.value(), "num_rows": meta.num_rows, "num_row_groups": meta.num_row_groups, "file_size": meta.file_size, "num_series": meta.num_series, "sequence": meta.sequence.map(|s| s.get()) })).expect("failed to serialize fixture SST metadata JSONL entry")).expect("failed to write fixture SST metadata JSONL entry"); - } - jsonl.flush().expect("failed to flush fixture files.jsonl"); - fs::write(args.out_dir.join("summary.json"), serde_json::to_vec_pretty(&serde_json::json!({ "case": case_name, "seed": seed, "table_index": table_index, "table": table.name, "database": table.database, "region_id": region_id.as_u64(), "table_dir": table_dir, "region_dir": region_dir, "sst_format": format!("{format:?}"), "sst_count": scenario.layout.sst_count, "rows_per_sst": scenario.layout.rows_per_sst, "row_group_size": scenario.layout.row_group_size, "total_rows": scenario.layout.sst_count * scenario.layout.rows_per_sst, "checkpoint_path": checkpoint_path, "files_jsonl_path": files_jsonl_path, "readback_validated": false, "metadata_source": "synthetic" })).expect("failed to serialize fixture summary")).expect("failed to write fixture summary.json"); - println!("Done. wrote {} SST file entries", manifest.files.len()); + query_perf_fixture::run().await; } diff --git a/src/cmd/src/bin/query_perf_fixture/case.rs b/src/cmd/src/bin/query_perf_fixture/case.rs new file mode 100644 index 00000000000..7dce735d7b9 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/case.rs @@ -0,0 +1,470 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::num::NonZeroUsize; + +use clap::ValueEnum; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct CaseFile { + pub(super) scenario: Scenario, +} + +#[derive(Debug, Deserialize, Serialize)] +#[serde(tag = "kind")] +pub(super) enum Scenario { + #[serde(rename = "direct_readable_sst")] + DirectReadableSst(DirectReadableSstScenario), + #[serde(rename = "prom_remote_write_then_query")] + PromRemoteWriteThenQuery(PromRemoteWriteThenQueryScenario), +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct DirectReadableSstScenario { + #[serde(default)] + pub(super) seed: Option, + pub(super) tables: Vec, + pub(super) layout: LayoutConfig, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct PromRemoteWriteThenQueryScenario { + #[serde(default)] + pub(super) queries: Vec, + pub(super) remote_write: PromRemoteWritePlan, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct PromRemoteWritePlan { + #[serde(default = "default_database")] + pub(super) database: String, + #[serde(alias = "metric_name")] + pub(super) metric: String, + #[serde(default = "default_physical_table")] + pub(super) physical_table: String, + #[serde(default = "default_series_count")] + pub(super) series_count: u64, + #[serde(default = "default_samples_per_series")] + pub(super) samples_per_series: u64, + #[serde(default = "default_start_unix_millis")] + pub(super) start_unix_millis: i64, + #[serde(default = "default_step_millis")] + pub(super) step_millis: i64, + #[serde(default = "default_chunk_series_count", alias = "batch_size")] + pub(super) chunk_series_count: u64, + #[serde(default = "default_timeout_seconds")] + pub(super) timeout_seconds: u64, + #[serde(default)] + pub(super) sample_chunk_size: Option, + #[serde(default = "default_flush_every_sample_chunks")] + pub(super) flush_every_sample_chunks: u64, + #[serde(default = "default_visibility_timeout_seconds")] + pub(super) visibility_timeout_seconds: u64, + #[serde(default)] + pub(super) prom_store: PromStoreConfig, + #[serde(default)] + pub(super) value: ValueConfig, + #[serde(default)] + pub(super) storage: Option, + #[serde(default)] + pub(super) read_bench: Option, +} + +pub(super) fn default_database() -> String { + "public".to_string() +} +pub(super) fn default_physical_table() -> String { + "greptime_physical_table".to_string() +} +pub(super) fn default_series_count() -> u64 { + 8 +} +pub(super) fn default_samples_per_series() -> u64 { + 30 +} +pub(super) fn default_start_unix_millis() -> i64 { + 1_704_067_200_000 +} +pub(super) fn default_step_millis() -> i64 { + 15_000 +} +pub(super) fn default_chunk_series_count() -> u64 { + 8 +} +pub(super) fn default_timeout_seconds() -> u64 { + 60 +} +pub(super) fn default_flush_every_sample_chunks() -> u64 { + 1 +} +pub(super) fn default_visibility_timeout_seconds() -> u64 { + 30 +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct PromStoreConfig { + #[serde(default = "default_pending_rows_flush_interval")] + pub(super) pending_rows_flush_interval: String, + #[serde(default = "default_max_batch_rows")] + pub(super) max_batch_rows: u64, + #[serde(default = "default_max_concurrent_flushes")] + pub(super) max_concurrent_flushes: u64, + #[serde(default = "default_worker_channel_capacity")] + pub(super) worker_channel_capacity: u64, + #[serde(default = "default_max_inflight_requests")] + pub(super) max_inflight_requests: u64, +} + +impl Default for PromStoreConfig { + fn default() -> Self { + Self { + pending_rows_flush_interval: default_pending_rows_flush_interval(), + max_batch_rows: default_max_batch_rows(), + max_concurrent_flushes: default_max_concurrent_flushes(), + worker_channel_capacity: default_worker_channel_capacity(), + max_inflight_requests: default_max_inflight_requests(), + } + } +} + +pub(super) fn default_pending_rows_flush_interval() -> String { + "1s".to_string() +} +pub(super) fn default_max_batch_rows() -> u64 { + 100000 +} +pub(super) fn default_max_concurrent_flushes() -> u64 { + 256 +} +pub(super) fn default_worker_channel_capacity() -> u64 { + 65526 +} +pub(super) fn default_max_inflight_requests() -> u64 { + 3000 +} + +#[derive(Debug, Clone, Copy, Deserialize, Serialize, ValueEnum, Eq, PartialEq)] +#[serde(rename_all = "snake_case")] +#[value(rename_all = "snake_case")] +pub(super) enum ValuePattern { + Linear, + Constant, + Modulo, + Unique, + SeededRandom, + RunLength, + QuantizedSignal, + SignalWithSporadicStalls, + MixedSignalRepeated, +} + +impl std::fmt::Display for ValuePattern { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{}", + serde_json::to_value(self).unwrap().as_str().unwrap() + ) + } +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct ValueConfig { + #[serde(default = "default_value_pattern")] + pub(super) pattern: ValuePattern, + #[serde(default)] + pub(super) base: f64, + #[serde(default = "default_value_step")] + pub(super) step: f64, + #[serde(default = "default_value_cardinality")] + pub(super) cardinality: u64, + #[serde(default)] + pub(super) seed: u64, + #[serde(default = "default_value_run_length")] + pub(super) run_length: u64, + #[serde(default = "default_value_stall_every")] + pub(super) stall_every: u64, + #[serde(default = "default_value_stall_length")] + pub(super) stall_length: u64, + #[serde(default = "default_value_mixed_every")] + pub(super) mixed_every: u64, +} +impl Default for ValueConfig { + fn default() -> Self { + Self { + pattern: default_value_pattern(), + base: 0.0, + step: default_value_step(), + cardinality: default_value_cardinality(), + seed: 0, + run_length: default_value_run_length(), + stall_every: default_value_stall_every(), + stall_length: default_value_stall_length(), + mixed_every: default_value_mixed_every(), + } + } +} +pub(super) fn default_value_pattern() -> ValuePattern { + ValuePattern::Linear +} +pub(super) fn default_value_step() -> f64 { + 0.125 +} +pub(super) fn default_value_cardinality() -> u64 { + 97 +} +pub(super) fn default_value_run_length() -> u64 { + 8 +} +pub(super) fn default_value_stall_every() -> u64 { + 100 +} +pub(super) fn default_value_stall_length() -> u64 { + 16 +} +pub(super) fn default_value_mixed_every() -> u64 { + 5 +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct StorageConfig { + #[serde(default = "default_true")] + pub(super) inspect: bool, + #[serde(default = "default_storage_column")] + pub(super) column: String, + #[serde(default)] + pub(super) root_suffix: Option, + #[serde(default)] + pub(super) include_metadata_files: bool, + #[serde(default = "default_min_files")] + pub(super) min_files: u64, + #[serde(default = "default_min_files")] + pub(super) min_files_with_column: u64, + #[serde(default)] + pub(super) require_encodings: Vec, + #[serde(default)] + pub(super) forbid_encodings: Vec, + #[serde(default)] + pub(super) max_total_file_size_bytes: Option, + #[serde(default)] + pub(super) max_column_compressed_size_bytes: Option, + #[serde(default)] + pub(super) max_column_uncompressed_size_bytes: Option, + #[serde(default)] + pub(super) max_candidate_total_file_size_regression_pct: Option, + #[serde(default)] + pub(super) max_candidate_column_compressed_size_regression_pct: Option, + #[serde(default)] + pub(super) max_candidate_column_uncompressed_size_regression_pct: Option, + #[serde(skip_deserializing, default)] + pub(super) planned_thresholds: Vec, +} + +#[derive(Debug, Serialize)] +pub(super) struct StorageThresholdPlan { + pub(super) threshold: String, + pub(super) status: String, + pub(super) value: serde_json::Value, +} + +impl StorageConfig { + pub(super) fn populate_planned_thresholds(&mut self) { + let mut planned = Vec::new(); + planned.push(StorageThresholdPlan::new("min_files", self.min_files)); + planned.push(StorageThresholdPlan::new( + "min_files_with_column", + self.min_files_with_column, + )); + if !self.require_encodings.is_empty() { + planned.push(StorageThresholdPlan::new( + "require_encodings", + &self.require_encodings, + )); + } + if !self.forbid_encodings.is_empty() { + planned.push(StorageThresholdPlan::new( + "forbid_encodings", + &self.forbid_encodings, + )); + } + macro_rules! push_optional { + ($name:literal, $value:expr) => { + if let Some(value) = $value { + planned.push(StorageThresholdPlan::new($name, value)); + } + }; + } + push_optional!("max_total_file_size_bytes", self.max_total_file_size_bytes); + push_optional!( + "max_column_compressed_size_bytes", + self.max_column_compressed_size_bytes + ); + push_optional!( + "max_column_uncompressed_size_bytes", + self.max_column_uncompressed_size_bytes + ); + push_optional!( + "max_candidate_total_file_size_regression_pct", + self.max_candidate_total_file_size_regression_pct + ); + push_optional!( + "max_candidate_column_compressed_size_regression_pct", + self.max_candidate_column_compressed_size_regression_pct + ); + push_optional!( + "max_candidate_column_uncompressed_size_regression_pct", + self.max_candidate_column_uncompressed_size_regression_pct + ); + self.planned_thresholds = planned; + } +} + +impl StorageThresholdPlan { + fn new(threshold: &str, value: T) -> Self { + Self { + threshold: threshold.to_string(), + status: "planned".to_string(), + value: serde_json::to_value(value).expect("storage threshold value must serialize"), + } + } +} +pub(super) fn default_true() -> bool { + true +} +pub(super) fn default_storage_column() -> String { + "greptime_value".to_string() +} +pub(super) fn default_min_files() -> u64 { + 1 +} + +#[derive(Debug, Deserialize, Serialize, Clone)] +pub(super) struct ReadBenchConfig { + #[serde(default = "default_true")] + pub(super) enabled: bool, + #[serde(default = "default_true")] + pub(super) parquetbench: bool, + #[serde(default = "default_true")] + pub(super) scanbench: bool, + #[serde(default = "default_iterations")] + pub(super) iterations: u64, + #[serde(default)] + pub(super) projection: Vec, + #[serde(default = "default_parquet_reader")] + pub(super) parquet_reader: String, + #[serde(default = "default_scan_scanner")] + pub(super) scan_scanner: String, + #[serde(default = "default_parallelism")] + pub(super) parallelism: u64, + #[serde(default)] + pub(super) max_files: Option, + #[serde(flatten)] + pub(super) thresholds: HashMap, +} + +impl Default for ReadBenchConfig { + fn default() -> Self { + Self { + enabled: true, + parquetbench: true, + scanbench: true, + iterations: default_iterations(), + projection: vec![], + parquet_reader: default_parquet_reader(), + scan_scanner: default_scan_scanner(), + parallelism: default_parallelism(), + max_files: None, + thresholds: HashMap::new(), + } + } +} + +pub(super) fn default_iterations() -> u64 { + 7 +} +pub(super) fn default_parquet_reader() -> String { + "direct".to_string() +} +pub(super) fn default_scan_scanner() -> String { + "seq".to_string() +} +pub(super) fn default_parallelism() -> u64 { + 1 +} + +impl Scenario { + pub(super) fn kind(&self) -> &'static str { + match self { + Scenario::DirectReadableSst(_) => "direct_readable_sst", + Scenario::PromRemoteWriteThenQuery(_) => "prom_remote_write_then_query", + } + } + + pub(super) fn direct_readable_sst(&self) -> &DirectReadableSstScenario { + match self { + Scenario::DirectReadableSst(scenario) => scenario, + _ => panic!("scenario is not direct_readable_sst"), + } + } +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct TableConfig { + pub(super) database: String, + pub(super) name: String, + pub(super) engine: String, + #[serde(default)] + pub(super) append_mode: Option, + #[serde(default)] + pub(super) sst_format: Option, + pub(super) primary_key: Vec, + pub(super) time_index: String, + pub(super) columns: Vec, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct ColumnConfig { + pub(super) name: String, + #[serde(rename = "type")] + pub(super) ty: String, + pub(super) semantic: String, + pub(super) distribution: Option, +} + +#[derive(Debug, Deserialize, Serialize)] +#[serde(tag = "kind")] +pub(super) enum Distribution { + #[serde(rename = "cardinality")] + Cardinality { + values: NonZeroUsize, + prefix: String, + }, + #[serde(rename = "deterministic_wave")] + DeterministicWave { min: f64, max: f64 }, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(super) struct LayoutConfig { + pub(super) regions: usize, + pub(super) sst_count: usize, + pub(super) rows_per_sst: usize, + pub(super) row_group_size: usize, + pub(super) series_count: NonZeroUsize, + pub(super) start_unix_nanos: i64, + pub(super) step_nanos: i64, + pub(super) time_range_layout: String, + pub(super) series_layout: String, +} diff --git a/src/cmd/src/bin/query_perf_fixture/direct_sst.rs b/src/cmd/src/bin/query_perf_fixture/direct_sst.rs new file mode 100644 index 00000000000..0f2e3b34876 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/direct_sst.rs @@ -0,0 +1,529 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::fs::{self}; +use std::io::{BufWriter, Write}; +use std::num::NonZeroU64; +use std::path::Path; +use std::sync::Arc; + +use api::v1::{OpType, SemanticType}; +use async_trait::async_trait; +use datatypes::arrow::array::{ + ArrayRef, BinaryDictionaryBuilder, Float64Array, RecordBatch, StringDictionaryBuilder, + TimestampMillisecondArray, TimestampNanosecondArray, UInt8Array, UInt64Array, +}; +use datatypes::arrow::datatypes::UInt32Type; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; +use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; +use mito2::access_layer::{FilePathProvider, Metrics, WriteType}; +use mito2::config::IndexConfig; +use mito2::manifest::action::{RegionCheckpoint, RegionManifest, RemovedFilesRecord}; +use mito2::read::FlatSource; +use mito2::sst::file::{FileMeta, RegionFileId}; +use mito2::sst::index::{Indexer, IndexerBuilder}; +use mito2::sst::parquet::writer::ParquetWriter; +use mito2::sst::parquet::{SstInfo, WriteOptions}; +use mito2::sst::{ + DEFAULT_WRITE_BUFFER_SIZE, FlatSchemaOptions, FormatType, to_flat_sst_arrow_schema, +}; +use object_store::ObjectStore; +use object_store::services::Fs as FsBuilder; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataRef}; +use store_api::path_utils::region_name; +use store_api::region_request::PathType; +use store_api::storage::{FileId, RegionId}; + +use crate::query_perf_fixture::DirectArgs; +use crate::query_perf_fixture::case::*; +use crate::query_perf_fixture::util::case_name_from_path; + +struct NoopIndexBuilder; + +#[async_trait] +impl IndexerBuilder for NoopIndexBuilder { + async fn build(&self, _file_id: FileId, _index_version: u64) -> Indexer { + Indexer::default() + } +} + +#[derive(Clone)] +struct FixedPathProvider { + table_dir: String, +} + +impl FilePathProvider for FixedPathProvider { + fn build_index_file_path(&self, file_id: RegionFileId) -> String { + mito2::sst::location::index_file_path_legacy(&self.table_dir, file_id, PathType::Bare) + } + fn build_index_file_path_with_version( + &self, + index_id: mito2::sst::file::RegionIndexId, + ) -> String { + mito2::sst::location::index_file_path(&self.table_dir, index_id, PathType::Bare) + } + fn build_sst_file_path(&self, file_id: RegionFileId) -> String { + mito2::sst::location::sst_file_path(&self.table_dir, file_id, PathType::Bare) + } +} + +fn semantic_type(s: &str) -> SemanticType { + match s.to_ascii_lowercase().as_str() { + "tag" => SemanticType::Tag, + "field" => SemanticType::Field, + "timestamp" => SemanticType::Timestamp, + other => panic!("unsupported semantic {other}"), + } +} + +fn concrete_type(ty: &str) -> ConcreteDataType { + match ty.to_ascii_uppercase().as_str() { + "STRING" => ConcreteDataType::string_datatype(), + "DOUBLE" => ConcreteDataType::float64_datatype(), + "UINT64" => ConcreteDataType::uint64_datatype(), + "TIMESTAMP(9)" => ConcreteDataType::timestamp_nanosecond_datatype(), + "TIMESTAMP(3)" => ConcreteDataType::timestamp_millisecond_datatype(), + other => panic!("unsupported column type {other}"), + } +} + +fn build_region_metadata(table: &TableConfig, region_id: RegionId) -> RegionMetadata { + let mut builder = store_api::metadata::RegionMetadataBuilder::new(region_id); + for (idx, col) in table.columns.iter().enumerate() { + let mut schema = ColumnSchema::new( + &col.name, + concrete_type(&col.ty), + col.semantic != "timestamp", + ); + if col.semantic == "tag" { + schema = schema.with_inverted_index(true); + } + if col.semantic == "tag" || col.semantic == "field" { + schema = schema + .with_skipping_options(SkippingIndexOptions { + granularity: 1, + ..Default::default() + }) + .expect("valid skipping index options for query perf fixture"); + } + if col.name == table.time_index { + schema = schema.with_time_index(true); + } + builder.push_column_metadata(ColumnMetadata { + column_schema: schema, + semantic_type: semantic_type(&col.semantic), + column_id: idx as u32, + }); + } + let pk_ids = table + .primary_key + .iter() + .map(|name| { + table + .columns + .iter() + .position(|c| c.name == *name) + .unwrap_or_else(|| panic!("primary key column {name} not found")) as u32 + }) + .collect(); + builder.primary_key(pk_ids); + builder.primary_key_encoding(PrimaryKeyEncoding::Dense); + builder + .build() + .expect("region metadata should be valid for fixture table config") +} + +fn encode_dense_primary_key(table: &TableConfig, values: &HashMap) -> Vec { + let fields = table + .primary_key + .iter() + .enumerate() + .map(|(idx, _)| { + ( + idx as u32, + SortField::new(ConcreteDataType::string_datatype()), + ) + }) + .collect(); + let converter = DensePrimaryKeyCodec::with_fields(fields); + converter + .encode( + table + .primary_key + .iter() + .map(|name| datatypes::value::ValueRef::String(values[name].as_str())), + ) + .expect("dense primary key encoding should match fixture primary key fields") +} + +fn tag_value(col: &ColumnConfig, series: usize) -> String { + match col + .distribution + .as_ref() + .expect("tag distribution is required") + { + Distribution::Cardinality { values, prefix } => { + format!("{}{}", prefix, series % values.get()) + } + _ => panic!("tag column {} requires cardinality distribution", col.name), + } +} + +fn series_for_row(layout: &LayoutConfig, sst_idx: usize, base_row: usize, row: usize) -> usize { + match layout.series_layout.as_str() { + "round_robin" | "timestamp_major" => (base_row + row) % layout.series_count.get(), + "per_sst" => sst_idx % layout.series_count.get(), + other => panic!("unsupported series_layout {other}"), + } +} + +fn timestamp_for_row(layout: &LayoutConfig, base_row: usize, row: usize) -> i64 { + let logical_row = base_row + row; + let timestamp_step = if layout.series_layout == "timestamp_major" { + logical_row / layout.series_count.get() + } else { + logical_row + }; + + layout.start_unix_nanos + (timestamp_step as i64 * layout.step_nanos) +} + +fn wave_value(min: f64, max: f64, row: usize) -> f64 { + let span = max - min; + let phase = (row % 1024) as f64 / 1023.0; + min + span * (0.5 - 0.5 * (std::f64::consts::TAU * phase).cos()) +} + +fn generate_record_batch( + table: &TableConfig, + metadata: &RegionMetadataRef, + layout: &LayoutConfig, + sst_idx: usize, + sequence: u64, +) -> RecordBatch { + let flat_schema = to_flat_sst_arrow_schema(metadata, &FlatSchemaOptions::default()); + let rows = layout.rows_per_sst; + let mut columns = Vec::new(); + let base_row = sst_idx * rows; + + for col in &table.columns { + match (col.semantic.as_str(), col.ty.to_ascii_uppercase().as_str()) { + ("tag", "STRING") => { + let mut b = StringDictionaryBuilder::::new(); + for row in 0..rows { + let series = series_for_row(layout, sst_idx, base_row, row); + let value = tag_value(col, series); + b.append_value(value); + } + columns.push(Arc::new(b.finish()) as ArrayRef); + } + ("field", "DOUBLE") => { + let (min, max) = match col.distribution.as_ref() { + Some(Distribution::DeterministicWave { min, max }) => (*min, *max), + _ => (0.0, 1.0), + }; + columns.push(Arc::new(Float64Array::from( + (0..rows) + .map(|r| wave_value(min, max, base_row + r)) + .collect::>(), + )) as ArrayRef); + } + ("field", "UINT64") => columns.push(Arc::new(UInt64Array::from( + (0..rows).map(|r| (base_row + r) as u64).collect::>(), + )) as ArrayRef), + ("timestamp", "TIMESTAMP(9)") => columns.push(Arc::new(TimestampNanosecondArray::from( + (0..rows) + .map(|r| timestamp_for_row(layout, base_row, r)) + .collect::>(), + )) as ArrayRef), + ("timestamp", "TIMESTAMP(3)") => { + columns.push(Arc::new(TimestampMillisecondArray::from( + (0..rows) + .map(|r| timestamp_for_row(layout, base_row, r) / 1_000_000) + .collect::>(), + )) as ArrayRef) + } + other => panic!("unsupported column combination {other:?}"), + } + } + + let mut pk_builder = BinaryDictionaryBuilder::::new(); + for row in 0..rows { + let series = series_for_row(layout, sst_idx, base_row, row); + let values = table + .columns + .iter() + .filter(|c| c.semantic == "tag") + .map(|c| (c.name.clone(), tag_value(c, series))) + .collect(); + pk_builder + .append(encode_dense_primary_key(table, &values)) + .expect("append generated dense primary key to Arrow dictionary builder"); + } + columns.push(Arc::new(pk_builder.finish())); + columns.push(Arc::new(UInt64Array::from_value(sequence, rows))); + columns.push(Arc::new(UInt8Array::from_value(OpType::Put as u8, rows))); + RecordBatch::try_new(flat_schema, columns) + .expect("generated fixture columns should match flat SST Arrow schema") +} + +fn file_meta_from_sst_info( + info: &SstInfo, + region_id: RegionId, + file_id: FileId, + sequence: u64, +) -> FileMeta { + FileMeta { + region_id, + file_id, + time_range: info.time_range, + level: 0, + file_size: info.file_size, + max_row_group_uncompressed_size: info.max_row_group_uncompressed_size, + available_indexes: Default::default(), + indexes: Default::default(), + index_file_size: 0, + index_version: 0, + num_rows: info.num_rows as u64, + num_row_groups: info.num_row_groups, + sequence: NonZeroU64::new(sequence), + partition_expr: None, + num_series: info.num_series, + primary_key_min: None, + primary_key_max: None, + } +} + +fn deterministic_file_id(seed: u64, index: usize) -> FileId { + let seed_hi = ((seed >> 16) & 0xffff) as u16; + let seed_lo = (seed & 0xffff) as u16; + FileId::parse_str(&format!( + "00000000-0000-{seed_hi:04x}-{seed_lo:04x}-{index:012x}" + )) + .expect("deterministic file id must be a valid UUID") +} + +fn remap_sst_file( + object_store_dir: &Path, + table_dir: &str, + region_id: RegionId, + old_file_id: FileId, + new_file_id: FileId, +) { + if old_file_id == new_file_id { + return; + } + let old_path = object_store_dir.join(mito2::sst::location::sst_file_path( + table_dir, + RegionFileId::new(region_id, old_file_id), + PathType::Bare, + )); + let new_path = object_store_dir.join(mito2::sst::location::sst_file_path( + table_dir, + RegionFileId::new(region_id, new_file_id), + PathType::Bare, + )); + if let Some(parent) = new_path.parent() { + fs::create_dir_all(parent).expect("failed to create remapped SST parent directory"); + } + fs::rename(&old_path, &new_path).unwrap_or_else(|err| { + panic!( + "failed to remap SST file {} -> {}: {err}", + old_path.display(), + new_path.display() + ) + }); +} +pub(super) async fn run_direct_sst(args: DirectArgs) { + let case_path = args.case.as_ref().expect("--case is required"); + let out_dir = args.out_dir.as_ref().expect("--out-dir is required"); + let case_text = fs::read_to_string(case_path).expect("failed to read query perf case TOML"); + let case: CaseFile = toml::from_str(&case_text).expect("failed to parse query perf case TOML"); + let case_name = case_name_from_path(case_path); + let scenario = case.scenario.direct_readable_sst(); + if scenario.tables.is_empty() || scenario.layout.regions != 1 { + panic!("fixture generator supports one or more tables and exactly one region per table"); + } + if scenario.layout.time_range_layout != "non_overlapping_per_sst" { + panic!("MVP supports time_range_layout=non_overlapping_per_sst"); + } + if scenario.layout.series_layout == "timestamp_major" + && !scenario + .layout + .rows_per_sst + .is_multiple_of(scenario.layout.series_count.get()) + { + panic!( + "series_layout=timestamp_major requires rows_per_sst to be divisible by series_count" + ); + } + if scenario.layout.sst_count > 1000 && !args.allow_large { + panic!("sst_count exceeds 1000; pass --allow-large"); + } + let seed = scenario.seed.unwrap_or(0); + let table_index = match args.table.as_deref() { + Some(name) => scenario + .tables + .iter() + .position(|table| table.name == name) + .unwrap_or_else(|| panic!("--table {name} was not found in case {case_name}")), + None if scenario.tables.len() == 1 => 0, + None => panic!( + "case {} contains {} tables; pass --table to choose one", + case_name, + scenario.tables.len() + ), + }; + let table = &scenario.tables[table_index]; + if table.engine != "mito" { + panic!("MVP supports engine=mito"); + } + let region_id = RegionId::from(args.region_id); + let table_dir = args + .table_dir + .unwrap_or_else(|| format!("data/{}/{}/", table.database, table.name)); + if Path::new(&table_dir).is_absolute() { + panic!("--table-dir must be relative"); + } + if table_dir.trim_end_matches('/').rsplit('/').next() + == Some(region_name(region_id.table_id(), region_id.region_sequence()).as_str()) + { + panic!("--table-dir must be table-level directory"); + } + let region_dir = + mito2::sst::location::region_dir_from_table_dir(&table_dir, region_id, PathType::Bare); + let format = match table.sst_format.as_deref().unwrap_or("flat") { + "flat" => FormatType::Flat, + "primary_key" | "primary-key" => FormatType::PrimaryKey, + other => panic!("unsupported sst_format {other}"), + }; + println!( + "query_perf_fixture case={} scenario={} table={}.{} ssts={} rows_per_sst={} format={format:?}", + case_name, + case.scenario.kind(), + table.database, + table.name, + scenario.layout.sst_count, + scenario.layout.rows_per_sst + ); + println!( + "out_dir={} table_dir={} region_dir={}", + out_dir.display(), + table_dir, + region_dir + ); + if args.dry_run { + return; + } + + let obj_store_dir = out_dir.join("object-store"); + let manifest_dir = out_dir.join("manifest"); + fs::create_dir_all(&obj_store_dir).expect("failed to create fixture object-store directory"); + fs::create_dir_all(&manifest_dir).expect("failed to create fixture manifest directory"); + let ostorage = ObjectStore::new(FsBuilder::default().root(&obj_store_dir.to_string_lossy())) + .expect("failed to create filesystem object store for fixture output") + .finish(); + let metadata: RegionMetadataRef = Arc::new(build_region_metadata(table, region_id)); + let mut files = HashMap::with_capacity(scenario.layout.sst_count); + let mut next_file_index = 1; + for i in 0..scenario.layout.sst_count { + let sequence = 1000 + i as u64; + let batch = generate_record_batch(table, &metadata, &scenario.layout, i, sequence); + let source = + FlatSource::new_iter(batch.schema(), Box::new(vec![batch].into_iter().map(Ok))); + let mut metrics = Metrics::new(WriteType::Flush); + let mut writer = ParquetWriter::new_with_object_store( + ostorage.clone(), + metadata.clone(), + IndexConfig::default(), + NoopIndexBuilder, + FixedPathProvider { + table_dir: table_dir.clone(), + }, + &mut metrics, + ) + .await; + let opts = WriteOptions { + write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, + row_group_size: scenario.layout.row_group_size, + max_file_size: None, + }; + let infos = match format { + FormatType::Flat => writer.write_all_flat(source, Some(sequence), &opts).await, + FormatType::PrimaryKey => { + writer + .write_all_flat_as_primary_key(source, Some(sequence), &opts) + .await + } + } + .expect("failed to write fixture SST through Mito parquet writer"); + for info in infos { + let file_id = deterministic_file_id( + seed.wrapping_add((table_index as u64) << 16), + next_file_index, + ); + next_file_index += 1; + remap_sst_file(&obj_store_dir, &table_dir, region_id, info.file_id, file_id); + files.insert( + file_id, + file_meta_from_sst_info(&info, region_id, file_id, sequence), + ); + } + } + let overall_seq = files + .values() + .filter_map(|m| m.sequence) + .map(|s| s.get()) + .max() + .unwrap_or(0); + let manifest = RegionManifest { + metadata, + files, + removed_files: RemovedFilesRecord::default(), + flushed_entry_id: overall_seq, + flushed_sequence: overall_seq, + committed_sequence: Some(overall_seq), + manifest_version: args.checkpoint_version, + truncated_entry_id: None, + compaction_time_window: None, + sst_format: format, + append_mode: table.append_mode, + }; + let checkpoint = RegionCheckpoint { + last_version: args.checkpoint_version, + compacted_actions: manifest.files.len(), + checkpoint: Some(manifest.clone()), + }; + let checkpoint_bytes = checkpoint + .encode() + .expect("failed to encode fixture region checkpoint"); + let checkpoint_path = manifest_dir.join(format!("{:020}.checkpoint", args.checkpoint_version)); + fs::write(&checkpoint_path, &checkpoint_bytes).expect("failed to write checkpoint file"); + fs::write(manifest_dir.join("_last_checkpoint"), serde_json::to_vec_pretty(&serde_json::json!({ "size": checkpoint_bytes.len(), "version": args.checkpoint_version, "checksum": null, "extend_metadata": {} })).expect("failed to serialize _last_checkpoint metadata")).expect("failed to write _last_checkpoint file"); + let files_jsonl_path = out_dir.join("files.jsonl"); + let mut jsonl = BufWriter::new( + fs::File::create(&files_jsonl_path).expect("failed to create fixture files.jsonl"), + ); + let mut entries: Vec<_> = manifest.files.iter().collect(); + entries.sort_by_key(|(id, _)| id.to_string()); + for (file_id, meta) in entries { + writeln!(jsonl, "{}", serde_json::to_string(&serde_json::json!({ "file_id": file_id.to_string(), "region_id": meta.region_id.as_u64(), "object_path": mito2::sst::location::sst_file_path(&table_dir, RegionFileId::new(meta.region_id, *file_id), PathType::Bare), "time_range_start": meta.time_range.0.value(), "time_range_end": meta.time_range.1.value(), "num_rows": meta.num_rows, "num_row_groups": meta.num_row_groups, "file_size": meta.file_size, "num_series": meta.num_series, "sequence": meta.sequence.map(|s| s.get()) })).expect("failed to serialize fixture SST metadata JSONL entry")).expect("failed to write fixture SST metadata JSONL entry"); + } + jsonl.flush().expect("failed to flush fixture files.jsonl"); + fs::write(out_dir.join("summary.json"), serde_json::to_vec_pretty(&serde_json::json!({ "case": case_name, "seed": seed, "table_index": table_index, "table": table.name, "database": table.database, "region_id": region_id.as_u64(), "table_dir": table_dir, "region_dir": region_dir, "sst_format": format!("{format:?}"), "sst_count": scenario.layout.sst_count, "rows_per_sst": scenario.layout.rows_per_sst, "row_group_size": scenario.layout.row_group_size, "total_rows": scenario.layout.sst_count * scenario.layout.rows_per_sst, "checkpoint_path": checkpoint_path, "files_jsonl_path": files_jsonl_path, "readback_validated": false, "metadata_source": "synthetic" })).expect("failed to serialize fixture summary")).expect("failed to write fixture summary.json"); + println!("Done. wrote {} SST file entries", manifest.files.len()); +} diff --git a/src/cmd/src/bin/query_perf_fixture/inspect_footer.rs b/src/cmd/src/bin/query_perf_fixture/inspect_footer.rs new file mode 100644 index 00000000000..a52005c4e30 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/inspect_footer.rs @@ -0,0 +1,181 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::BTreeSet; +use std::fs::{self, File}; +use std::path::{Path, PathBuf}; + +use clap::Args as ClapArgs; +use parquet::file::reader::{FileReader, SerializedFileReader}; +use serde::Serialize; + +#[derive(Debug, ClapArgs)] +pub(super) struct InspectFooterArgs { + #[arg(long)] + root: PathBuf, + #[arg(long, default_value = "greptime_value")] + column: String, + #[arg(long)] + include_metadata_files: bool, +} + +#[derive(Debug, Serialize)] +struct FooterReport { + root: String, + summary: FooterSummary, + files: Vec, +} +#[derive(Debug, Default, Serialize)] +struct FooterSummary { + column: String, + file_count: usize, + files_with_column: usize, + total_file_size: u64, + total_rows: i64, + column_compressed_size: i64, + column_uncompressed_size: i64, + unique_encodings: Vec, +} +#[derive(Debug, Serialize)] +struct FooterFileReport { + path: String, + relative_path: String, + file_size: u64, + num_rows: i64, + num_row_groups: usize, + columns: Vec, +} +#[derive(Debug, Serialize)] +struct FooterColumnChunkReport { + row_group_index: usize, + column_index: usize, + column_path: String, + encodings: Vec, + compression: String, + compressed_size: i64, + uncompressed_size: i64, + num_values: i64, +} + +pub(super) fn run_inspect_footer( + args: InspectFooterArgs, +) -> Result<(), Box> { + let files = collect_parquet_files(&args.root, args.include_metadata_files)?; + let reports = files + .iter() + .map(|p| inspect_file(&args.root, p, &args.column)) + .collect::, _>>()?; + println!( + "{}", + serde_json::to_string_pretty(&FooterReport { + root: args.root.display().to_string(), + summary: summarize_footer(&args.column, &reports), + files: reports + })? + ); + Ok(()) +} +fn collect_parquet_files( + root: &Path, + include_metadata_files: bool, +) -> Result, Box> { + let mut files = Vec::new(); + collect_parquet_files_inner(root, root, include_metadata_files, &mut files)?; + files.sort(); + Ok(files) +} +fn collect_parquet_files_inner( + root: &Path, + path: &Path, + include_metadata_files: bool, + files: &mut Vec, +) -> Result<(), Box> { + if path.is_dir() { + for entry in fs::read_dir(path)? { + collect_parquet_files_inner(root, &entry?.path(), include_metadata_files, files)?; + } + } else if path.extension().is_some_and(|ext| ext == "parquet") + && (include_metadata_files + || !path + .strip_prefix(root) + .unwrap_or(path) + .components() + .any(|c| c.as_os_str() == "metadata")) + { + files.push(path.to_path_buf()); + } + Ok(()) +} +fn inspect_file( + root: &Path, + path: &Path, + column: &str, +) -> Result> { + let file_size = fs::metadata(path)?.len(); + let reader = SerializedFileReader::new(File::open(path)?)?; + let metadata = reader.metadata().file_metadata(); + let row_groups = reader.metadata().row_groups(); + let mut columns = Vec::new(); + for (row_group_index, rg) in row_groups.iter().enumerate() { + for (column_index, chunk) in rg.columns().iter().enumerate() { + let column_path = chunk.column_path().string(); + if column_path == column { + columns.push(FooterColumnChunkReport { + row_group_index, + column_index, + column_path, + encodings: chunk.encodings().map(|e| format!("{e:?}")).collect(), + compression: format!("{:?}", chunk.compression()), + compressed_size: chunk.compressed_size(), + uncompressed_size: chunk.uncompressed_size(), + num_values: chunk.num_values(), + }); + } + } + } + Ok(FooterFileReport { + path: path.display().to_string(), + relative_path: path + .strip_prefix(root) + .unwrap_or(path) + .display() + .to_string(), + file_size, + num_rows: metadata.num_rows(), + num_row_groups: row_groups.len(), + columns, + }) +} +fn summarize_footer(column: &str, files: &[FooterFileReport]) -> FooterSummary { + let mut encodings = BTreeSet::new(); + let mut s = FooterSummary { + column: column.to_string(), + file_count: files.len(), + ..Default::default() + }; + for file in files { + s.total_file_size += file.file_size; + s.total_rows += file.num_rows; + if !file.columns.is_empty() { + s.files_with_column += 1; + } + for c in &file.columns { + s.column_compressed_size += c.compressed_size; + s.column_uncompressed_size += c.uncompressed_size; + encodings.extend(c.encodings.iter().cloned()); + } + } + s.unique_encodings = encodings.into_iter().collect(); + s +} diff --git a/src/cmd/src/bin/query_perf_fixture/mod.rs b/src/cmd/src/bin/query_perf_fixture/mod.rs new file mode 100644 index 00000000000..00cef327c78 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/mod.rs @@ -0,0 +1,133 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +mod case; +mod direct_sst; +mod inspect_footer; +mod prom_remote_write; +mod util; + +use std::fs; +use std::path::PathBuf; + +use case::*; +use clap::{Args as ClapArgs, Parser, Subcommand}; +use direct_sst::run_direct_sst; +use inspect_footer::{InspectFooterArgs, run_inspect_footer}; +use prom_remote_write::{PromRemoteWriteArgs, run_prom_remote_write}; +use serde_json::json; + +#[derive(Parser, Debug)] +#[command(name = "query_perf_fixture")] +#[command(about = "Query performance fixture helper")] +struct Args { + #[command(subcommand)] + command: Option, + + #[command(flatten)] + legacy: DirectArgs, +} + +#[derive(Subcommand, Debug)] +enum Command { + DirectSst(DirectArgs), + PromRemoteWrite(PromRemoteWriteArgs), + InspectFooter(InspectFooterArgs), + Plan(PlanArgs), +} + +#[derive(ClapArgs, Debug, Clone)] +struct DirectArgs { + /// TOML case file, for example tests/perf/query_cases/.../case.toml. + #[arg(long, value_name = "PATH")] + case: Option, + + /// Output directory (creates object-store/, manifest/, files.jsonl, summary.json). + #[arg(long, value_name = "DIR")] + out_dir: Option, + + /// Region ID to synthesize. Defaults to table id 1024, region 0. + #[arg(long, default_value = "4398046511104")] + region_id: u64, + + /// Table directory relative to object-store root. Defaults to data/{database}/{table}/. + #[arg(long, value_name = "DIR")] + table_dir: Option, + + /// Table name to generate when the case contains multiple [[tables]]. + #[arg(long, value_name = "NAME")] + table: Option, + + /// Manifest/checkpoint version. + #[arg(long, default_value = "1000000")] + checkpoint_version: u64, + + /// Safety flag when sst_count exceeds 1000. + #[arg(long)] + allow_large: bool, + + /// Print plan only. + #[arg(long)] + dry_run: bool, +} + +#[derive(ClapArgs, Debug)] +struct PlanArgs { + #[arg(long, value_name = "PATH")] + case: PathBuf, +} + +pub async fn run() { + let args = Args::parse(); + match args.command { + Some(Command::DirectSst(direct)) => run_direct_sst(direct).await, + Some(Command::PromRemoteWrite(rw)) => run_prom_remote_write(rw) + .await + .expect("prom remote write failed"), + Some(Command::InspectFooter(inspect)) => { + run_inspect_footer(inspect).expect("inspect footer failed") + } + Some(Command::Plan(plan)) => run_plan(plan).expect("plan failed"), + None => run_direct_sst(args.legacy).await, + } +} + +fn run_plan(args: PlanArgs) -> Result<(), Box> { + let case_text = fs::read_to_string(&args.case)?; + let mut case: CaseFile = toml::from_str(&case_text)?; + if let Scenario::PromRemoteWriteThenQuery(s) = &mut case.scenario + && let Some(storage) = &mut s.remote_write.storage + { + storage.populate_planned_thresholds(); + if storage.inspect { + let mut rb = s.remote_write.read_bench.clone().unwrap_or_default(); + if rb.enabled && rb.projection.is_empty() { + rb.projection = vec![storage.column.clone()]; + } + s.remote_write.read_bench = Some(rb); + } else if s + .remote_write + .read_bench + .as_ref() + .is_some_and(|rb| rb.enabled) + { + return Err("scenario.remote_write.read_bench requires scenario.remote_write.storage.inspect = true".into()); + } + } + println!( + "{}", + serde_json::to_string_pretty(&json!({"schema_version": 1, "scenario": case.scenario}))? + ); + Ok(()) +} diff --git a/src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs b/src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs new file mode 100644 index 00000000000..1f36b0c2d23 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs @@ -0,0 +1,333 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use clap::Args as ClapArgs; +use prost::Message; +use serde_json::json; +use servers::prom_store::snappy_compress; + +use crate::query_perf_fixture::case::ValuePattern; + +#[derive(Debug, ClapArgs)] +pub(super) struct PromRemoteWriteArgs { + #[arg(long, default_value = "http://127.0.0.1:4000/v1/prometheus/write")] + endpoint: String, + #[arg(long, default_value = "public")] + database: String, + #[arg(long)] + metric: String, + #[arg(long, default_value = "greptime_physical_table")] + physical_table: String, + #[arg(long, default_value_t = 8)] + series_count: u64, + #[arg(long, default_value_t = 30)] + samples_per_series: u64, + #[arg(long, default_value_t = 1_704_067_200_000)] + start_unix_millis: i64, + #[arg(long, default_value_t = 15_000)] + step_millis: i64, + #[arg(long, alias = "batch-size", default_value_t = 8)] + chunk_series_count: u64, + #[arg(long, default_value_t = 60)] + timeout_seconds: u64, + #[arg(long, default_value_t = ValuePattern::Linear)] + value_pattern: ValuePattern, + #[arg(long, default_value_t = 0.0)] + value_base: f64, + #[arg(long, default_value_t = 0.125)] + value_step: f64, + #[arg(long, default_value_t = 97)] + value_cardinality: u64, + #[arg(long, default_value_t = 0)] + value_seed: u64, + #[arg(long, default_value_t = 8)] + value_run_length: u64, + #[arg(long, default_value_t = 100)] + value_stall_every: u64, + #[arg(long, default_value_t = 16)] + value_stall_length: u64, + #[arg(long, default_value_t = 5)] + value_mixed_every: u64, + #[arg(long, default_value_t = 0)] + value_sample_offset: u64, + #[arg(long)] + value_total_samples_per_series: Option, +} + +pub(super) async fn run_prom_remote_write( + args: PromRemoteWriteArgs, +) -> Result<(), Box> { + use api::prom_store::remote::WriteRequest; + let started = std::time::Instant::now(); + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(args.timeout_seconds)) + .build()?; + let mut batches = 0_u64; + let mut rows = 0_u64; + let mut http_statuses = Vec::new(); + for first in (0..args.series_count).step_by(args.chunk_series_count.max(1) as usize) { + let last = (first + args.chunk_series_count.max(1)).min(args.series_count); + let timeseries = (first..last) + .map(|series_idx| make_prom_series(&args, series_idx)) + .collect::>(); + rows += timeseries + .iter() + .map(|ts| ts.samples.len() as u64) + .sum::(); + let body = snappy_compress( + &WriteRequest { + timeseries, + ..Default::default() + } + .encode_to_vec(), + )?; + let response = client + .post(&args.endpoint) + .header("Content-Encoding", "snappy") + .body(body) + .send() + .await?; + let status = response.status(); + http_statuses.push(status.as_u16()); + if !status.is_success() { + return Err(format!( + "remote-write failed with status {status}: {}", + response.text().await.unwrap_or_default() + ) + .into()); + } + batches += 1; + } + println!( + "{}", + json!({"status":"ok","endpoint":args.endpoint,"database":args.database,"metric":args.metric,"physical_table":args.physical_table,"series_count":args.series_count,"samples_per_series":args.samples_per_series,"rows":rows,"samples_written":rows,"batches":batches,"elapsed_seconds":started.elapsed().as_secs_f64(),"http_statuses":http_statuses,"value":{"pattern":args.value_pattern.to_string(),"base":args.value_base,"step":args.value_step,"cardinality":args.value_cardinality,"seed":args.value_seed,"run_length":args.value_run_length,"stall_every":args.value_stall_every,"stall_length":args.value_stall_length,"mixed_every":args.value_mixed_every,"sample_offset":args.value_sample_offset,"total_samples_per_series":args.value_total_samples_per_series.unwrap_or(args.samples_per_series)}}) + ); + Ok(()) +} + +fn make_prom_series( + args: &PromRemoteWriteArgs, + series_idx: u64, +) -> api::prom_store::remote::TimeSeries { + let mut labels = vec![ + prom_label("__name__", &args.metric), + prom_label("x_greptime_database", &args.database), + prom_label("x_greptime_physical_table", &args.physical_table), + prom_label("host", &format!("host{:04}", series_idx % 1024)), + prom_label("instance", &format!("instance{:06}", series_idx)), + ]; + labels.sort_unstable_by(|l, r| l.name.cmp(&r.name)); + let samples = (0..args.samples_per_series) + .map(|sample_idx| api::prom_store::remote::Sample { + value: deterministic_prom_value(args, series_idx, sample_idx), + timestamp: args.start_unix_millis + (sample_idx as i64 * args.step_millis), + }) + .collect(); + api::prom_store::remote::TimeSeries { + labels, + samples, + ..Default::default() + } +} +fn prom_label(name: &str, value: &str) -> api::prom_store::remote::Label { + api::prom_store::remote::Label { + name: name.to_string(), + value: value.to_string(), + } +} +fn deterministic_prom_value(args: &PromRemoteWriteArgs, series_idx: u64, sample_idx: u64) -> f64 { + let ordinal = series_idx + * args + .value_total_samples_per_series + .unwrap_or(args.samples_per_series) + + args.value_sample_offset + + sample_idx; + let local = args.value_sample_offset + sample_idx; + match args.value_pattern { + ValuePattern::Linear => { + args.value_base + (series_idx % 97) as f64 + local as f64 * args.value_step + } + ValuePattern::Constant => args.value_base, + ValuePattern::Modulo => { + args.value_base + (ordinal % args.value_cardinality.max(1)) as f64 * args.value_step + } + ValuePattern::Unique => args.value_base + ordinal as f64 * args.value_step, + ValuePattern::SeededRandom => { + args.value_base + + (splitmix64(ordinal ^ args.value_seed) % args.value_cardinality.max(1)) as f64 + * args.value_step + } + ValuePattern::RunLength => { + args.value_base + + ((ordinal / args.value_run_length.max(1)) % args.value_cardinality.max(1)) as f64 + * args.value_step + } + ValuePattern::QuantizedSignal => { + args.value_base + + (((series_idx % args.value_cardinality.max(1)) + + (local / args.value_run_length.max(1))) + % args.value_cardinality.max(1)) as f64 + * args.value_step + } + ValuePattern::SignalWithSporadicStalls => { + let every = args.value_stall_every.max(1); + let phase = local % every; + let effective = if phase < args.value_stall_length.min(every) { + local - phase + } else { + local + }; + args.value_base + (series_idx % 97) as f64 + effective as f64 * args.value_step + } + ValuePattern::MixedSignalRepeated => { + if local.is_multiple_of(args.value_mixed_every.max(1)) { + args.value_base + } else { + args.value_base + (series_idx % 97) as f64 + local as f64 * args.value_step + } + } + } +} +fn splitmix64(mut value: u64) -> u64 { + value = value.wrapping_add(0x9e3779b97f4a7c15); + let mut mixed = value; + mixed = (mixed ^ (mixed >> 30)).wrapping_mul(0xbf58476d1ce4e5b9); + mixed = (mixed ^ (mixed >> 27)).wrapping_mul(0x94d049bb133111eb); + mixed ^ (mixed >> 31) +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use super::*; + + fn args(pattern: ValuePattern) -> PromRemoteWriteArgs { + PromRemoteWriteArgs { + endpoint: "http://127.0.0.1:4000/v1/prometheus/write".to_string(), + database: "public".to_string(), + metric: "test_metric".to_string(), + physical_table: "greptime_physical_table".to_string(), + series_count: 2, + samples_per_series: 64, + start_unix_millis: 1_704_067_200_000, + step_millis: 60_000, + chunk_series_count: 2, + timeout_seconds: 60, + value_pattern: pattern, + value_base: 0.0, + value_step: 1.0, + value_cardinality: 1_024, + value_seed: 8_444, + value_run_length: 16, + value_stall_every: 100, + value_stall_length: 16, + value_mixed_every: 5, + value_sample_offset: 0, + value_total_samples_per_series: None, + } + } + + #[test] + fn seeded_random_is_deterministic_in_range_and_highly_distinct() { + let mut args = args(ValuePattern::SeededRandom); + args.samples_per_series = 4_096; + args.value_cardinality = 29_491_200; + + let values = (0..args.samples_per_series) + .map(|sample_idx| deterministic_prom_value(&args, 0, sample_idx)) + .collect::>(); + let repeated = (0..args.samples_per_series) + .map(|sample_idx| deterministic_prom_value(&args, 0, sample_idx)) + .collect::>(); + + assert_eq!(values, repeated); + assert!(values.iter().all(|value| { + (args.value_base..=args.value_base + (args.value_cardinality - 1) as f64) + .contains(value) + })); + assert!( + values + .iter() + .map(|value| value.to_bits()) + .collect::>() + .len() + > 4_000 + ); + } + + #[test] + fn run_length_values_have_exact_sixteen_sample_runs() { + let args = args(ValuePattern::RunLength); + + for sample_idx in 0..64 { + assert_eq!( + deterministic_prom_value(&args, 0, sample_idx), + (sample_idx / 16) as f64 + ); + } + } + + #[test] + fn mixed_every_five_repeats_only_the_expected_positions() { + let mut args = args(ValuePattern::MixedSignalRepeated); + args.value_step = 0.125; + + for sample_idx in 0..20 { + let expected = if sample_idx % 5 == 0 { + 0.0 + } else { + 1.0 + sample_idx as f64 * 0.125 + }; + assert_eq!(deterministic_prom_value(&args, 1, sample_idx), expected); + } + } + + #[test] + fn chunked_values_match_monolithic_values_through_offsets() { + let mut monolithic = args(ValuePattern::SeededRandom); + monolithic.samples_per_series = 14_400; + monolithic.value_total_samples_per_series = Some(14_400); + + for series_idx in 0..2 { + let expected = (0..monolithic.samples_per_series) + .map(|sample_idx| deterministic_prom_value(&monolithic, series_idx, sample_idx)) + .collect::>(); + for chunk_idx in 0..10 { + let mut chunk = args(ValuePattern::SeededRandom); + chunk.samples_per_series = 1_440; + chunk.value_sample_offset = chunk_idx * 1_440; + chunk.value_total_samples_per_series = Some(14_400); + for sample_idx in 0..chunk.samples_per_series { + assert_eq!( + deterministic_prom_value(&chunk, series_idx, sample_idx), + expected[(chunk.value_sample_offset + sample_idx) as usize] + ); + } + } + } + } + + #[test] + fn values_continue_across_series_boundaries() { + let mut args = args(ValuePattern::Unique); + args.samples_per_series = 16; + args.value_total_samples_per_series = Some(16); + args.value_base = 1.0; + args.value_step = 0.5; + + assert_eq!(deterministic_prom_value(&args, 0, 15), 8.5); + assert_eq!(deterministic_prom_value(&args, 1, 0), 9.0); + } +} diff --git a/src/cmd/src/bin/query_perf_fixture/util.rs b/src/cmd/src/bin/query_perf_fixture/util.rs new file mode 100644 index 00000000000..2c86baa6839 --- /dev/null +++ b/src/cmd/src/bin/query_perf_fixture/util.rs @@ -0,0 +1,23 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::path::Path; + +pub(super) fn case_name_from_path(path: &Path) -> String { + path.parent() + .and_then(Path::file_name) + .and_then(|name| name.to_str()) + .unwrap_or("query_perf_case") + .to_string() +} diff --git a/src/cmd/src/datanode.rs b/src/cmd/src/datanode.rs index 744624a5084..c33be6c6f33 100644 --- a/src/cmd/src/datanode.rs +++ b/src/cmd/src/datanode.rs @@ -15,11 +15,10 @@ pub mod builder; #[allow(clippy::print_stdout)] pub(crate) mod objbench; -#[cfg(feature = "dev-tools")] #[allow(clippy::print_stdout)] -mod parquetbench; +pub mod parquetbench; #[allow(clippy::print_stdout)] -mod scanbench; +pub mod scanbench; use std::path::Path; use std::time::Duration; @@ -40,7 +39,6 @@ use tracing_appender::non_blocking::WorkerGuard; use crate::App; use crate::datanode::builder::InstanceBuilder; use crate::datanode::objbench::ObjbenchCommand; -#[cfg(feature = "dev-tools")] use crate::datanode::parquetbench::ParquetbenchCommand; use crate::datanode::scanbench::ScanbenchCommand; use crate::error::{ @@ -115,7 +113,6 @@ impl Command { SubCommand::Start(cmd) => cmd.load_options(global_options), // Bench commands are standalone utilities and don't need to load DatanodeOptions. SubCommand::Objbench(_) | SubCommand::Scanbench(_) => Self::default_bench_options(), - #[cfg(feature = "dev-tools")] SubCommand::Parquetbench(_) => Self::default_bench_options(), } } @@ -141,7 +138,6 @@ pub enum SubCommand { /// Scan benchmark tool - benchmarks scanning a region directly from storage Scanbench(ScanbenchCommand), /// Benchmark scanning a single parquet SST. - #[cfg(feature = "dev-tools")] Parquetbench(ParquetbenchCommand), } @@ -160,7 +156,6 @@ impl SubCommand { cmd.run().await?; std::process::exit(0); } - #[cfg(feature = "dev-tools")] SubCommand::Parquetbench(cmd) => { cmd.run().await?; std::process::exit(0); diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 5096c9221cd..7a28366eb86 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -552,6 +552,7 @@ impl StartCommand { flow_metadata_allocator: flow_metadata_allocator.clone(), region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl), soft_drop_enabled: false, + soft_drop_retention: None, }; let ddl_manager = DdlManager::try_new( diff --git a/src/cmd/src/user.rs b/src/cmd/src/user.rs index 6f7e28f9687..cf7a85b11d2 100644 --- a/src/cmd/src/user.rs +++ b/src/cmd/src/user.rs @@ -17,8 +17,9 @@ use std::io; use auth::{ - DEFAULT_PBKDF2_SHA256_ITERATIONS, MAX_PBKDF2_SHA256_SALT_LEN, + DEFAULT_PBKDF2_SHA256_ITERATIONS, DEFAULT_PBKDF2_SHA256_SALT_LEN, MAX_PBKDF2_SHA256_SALT_LEN, format_mysql_native_password_verifier, format_pbkdf2_sha256_password_verifier, + format_pg_scram_sha256_password_verifier, }; use clap::{ArgGroup, Parser, Subcommand, ValueEnum}; use rand::RngCore; @@ -65,15 +66,15 @@ pub struct HashPasswordCommand { #[clap(long)] password_stdin: bool, - /// PBKDF2-SHA256 iteration count. + /// PBKDF2-SHA256 / SCRAM-SHA-256 iteration count. #[clap(long, default_value_t = DEFAULT_PBKDF2_SHA256_ITERATIONS)] iterations: u32, - /// PBKDF2-SHA256 random salt length in bytes. - #[clap(long, default_value_t = 16)] + /// PBKDF2-SHA256 / SCRAM-SHA-256 random salt length in bytes. + #[clap(long, default_value_t = DEFAULT_PBKDF2_SHA256_SALT_LEN)] salt_len: usize, - /// PBKDF2-SHA256 salt as hex. Mainly useful for deterministic automation. + /// PBKDF2-SHA256 / SCRAM-SHA-256 salt as hex. Mainly useful for deterministic automation. #[clap(long)] salt_hex: Option, } @@ -83,6 +84,7 @@ pub struct HashPasswordCommand { enum PasswordFormat { Pbkdf2Sha256, MysqlNativePassword, + PgScramSha256, } impl HashPasswordCommand { @@ -98,6 +100,16 @@ impl HashPasswordCommand { PasswordFormat::MysqlNativePassword => { format_mysql_native_password_verifier(password.as_bytes()) } + PasswordFormat::PgScramSha256 => { + let salt = self.pbkdf2_salt()?; + format_pg_scram_sha256_password_verifier( + password.as_bytes(), + &salt, + self.iterations, + ) + .map_err(common_error::ext::BoxedError::new) + .context(error::OtherSnafu)? + } }; println!("{verifier}"); @@ -186,6 +198,29 @@ mod tests { ); } + #[test] + fn test_hash_password_command_with_pg_scram_sha256() { + let cmd = HashPasswordCommand { + format: PasswordFormat::PgScramSha256, + password: Some("password".to_string()), + password_stdin: false, + iterations: 4096, + salt_len: 16, + salt_hex: Some("73616c74".to_string()), + }; + + let password = cmd.read_password().unwrap(); + let salt = cmd.pbkdf2_salt().unwrap(); + let verifier = + format_pg_scram_sha256_password_verifier(password.as_bytes(), &salt, cmd.iterations) + .unwrap(); + + assert_eq!( + "pg_scram_sha256:4096:73616c74:945e1c466fc9932efadc23781edc5d1e78d5e10f005933652af1a6105154f084:b9bf0e811b1fb6793671c0cc3adedf7c75cd72291191092ad65878c5a02aad2c", + verifier + ); + } + #[test] fn test_hash_password_command_with_mysql_native_password() { let verifier = format_mysql_native_password_verifier("123456".as_bytes()); diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index f07a51eba48..c876a424733 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -12,13 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::io::Write; use std::time::Duration; use cmd::options::GreptimeOptions; use common_base::memory_limit::MemoryLimit; +use common_base::readable_size::ReadableSize; use common_config::{Configurable, DEFAULT_DATA_HOME, ENV_VAR_SEP}; use common_options::datanode::{ClientOptions, DatanodeClientOptions}; use common_telemetry::logging::{DEFAULT_LOGGING_DIR, DEFAULT_OTLP_HTTP_ENDPOINT, LoggingOptions}; +use common_test_util::temp_dir::create_named_temp_file; use common_wal::config::DatanodeWalConfig; use common_wal::config::raft_engine::RaftEngineConfig; use datanode::config::{DatanodeOptions, RegionEngineConfig, StorageConfig}; @@ -90,13 +93,13 @@ fn test_load_datanode_example_config() { region_engine: vec![ RegionEngineConfig::Mito(MitoConfig { auto_flush_interval: Duration::from_secs(3600), + default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), scan_memory_limit: MemoryLimit::Unlimited, ..Default::default() }), RegionEngineConfig::File(FileEngineConfig {}), RegionEngineConfig::Metric(MetricEngineConfig { - sparse_primary_key_encoding: true, experimental_enable_metric_value_split: false, flush_metadata_region_interval: Duration::from_secs(30), }), @@ -109,6 +112,7 @@ fn test_load_datanode_example_config() { level: Some("info".to_string()), dir: format!("{}/{}", DEFAULT_DATA_HOME, DEFAULT_LOGGING_DIR), otlp_endpoint: Some(DEFAULT_OTLP_HTTP_ENDPOINT.to_string()), + otlp_export_protocol: Some(common_telemetry::logging::OtlpExportProtocol::Http), tracing_sample_ratio: Some(Default::default()), ..Default::default() }, @@ -148,6 +152,7 @@ fn test_load_frontend_example_config() { level: Some("info".to_string()), dir: format!("{}/{}", DEFAULT_DATA_HOME, DEFAULT_LOGGING_DIR), otlp_endpoint: Some(DEFAULT_OTLP_HTTP_ENDPOINT.to_string()), + otlp_export_protocol: Some(common_telemetry::logging::OtlpExportProtocol::Http), tracing_sample_ratio: Some(Default::default()), ..Default::default() }, @@ -185,6 +190,11 @@ fn test_load_metasrv_example_config() { let options = GreptimeOptions::::load_layered_options(example_config.to_str(), "") .unwrap(); + assert!(!options.component.gc.experimental_soft_drop.enable); + assert_eq!( + Duration::from_secs(7 * 24 * 60 * 60), + options.component.gc.experimental_soft_drop.retention + ); let expected = GreptimeOptions:: { component: MetasrvOptions { selector: SelectorType::default(), @@ -198,6 +208,7 @@ fn test_load_metasrv_example_config() { dir: format!("{}/{}", DEFAULT_DATA_HOME, DEFAULT_LOGGING_DIR), level: Some("info".to_string()), otlp_endpoint: Some(DEFAULT_OTLP_HTTP_ENDPOINT.to_string()), + otlp_export_protocol: Some(common_telemetry::logging::OtlpExportProtocol::Http), tracing_sample_ratio: Some(Default::default()), ..Default::default() }, @@ -297,13 +308,13 @@ fn test_load_standalone_example_config() { region_engine: vec![ RegionEngineConfig::Mito(MitoConfig { auto_flush_interval: Duration::from_secs(3600), + default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), scan_memory_limit: MemoryLimit::Unlimited, ..Default::default() }), RegionEngineConfig::File(FileEngineConfig {}), RegionEngineConfig::Metric(MetricEngineConfig { - sparse_primary_key_encoding: true, experimental_enable_metric_value_split: false, flush_metadata_region_interval: Duration::from_secs(30), }), @@ -316,6 +327,7 @@ fn test_load_standalone_example_config() { level: Some("info".to_string()), dir: format!("{}/{}", DEFAULT_DATA_HOME, DEFAULT_LOGGING_DIR), otlp_endpoint: Some(DEFAULT_OTLP_HTTP_ENDPOINT.to_string()), + otlp_export_protocol: Some(common_telemetry::logging::OtlpExportProtocol::Http), tracing_sample_ratio: Some(Default::default()), ..Default::default() }, @@ -381,3 +393,42 @@ fn test_load_heartbeat_env_vars_from_env() { similar_asserts::assert_eq!(standalone.component.heartbeat_env_vars, expected); }); } + +#[test] +fn test_load_metric_config_with_removed_sparse_primary_key_encoding() { + // The `sparse_primary_key_encoding` option was removed from the metric + // engine config and is now always-on. Existing user configs that still set + // it must be tolerated (ignored), not rejected. + let mut file = create_named_temp_file(); + let toml_str = r#" + node_id = 42 + + [meta_client] + metasrv_addrs = ["127.0.0.1:3002"] + + [[region_engine]] + [region_engine.metric] + sparse_primary_key_encoding = false + flush_metadata_region_interval = "30s" + "#; + write!(file, "{}", toml_str).unwrap(); + + let options = + GreptimeOptions::::load_layered_options(file.path().to_str(), "") + .expect("config with removed 'sparse_primary_key_encoding' should load without error"); + + // The unknown option is ignored; known metric engine options are still parsed/applied. + let metric = options + .component + .region_engine + .iter() + .find_map(|c| match c { + RegionEngineConfig::Metric(c) => Some(c), + _ => None, + }) + .expect("metric engine config should be present"); + assert_eq!( + metric.flush_metadata_region_interval, + Duration::from_secs(30) + ); +} diff --git a/src/common/datasource/src/file_format.rs b/src/common/datasource/src/file_format.rs index 028aeef14d6..61f838899ef 100644 --- a/src/common/datasource/src/file_format.rs +++ b/src/common/datasource/src/file_format.rs @@ -63,6 +63,7 @@ pub const FORMAT_SCHEMA_INFER_MAX_RECORD: &str = "schema_infer_max_record"; pub const FORMAT_HEADERS: &str = "headers"; pub const FORMAT_HAS_HEADER: &str = "has_header"; pub const FORMAT_SKIP_BAD_RECORDS: &str = "skip_bad_records"; +pub const FORMAT_STRICT_HEADERS: &str = "strict_headers"; pub const FORMAT_TYPE: &str = "format"; pub const FILE_PATTERN: &str = "pattern"; pub const TIMESTAMP_FORMAT: &str = "timestamp_format"; diff --git a/src/common/datasource/src/file_format/csv.rs b/src/common/datasource/src/file_format/csv.rs index 24b61340f5b..d016d4eabe9 100644 --- a/src/common/datasource/src/file_format/csv.rs +++ b/src/common/datasource/src/file_format/csv.rs @@ -49,6 +49,7 @@ const SKIP_BAD_RECORDS_BATCH_SIZE: usize = 1; pub struct CsvFormat { pub has_header: bool, pub skip_bad_records: bool, + pub strict_headers: bool, pub delimiter: u8, pub schema_infer_max_record: Option, pub compression_type: CompressionType, @@ -111,6 +112,16 @@ impl TryFrom<&HashMap> for CsvFormat { format.skip_bad_records = parse_bool(file_format::FORMAT_SKIP_BAD_RECORDS, skip_bad_records)?; }; + if let Some(strict_headers) = value.get(file_format::FORMAT_STRICT_HEADERS) { + format.strict_headers = parse_bool(file_format::FORMAT_STRICT_HEADERS, strict_headers)?; + }; + if format.strict_headers && !format.has_header { + return error::ParseFormatSnafu { + key: file_format::FORMAT_STRICT_HEADERS, + value: "strict_headers=true requires headers=true", + } + .fail(); + } if let Some(timestamp_format) = value.get(file_format::TIMESTAMP_FORMAT) { format.timestamp_format = Some(timestamp_format.clone()); } @@ -135,6 +146,7 @@ impl Default for CsvFormat { Self { has_header: true, skip_bad_records: false, + strict_headers: false, delimiter: b',', schema_infer_max_record: Some(file_format::DEFAULT_SCHEMA_INFER_MAX_RECORD), compression_type: CompressionType::Uncompressed, @@ -365,7 +377,8 @@ mod tests { use super::*; use crate::file_format::{ FORMAT_COMPRESSION_TYPE, FORMAT_DELIMITER, FORMAT_HAS_HEADER, FORMAT_HEADERS, - FORMAT_SCHEMA_INFER_MAX_RECORD, FORMAT_SKIP_BAD_RECORDS, FileFormat, file_to_stream, + FORMAT_SCHEMA_INFER_MAX_RECORD, FORMAT_SKIP_BAD_RECORDS, FORMAT_STRICT_HEADERS, FileFormat, + file_to_stream, }; use crate::test_util::{format_schema, test_store}; @@ -492,6 +505,7 @@ mod tests { delimiter: b'\t', has_header: false, skip_bad_records: false, + strict_headers: false, timestamp_format: None, time_format: None, date_format: None @@ -509,6 +523,17 @@ mod tests { } ); + let map = HashMap::from([(FORMAT_STRICT_HEADERS.to_string(), "true".to_string())]); + let format = CsvFormat::try_from(&map).unwrap(); + + assert_eq!( + format, + CsvFormat { + strict_headers: true, + ..CsvFormat::default() + } + ); + let map = HashMap::from([(FORMAT_HEADERS.to_string(), "true".to_string())]); let format = CsvFormat::try_from(&map).unwrap(); assert_eq!(format, CsvFormat::default()); @@ -544,6 +569,9 @@ mod tests { let map = HashMap::from([(FORMAT_HEADERS.to_string(), "yes".to_string())]); assert!(CsvFormat::try_from(&map).is_err()); + + let map = HashMap::from([(FORMAT_STRICT_HEADERS.to_string(), "yes".to_string())]); + assert!(CsvFormat::try_from(&map).is_err()); } #[test] @@ -555,6 +583,15 @@ mod tests { assert!(CsvFormat::try_from(&map).is_err()); } + #[test] + fn test_try_from_rejects_strict_headerless_csv() { + let map = HashMap::from([ + (FORMAT_HEADERS.to_string(), "false".to_string()), + (FORMAT_STRICT_HEADERS.to_string(), "true".to_string()), + ]); + assert!(CsvFormat::try_from(&map).is_err()); + } + #[tokio::test] async fn test_compressed_csv() { // Create test data diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 5134c40a804..90e631679a7 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -92,6 +92,7 @@ tokio-postgres-rustls = { version = "0.12", optional = true } tonic.workspace = true tracing.workspace = true typetag.workspace = true +uuid.workspace = true [dev-dependencies] chrono.workspace = true @@ -100,4 +101,3 @@ common-procedure-test.workspace = true common-test-util.workspace = true common-wal = { workspace = true, features = ["testing"] } datatypes.workspace = true -uuid.workspace = true diff --git a/src/common/meta/src/ddl.rs b/src/common/meta/src/ddl.rs index cc578b7552f..7f168e921e5 100644 --- a/src/common/meta/src/ddl.rs +++ b/src/common/meta/src/ddl.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::sync::Arc; +use std::time::Duration; use store_api::storage::{RegionId, TableId}; @@ -121,6 +122,8 @@ pub struct DdlContext { pub region_failure_detector_controller: RegionFailureDetectorControllerRef, /// Whether table drops should stop after tombstoning metadata. pub soft_drop_enabled: bool, + /// Fixed retention used to calculate new soft-drop deadlines. + pub soft_drop_retention: Option, } impl DdlContext { diff --git a/src/common/meta/src/ddl/drop_table.rs b/src/common/meta/src/ddl/drop_table.rs index f5b5bc7a22a..f3f085d1f24 100644 --- a/src/common/meta/src/ddl/drop_table.rs +++ b/src/common/meta/src/ddl/drop_table.rs @@ -26,6 +26,7 @@ use common_procedure::{ }; use common_telemetry::info; use common_telemetry::tracing::warn; +use common_time::util::current_time_millis; use common_wal::options::WalOptions; use serde::{Deserialize, Serialize}; use snafu::{OptionExt, ResultExt}; @@ -33,13 +34,14 @@ use store_api::storage::RegionNumber; use strum::AsRefStr; use table::metadata::TableId; use table::table_reference::TableReference; +use uuid::Uuid; use self::executor::DropTableExecutor; use crate::ddl::DdlContext; use crate::ddl::utils::{convert_region_routes_to_detecting_regions, map_to_procedure_error}; use crate::error::{self, Result}; use crate::key::table_route::TableRouteValue; -use crate::lock_key::{CatalogLock, SchemaLock, TableLock}; +use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; use crate::metrics; use crate::region_keeper::OperatingRegionGuard; use crate::rpc::ddl::DropTableTask; @@ -60,7 +62,7 @@ impl DropTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::DropTable"; pub fn new(task: DropTableTask, context: DdlContext) -> Self { - let data = DropTableData::new(task); + let data = DropTableData::new(task, context.soft_drop_enabled, context.soft_drop_retention); let executor = data.build_executor(); Self { context, @@ -71,7 +73,16 @@ impl DropTableProcedure { } pub fn from_json(json: &str, context: DdlContext) -> ProcedureResult { - let data: DropTableData = serde_json::from_str(json).context(FromJsonSnafu)?; + let mut data: DropTableData = serde_json::from_str(json).context(FromJsonSnafu)?; + if data.state == DropTableState::Prepare + && data.soft_drop_enabled + && data.dropped_at.is_none() + && data.soft_drop_retention_millis.is_none() + { + data.soft_drop_retention_millis = context + .soft_drop_retention + .and_then(|retention| i64::try_from(retention.as_millis()).ok()); + } let executor = data.build_executor(); Ok(Self { @@ -83,10 +94,36 @@ impl DropTableProcedure { } pub(crate) async fn on_prepare(&mut self) -> Result { - if self.executor.on_prepare(&self.context).await?.stop() { + if self + .executor + .on_prepare(&self.context, self.data.soft_drop_enabled) + .await? + .stop() + { return Ok(Status::done()); } self.fill_table_metadata().await?; + if self.data.soft_drop_enabled && self.data.dropped_at.is_none() { + let dropped_at = current_time_millis(); + let retention_millis = + self.data + .soft_drop_retention_millis + .context(error::UnexpectedSnafu { + err_msg: "Soft-drop retention is missing from the procedure".to_string(), + })?; + let retention_expires_at = dropped_at.checked_add(retention_millis).context( + error::UnexpectedSnafu { + err_msg: format!( + "Soft-drop retention deadline overflows i64: dropped_at={dropped_at}, retention_millis={retention_millis}" + ), + }, + )?; + self.data.dropped_at = Some(dropped_at); + self.data.retention_expires_at = Some(retention_expires_at); + } + if self.data.soft_drop_enabled && self.data.drop_generation.is_none() { + self.data.drop_generation = Some(Uuid::new_v4().to_string()); + } self.data.state = DropTableState::DeleteMetadata; Ok(Status::executing(true)) @@ -139,6 +176,9 @@ impl DropTableProcedure { &self.context, table_route_value, &self.data.region_wal_options, + self.data.dropped_at, + self.data.retention_expires_at, + self.data.drop_generation.as_deref(), ) .await?; info!("Deleted table metadata for table {table_id}"); @@ -173,12 +213,13 @@ impl DropTableProcedure { .await?; } - if self.context.soft_drop_enabled { + if self.data.soft_drop_enabled { self.executor .on_close_regions( &self.context.node_manager, &self.context.leader_region_registry, &self.data.physical_region_routes, + true, ) .await?; self.context @@ -285,6 +326,7 @@ impl Procedure for DropTableProcedure { let lock_key = vec![ CatalogLock::Read(table_ref.catalog).into(), SchemaLock::read(table_ref.catalog, table_ref.schema).into(), + TableNameLock::new(table_ref.catalog, table_ref.schema, table_ref.table).into(), TableLock::Write(table_id).into(), ]; @@ -328,10 +370,24 @@ pub struct DropTableData { pub region_wal_options: HashMap, #[serde(default)] pub allow_rollback: bool, + #[serde(default)] + pub soft_drop_enabled: bool, + #[serde(default)] + pub dropped_at: Option, + #[serde(default)] + pub retention_expires_at: Option, + #[serde(default)] + pub soft_drop_retention_millis: Option, + #[serde(default)] + pub drop_generation: Option, } impl DropTableData { - pub fn new(task: DropTableTask) -> Self { + pub fn new( + task: DropTableTask, + soft_drop_enabled: bool, + soft_drop_retention: Option, + ) -> Self { Self { state: DropTableState::Prepare, task, @@ -339,6 +395,12 @@ impl DropTableData { physical_table_id: None, region_wal_options: HashMap::new(), allow_rollback: false, + soft_drop_enabled, + dropped_at: None, + retention_expires_at: None, + soft_drop_retention_millis: soft_drop_retention + .and_then(|retention| i64::try_from(retention.as_millis()).ok()), + drop_generation: None, } } diff --git a/src/common/meta/src/ddl/drop_table/executor.rs b/src/common/meta/src/ddl/drop_table/executor.rs index 1e85a64ef60..61e74d1b5cf 100644 --- a/src/common/meta/src/ddl/drop_table/executor.rs +++ b/src/common/meta/src/ddl/drop_table/executor.rs @@ -15,8 +15,8 @@ use std::collections::{HashMap, HashSet}; use api::v1::region::{ - CloseRequest as PbCloseRegionRequest, DropRequest as PbDropRegionRequest, RegionRequest, - RegionRequestHeader, region_request, + CleanUpRequest as PbCleanUpRequest, CloseRequest as PbCloseRegionRequest, + DropRequest as PbDropRegionRequest, RegionRequest, RegionRequestHeader, region_request, }; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; @@ -26,14 +26,17 @@ use common_wal::options::WalOptions; use futures::future::join_all; use snafu::ensure; use store_api::storage::{RegionId, RegionNumber}; -use table::metadata::TableId; +use table::metadata::{TableId, TableInfo}; use table::table_name::TableName; use crate::cache_invalidator::Context; -use crate::ddl::DdlContext; -use crate::ddl::utils::{add_peer_context_if_needed, convert_region_routes_to_detecting_regions}; +use crate::ddl::utils::{ + add_peer_context_if_needed, convert_region_routes_to_detecting_regions, region_storage_path, +}; +use crate::ddl::{CreateRequestBuilder, DdlContext, build_template_from_raw_table_info}; use crate::error::{self, Result}; use crate::instruction::CacheIdent; +use crate::key::DroppedTableLifecycle; use crate::key::table_name::TableNameKey; use crate::key::table_route::TableRouteValue; use crate::node_manager::NodeManagerRef; @@ -84,7 +87,11 @@ impl DropTableExecutor { /// - Throws an error if table not exists and `drop_if_exists` is `false`. /// - Rejects dropping a recreated live table while an older tombstone still owns the same /// fully qualified name. - pub async fn on_prepare(&self, ctx: &DdlContext) -> Result> { + pub async fn on_prepare( + &self, + ctx: &DdlContext, + soft_drop_enabled: bool, + ) -> Result> { let table_ref = self.table.table_ref(); let exist = ctx @@ -108,12 +115,12 @@ impl DropTableExecutor { } ); - if ctx.soft_drop_enabled - && let Some(dropped_table) = ctx - .table_metadata_manager - .get_dropped_table(&self.table) - .await? + if let Some(dropped_table) = ctx + .table_metadata_manager + .get_dropped_table(&self.table) + .await? && dropped_table.table_id != self.table_id + && (soft_drop_enabled || dropped_table.dropped_at.is_some()) { return error::TableNameTombstoneConflictSnafu { table_name: table_ref.to_string(), @@ -132,13 +139,21 @@ impl DropTableExecutor { ctx: &DdlContext, table_route_value: &TableRouteValue, region_wal_options: &HashMap, + dropped_at: Option, + retention_expires_at: Option, + drop_generation: Option<&str>, ) -> Result<()> { ctx.table_metadata_manager - .delete_table_metadata( + .delete_table_metadata_with_retention_and_generation( self.table_id, &self.table, table_route_value, region_wal_options, + DroppedTableLifecycle { + dropped_at, + retention_expires_at, + drop_generation, + }, ) .await } @@ -313,6 +328,7 @@ impl DropTableExecutor { }), body: Some(region_request::Body::Close(PbCloseRegionRequest { region_id: region_id.as_u64(), + flush_on_close: false, })), }; @@ -347,43 +363,53 @@ impl DropTableExecutor { Ok(()) } - /// Closes all table regions on datanodes without deleting region files or metadata tombstones. - pub async fn on_close_regions( + /// Cleans leader regions on datanodes without reopening them as live regions. + pub async fn on_cleanup_regions_offline( &self, node_manager: &NodeManagerRef, leader_region_registry: &LeaderRegionRegistryRef, + table_info: &TableInfo, region_routes: &[RegionRoute], + region_wal_options: &HashMap, ) -> Result<()> { - let table_id = self.table_id; - let mut seen_peer_ids = HashSet::new(); - let peers = find_leaders(region_routes) - .into_iter() - .chain(find_followers(region_routes)) - .filter(|peer| seen_peer_ids.insert(peer.id)); - let mut close_region_tasks = Vec::new(); + let template = build_template_from_raw_table_info(table_info)?; + let builder = CreateRequestBuilder::new(template, None); + let storage_path = region_storage_path(&self.table.catalog_name, &self.table.schema_name); - for datanode in peers { + let leaders = find_leaders(region_routes); + let mut cleanup_region_tasks = Vec::with_capacity(leaders.len()); + let table_id = self.table_id; + for datanode in leaders { let requester = node_manager.datanode(&datanode).await; - let region_ids = find_leader_regions(region_routes, &datanode) - .into_iter() - .chain(find_follower_regions(region_routes, &datanode)) - .map(|region_number| RegionId::new(table_id, region_number)); + let regions = find_leader_regions(region_routes, &datanode); + let region_ids = regions + .iter() + .map(|region_number| RegionId::new(table_id, *region_number)) + .collect::>(); for region_id in region_ids { - debug!("Closing region {region_id} on Datanode {datanode:?}"); + debug!("Cleaning region {region_id} offline on Datanode {datanode:?}"); + let create_request = builder.build_one( + region_id, + storage_path.clone(), + region_wal_options, + &HashMap::new(), + )?; let request = RegionRequest { header: Some(RegionRequestHeader { tracing_context: TracingContext::from_current_span().to_w3c(), ..Default::default() }), - body: Some(region_request::Body::Close(PbCloseRegionRequest { - region_id: region_id.as_u64(), + body: Some(region_request::Body::CleanUp(PbCleanUpRequest { + region_id: create_request.region_id, + engine: create_request.engine, + path: create_request.path, + options: create_request.options, })), }; - let datanode = datanode.clone(); let requester = requester.clone(); - close_region_tasks.push(async move { + cleanup_region_tasks.push(async move { if let Err(err) = requester.handle(request).await && err.status_code() != StatusCode::RegionNotFound { @@ -394,6 +420,84 @@ impl DropTableExecutor { } } + join_all(cleanup_region_tasks) + .await + .into_iter() + .collect::>>()?; + + let region_ids = operating_leader_regions(region_routes); + leader_region_registry.batch_delete(region_ids.into_iter().map(|(region_id, _)| region_id)); + + Ok(()) + } + + /// Closes all table regions on datanodes without deleting region files or metadata tombstones. + /// When `flush_leaders_on_close` is set, only leader regions are flushed before close. + pub async fn on_close_regions( + &self, + node_manager: &NodeManagerRef, + leader_region_registry: &LeaderRegionRegistryRef, + region_routes: &[RegionRoute], + flush_leaders_on_close: bool, + ) -> Result<()> { + let table_id = self.table_id; + let mut seen_peer_ids = HashSet::new(); + let peers = find_leaders(region_routes) + .into_iter() + .chain(find_followers(region_routes)) + .filter(|peer| seen_peer_ids.insert(peer.id)); + let close_region_tasks = peers.map(|datanode| { + let region_ids = find_leader_regions(region_routes, &datanode) + .into_iter() + .map(|region_number| { + ( + RegionId::new(table_id, region_number), + flush_leaders_on_close, + ) + }) + .chain( + find_follower_regions(region_routes, &datanode) + .into_iter() + .map(|region_number| (RegionId::new(table_id, region_number), false)), + ) + .collect::>(); + + async move { + let requester = node_manager.datanode(&datanode).await; + let close_region_tasks = + region_ids.into_iter().map(|(region_id, flush_on_close)| { + debug!("Closing region {region_id} on Datanode {datanode:?}"); + let request = RegionRequest { + header: Some(RegionRequestHeader { + tracing_context: TracingContext::from_current_span().to_w3c(), + ..Default::default() + }), + body: Some(region_request::Body::Close(PbCloseRegionRequest { + region_id: region_id.as_u64(), + flush_on_close, + })), + }; + + let datanode = datanode.clone(); + let requester = requester.clone(); + async move { + if let Err(err) = requester.handle(request).await + && err.status_code() != StatusCode::RegionNotFound + { + return Err(add_peer_context_if_needed(datanode)(err)); + } + Ok(()) + } + }); + + join_all(close_region_tasks) + .await + .into_iter() + .collect::>>()?; + Ok(()) + } + }); + join_all(close_region_tasks) .await .into_iter() @@ -469,7 +573,7 @@ mod tests { 1024, true, ); - let ctrl = executor.on_prepare(&ctx).await.unwrap(); + let ctrl = executor.on_prepare(&ctx, false).await.unwrap(); assert!(ctrl.stop()); // Drops a non-exists table @@ -478,7 +582,7 @@ mod tests { 1024, false, ); - let err = executor.on_prepare(&ctx).await.unwrap_err(); + let err = executor.on_prepare(&ctx, false).await.unwrap_err(); assert_matches!(err, error::Error::TableNotFound { .. }); // Drops a exists table @@ -496,7 +600,7 @@ mod tests { ) .await .unwrap(); - let ctrl = executor.on_prepare(&ctx).await.unwrap(); + let ctrl = executor.on_prepare(&ctx, false).await.unwrap(); assert!(!ctrl.stop()); } } diff --git a/src/common/meta/src/ddl/drop_table/metadata.rs b/src/common/meta/src/ddl/drop_table/metadata.rs index 6289b6c61ab..213ebb461b3 100644 --- a/src/common/meta/src/ddl/drop_table/metadata.rs +++ b/src/common/meta/src/ddl/drop_table/metadata.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use common_catalog::consts::FILE_ENGINE; use common_catalog::format_full_table_name; use snafu::{OptionExt, ensure}; use store_api::metric_engine_consts::METRIC_ENGINE_NAME; @@ -49,7 +50,7 @@ impl DropTableProcedure { ); // TODO(hl): support soft-dropping logical tables. ensure!( - !(self.context.soft_drop_enabled + !(self.data.soft_drop_enabled && is_metric_engine_logical_table( &table_info_value.table_info, &table_route_value @@ -61,6 +62,12 @@ impl DropTableProcedure { if physical_table_id == self.data.table_id() { let engine = table_info_value.table_info.meta.engine; + ensure!( + !(self.data.soft_drop_enabled && engine == FILE_ENGINE), + error::UnsupportedSnafu { + operation: "soft-dropping file-engine tables".to_string() + } + ); // rollback only if dropping the metric physical table fails self.data.allow_rollback = engine.as_str() == METRIC_ENGINE_NAME; diff --git a/src/common/meta/src/ddl/purge_dropped_table.rs b/src/common/meta/src/ddl/purge_dropped_table.rs index b95e5db976d..040a0d8e60d 100644 --- a/src/common/meta/src/ddl/purge_dropped_table.rs +++ b/src/common/meta/src/ddl/purge_dropped_table.rs @@ -19,6 +19,7 @@ use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ Context as ProcedureContext, LockKey, Procedure, Result as ProcedureResult, Status, }; +use common_time::util::current_time_millis; use common_wal::options::WalOptions; use serde::{Deserialize, Serialize}; use snafu::{ResultExt, ensure}; @@ -29,12 +30,12 @@ use table::table_name::TableName; use crate::ddl::DdlContext; use crate::ddl::drop_table::executor::DropTableExecutor; -use crate::ddl::undrop_table::open_regions_ignore_region_not_found; use crate::ddl::utils::{ convert_region_routes_to_detecting_regions, is_metric_engine_logical_table, map_to_procedure_error, }; use crate::error::{self, Result}; +use crate::key::DroppedTableMetadata; use crate::key::table_route::TableRouteValue; use crate::lock_key::TableLock; use crate::rpc::ddl::PurgeDroppedTableTask; @@ -47,6 +48,7 @@ pub struct PurgeDroppedTableProcedure { impl PurgeDroppedTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::PurgeDroppedTable"; + pub const EXPIRED_TYPE_NAME: &'static str = "metasrv-procedure::PurgeExpiredDroppedTable"; pub fn new(task: PurgeDroppedTableTask, context: DdlContext) -> Self { Self { @@ -55,12 +57,19 @@ impl PurgeDroppedTableProcedure { } } + pub fn new_if_expired(task: PurgeDroppedTableTask, context: DdlContext) -> Self { + Self { + context, + data: PurgeDroppedTableData::new_if_expired(task), + } + } + pub fn from_json(json: &str, context: DdlContext) -> ProcedureResult { let data: PurgeDroppedTableData = serde_json::from_str(json).context(FromJsonSnafu)?; Ok(Self { context, data }) } - async fn on_prepare(&mut self) -> Result { + async fn eligible_dropped_table(&self) -> Result> { let dropped_table = self .context .table_metadata_manager @@ -68,8 +77,25 @@ impl PurgeDroppedTableProcedure { .await?; let Some(dropped_table) = dropped_table else { - return Ok(Status::done()); + return Ok(None); }; + if self.data.check_expired { + if !self.context.soft_drop_enabled { + return Ok(None); + } + let retention_millis = self + .context + .soft_drop_retention + .and_then(|retention| i64::try_from(retention.as_millis()).ok()); + let expires_at = dropped_table.retention_expires_at.or_else(|| { + dropped_table.dropped_at.and_then(|dropped_at| { + retention_millis.and_then(|retention| dropped_at.checked_add(retention)) + }) + }); + if expires_at.is_none_or(|expires_at| expires_at > current_time_millis()) { + return Ok(None); + } + } ensure!( !is_metric_engine_logical_table( &dropped_table.table_info_value.table_info, @@ -79,41 +105,66 @@ impl PurgeDroppedTableProcedure { operation: "purging metric logical tables".to_string() } ); + Ok(Some(dropped_table)) + } + + fn update_dropped_table(&mut self, dropped_table: DroppedTableMetadata) { self.data.table_id = Some(dropped_table.table_id); self.data.table_name = Some(dropped_table.table_name); self.data.table_info = Some(dropped_table.table_info_value.table_info); self.data.table_route_value = Some(dropped_table.table_route_value); self.data.region_wal_options = dropped_table.region_wal_options; - self.data.state = PurgeDroppedTableState::OpenRegions; - Ok(Status::executing(true)) + self.data.drop_generation = dropped_table.drop_generation; + self.data.drop_generation_loaded = true; } - async fn on_open_regions(&mut self) -> Result { - if let Some(region_routes) = self.data.physical_region_routes() { - open_regions_ignore_region_not_found( - &self.context, - self.data.table_id(), - self.data.table_name(), - self.data.table_info(), - region_routes, - &self.data.region_wal_options, - ) - .await?; - } + async fn on_prepare(&mut self) -> Result { + let Some(dropped_table) = self.eligible_dropped_table().await? else { + return Ok(Status::done()); + }; + self.update_dropped_table(dropped_table); self.data.state = PurgeDroppedTableState::DropRegions; Ok(Status::executing(true)) } async fn on_drop_regions(&mut self) -> Result { + let expected_generation = self.data.drop_generation.as_deref().unwrap_or_default(); + let purge_claim = self + .context + .table_metadata_manager + .dropped_table_purge_claim(self.data.task.table_id) + .await?; + if let Some(purge_claim) = purge_claim { + if purge_claim != expected_generation { + return Ok(Status::done()); + } + self.data.purging_claimed = true; + } else { + // Revalidate after recovery before crossing the durable cleanup boundary. + let Some(dropped_table) = self.eligible_dropped_table().await? else { + return Ok(Status::done()); + }; + if dropped_table.drop_generation != self.data.drop_generation { + return Ok(Status::done()); + } + self.update_dropped_table(dropped_table); + self.context + .table_metadata_manager + .mark_dropped_table_purging( + self.data.task.table_id, + self.data.drop_generation.as_deref(), + ) + .await?; + self.data.purging_claimed = true; + } if let Some(region_routes) = self.data.physical_region_routes() { self.executor() - .on_drop_regions( + .on_cleanup_regions_offline( &self.context.node_manager, &self.context.leader_region_registry, + self.data.table_info(), region_routes, - false, - false, - false, + &self.data.region_wal_options, ) .await?; self.context @@ -127,6 +178,19 @@ impl PurgeDroppedTableProcedure { } async fn on_delete_tombstone(&mut self) -> Result { + if self.data.purging_claimed { + let expected_generation = self.data.drop_generation.as_deref().unwrap_or_default(); + if self + .context + .table_metadata_manager + .dropped_table_purge_claim(self.data.task.table_id) + .await? + .as_deref() + != Some(expected_generation) + { + return Ok(Status::done()); + } + } self.context .table_metadata_manager .delete_table_metadata_tombstone( @@ -147,17 +211,25 @@ impl PurgeDroppedTableProcedure { #[async_trait] impl Procedure for PurgeDroppedTableProcedure { fn type_name(&self) -> &str { - Self::TYPE_NAME + if self.data.check_expired { + Self::EXPIRED_TYPE_NAME + } else { + Self::TYPE_NAME + } } fn recover(&mut self) -> ProcedureResult<()> { + if self.data.state == PurgeDroppedTableState::DropRegions + && !self.data.drop_generation_loaded + { + self.data.state = PurgeDroppedTableState::Prepare; + } Ok(()) } async fn execute(&mut self, _: &ProcedureContext) -> ProcedureResult { match self.data.state { PurgeDroppedTableState::Prepare => self.on_prepare().await, - PurgeDroppedTableState::OpenRegions => self.on_open_regions().await, PurgeDroppedTableState::DropRegions => self.on_drop_regions().await, PurgeDroppedTableState::DeleteTombstone => self.on_delete_tombstone().await, } @@ -183,6 +255,14 @@ pub struct PurgeDroppedTableData { table_route_value: Option, #[serde(default)] region_wal_options: HashMap, + #[serde(default)] + check_expired: bool, + #[serde(default)] + drop_generation: Option, + #[serde(default)] + drop_generation_loaded: bool, + #[serde(default)] + purging_claimed: bool, } impl PurgeDroppedTableData { @@ -195,6 +275,17 @@ impl PurgeDroppedTableData { table_info: None, table_route_value: None, region_wal_options: HashMap::new(), + check_expired: false, + drop_generation: None, + drop_generation_loaded: false, + purging_claimed: false, + } + } + + fn new_if_expired(task: PurgeDroppedTableTask) -> Self { + Self { + check_expired: true, + ..Self::new(task) } } @@ -206,14 +297,14 @@ impl PurgeDroppedTableData { self.table_name.as_ref().unwrap() } - fn table_info(&self) -> &TableInfo { - self.table_info.as_ref().unwrap() - } - fn table_route_value(&self) -> &TableRouteValue { self.table_route_value.as_ref().unwrap() } + fn table_info(&self) -> &TableInfo { + self.table_info.as_ref().unwrap() + } + fn physical_region_routes(&self) -> Option<&[RegionRoute]> { match self.table_route_value() { TableRouteValue::Physical(route) => Some(&route.region_routes), @@ -225,7 +316,6 @@ impl PurgeDroppedTableData { #[derive(Debug, Serialize, Deserialize, AsRefStr, PartialEq)] enum PurgeDroppedTableState { Prepare, - OpenRegions, DropRegions, DeleteTombstone, } diff --git a/src/common/meta/src/ddl/tests/drop_table.rs b/src/common/meta/src/ddl/tests/drop_table.rs index fd41dd5126d..dd960df75c4 100644 --- a/src/common/meta/src/ddl/tests/drop_table.rs +++ b/src/common/meta/src/ddl/tests/drop_table.rs @@ -14,19 +14,18 @@ use std::assert_matches; use std::collections::{HashMap, HashSet}; -use std::sync::{Arc, Mutex as StdMutex}; +use std::sync::Arc; +use std::time::Duration; -use api::region::RegionResponse; use api::v1::region::{RegionRequest, region_request}; use async_trait::async_trait; -use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; -use common_error::ext::{BoxedError, ErrorExt, StackError}; +use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, FILE_ENGINE}; +use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; -use common_procedure::Procedure; +use common_procedure::{Procedure, StringKey}; use common_procedure_test::{ execute_procedure_until, execute_procedure_until_done, new_test_procedure_context, }; -use snafu::ResultExt; use store_api::region_engine::RegionRole; use store_api::storage::RegionId; use table::metadata::TableId; @@ -43,20 +42,297 @@ use crate::ddl::test_util::{ put_datanode_address, test_create_logical_table_task, test_create_physical_table_task, }; use crate::ddl::undrop_table::UndropTableProcedure; -use crate::ddl::{DetectingRegion, RegionFailureDetectorController, TableMetadata}; -use crate::error::{self, Error}; +use crate::ddl::{DdlContext, DetectingRegion, RegionFailureDetectorController, TableMetadata}; +use crate::error::Error; +use crate::key::MetadataKey; use crate::key::table_name::TableNameKey; use crate::key::table_route::TableRouteValue; +use crate::kv_backend::KvBackend; use crate::kv_backend::memory::MemoryKvBackend; use crate::peer::Peer; use crate::rpc::ddl::{DropTableTask, PurgeDroppedTableTask, UndropTableTask}; use crate::rpc::router::{Region, RegionRoute}; +use crate::rpc::store::{BatchDeleteRequest, PutRequest}; use crate::test_util::{MockDatanodeManager, new_ddl_context, new_ddl_context_with_kv_backend}; +fn dropped_at_marker_key(table_id: TableId) -> String { + format!("__tombstone/__dropped_at/{table_id}") +} + +fn retention_expires_at_marker_key(table_id: TableId) -> String { + format!("__tombstone/__retention_expires_at/{table_id}") +} + +fn purging_marker_key(table_id: TableId) -> String { + format!("__tombstone/__purging/{table_id}") +} + +fn drop_generation_marker_key(table_id: TableId) -> String { + format!("__tombstone/__drop_generation/{table_id}") +} + +async fn create_dropped_table( + context: &DdlContext, + table_id: TableId, + dropped_at: Option, + retention_expires_at: Option, +) { + let task = test_create_table_task("foo", table_id); + let table_name = task.table_name(); + let table_route = TableRouteValue::physical(vec![]); + context + .table_metadata_manager + .create_table_metadata(task.table_info, table_route.clone(), HashMap::new()) + .await + .unwrap(); + context + .table_metadata_manager + .delete_table_metadata_with_retention( + table_id, + &table_name, + &table_route, + &HashMap::new(), + dropped_at, + retention_expires_at, + ) + .await + .unwrap(); +} + +async fn undrop_at_restore_metadata( + context: &DdlContext, + table_id: TableId, +) -> UndropTableProcedure { + create_dropped_table(context, table_id, Some(1), Some(i64::MAX)).await; + let mut undrop = UndropTableProcedure::new(new_undrop_table_task(table_id), context.clone()); + let procedure_context = new_test_procedure_context(); + undrop.execute(&procedure_context).await.unwrap(); + undrop.execute(&procedure_context).await.unwrap(); + undrop +} + +fn legacy_undrop_snapshot(procedure: &UndropTableProcedure) -> String { + let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + let data = data.as_object_mut().unwrap(); + for field in [ + "dropped_at", + "retention_expires_at", + "drop_generation", + "tombstone_identity_loaded", + ] { + data.remove(field); + } + serde_json::to_string(data).unwrap() +} + +#[test] +fn test_old_drop_table_json_defaults_to_hard_drop() { + let mut runtime_context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + runtime_context.soft_drop_enabled = true; + let procedure = DropTableProcedure::new( + new_drop_table_task("foo", 1024, false), + runtime_context.clone(), + ); + let mut old_data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + old_data + .as_object_mut() + .unwrap() + .remove("soft_drop_enabled"); + old_data.as_object_mut().unwrap().remove("dropped_at"); + + let recovered = DropTableProcedure::from_json(&old_data.to_string(), runtime_context).unwrap(); + let recovered_data: serde_json::Value = + serde_json::from_str(&recovered.dump().unwrap()).unwrap(); + + assert_eq!(recovered_data["soft_drop_enabled"], false); + assert_eq!(recovered_data["dropped_at"], serde_json::Value::Null); +} + +#[test] +fn test_drop_table_lock_key_includes_table_name() { + let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); + + let keys = procedure + .lock_key() + .keys_to_lock() + .cloned() + .collect::>(); + assert!( + keys.iter().any(|key| matches!( + key, + StringKey::Exclusive(key) if key == "__table_name_lock/greptime.public.foo" + )), + "drop lock keys should include the table name: {keys:?}" + ); + assert!( + keys.iter().any(|key| matches!( + key, + StringKey::Exclusive(key) if key == "__table_lock/1024" + )), + "drop lock keys should include the table id: {keys:?}" + ); +} + +#[tokio::test] +async fn test_recovered_soft_drop_preserves_persisted_mode() { + let (tx, mut rx) = mpsc::channel(8); + let node_manager = Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(tx))); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let mut soft_context = + new_ddl_context_with_kv_backend(node_manager.clone(), kv_backend.clone()); + soft_context.soft_drop_enabled = true; + let table_id = 1024; + let task = test_create_table_task("foo", table_id); + soft_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(RegionId::new(table_id, 1)), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }]), + HashMap::new(), + ) + .await + .unwrap(); + let procedure = DropTableProcedure::new( + new_drop_table_task("foo", table_id, false), + soft_context.clone(), + ); + + let hard_context = new_ddl_context_with_kv_backend(node_manager, kv_backend); + let mut recovered = + DropTableProcedure::from_json(&procedure.dump().unwrap(), hard_context).unwrap(); + execute_procedure_until_done(&mut recovered).await; + + assert!( + soft_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_some() + ); + let (_, request) = rx.try_recv().unwrap(); + assert_matches!(request.body, Some(region_request::Body::Close(_))); +} + +#[tokio::test] +async fn test_legacy_soft_drop_prepare_preserves_persisted_mode() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let mut context = new_ddl_context(node_manager); + context.soft_drop_enabled = false; + context.soft_drop_retention = Some(Duration::from_millis(100)); + let table_id = 1024; + let task = test_create_table_task("foo", table_id); + context + .table_metadata_manager + .create_table_metadata( + task.table_info, + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + + let mut persisted: serde_json::Value = serde_json::from_str( + &DropTableProcedure::new(new_drop_table_task("foo", table_id, false), context.clone()) + .dump() + .unwrap(), + ) + .unwrap(); + persisted["soft_drop_enabled"] = true.into(); + persisted + .as_object_mut() + .unwrap() + .remove("soft_drop_retention_millis"); + persisted + .as_object_mut() + .unwrap() + .remove("retention_expires_at"); + + let mut recovered = + DropTableProcedure::from_json(&persisted.to_string(), context.clone()).unwrap(); + recovered.on_prepare().await.unwrap(); + + assert!(recovered.data.soft_drop_enabled); + assert_eq!(Some(100), recovered.data.soft_drop_retention_millis); + assert!(recovered.data.dropped_at.is_some()); + assert!(recovered.data.retention_expires_at.is_some()); +} + +#[tokio::test] +async fn test_disabled_soft_drop_creates_hard_drop_procedure() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let mut context = new_ddl_context(node_manager); + context.soft_drop_enabled = false; + context.soft_drop_retention = Some(Duration::from_millis(100)); + + let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); + + assert!(!procedure.data.soft_drop_enabled); + assert_eq!(None, procedure.data.dropped_at); + assert_eq!(None, procedure.data.retention_expires_at); +} + +#[tokio::test] +async fn test_recovered_hard_drop_ignores_soft_runtime_context() { + let (tx, mut rx) = mpsc::channel(8); + let node_manager = Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(tx))); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let hard_context = new_ddl_context_with_kv_backend(node_manager.clone(), kv_backend.clone()); + let table_id = 1024; + let task = test_create_table_task("foo", table_id); + hard_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(RegionId::new(table_id, 1)), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }]), + HashMap::new(), + ) + .await + .unwrap(); + let procedure = DropTableProcedure::new( + new_drop_table_task("foo", table_id, false), + hard_context.clone(), + ); + + let mut soft_context = new_ddl_context_with_kv_backend(node_manager, kv_backend); + soft_context.soft_drop_enabled = true; + let mut recovered = + DropTableProcedure::from_json(&procedure.dump().unwrap(), soft_context).unwrap(); + execute_procedure_until_done(&mut recovered).await; + + assert!( + hard_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_none() + ); + let (_, request) = rx.try_recv().unwrap(); + assert_matches!(request.body, Some(region_request::Body::Drop(_))); +} + #[tokio::test] async fn test_on_prepare_table_not_exists_err() { let node_manager = Arc::new(MockDatanodeManager::new(())); - let ddl_context = new_ddl_context(node_manager); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + ddl_context.soft_drop_retention = Some(Duration::from_millis(100)); let table_name = "foo"; let table_id = 1024; let task = test_create_table_task(table_name, table_id); @@ -73,8 +349,86 @@ async fn test_on_prepare_table_not_exists_err() { let task = new_drop_table_task("bar", table_id, false); let mut procedure = DropTableProcedure::new(task, ddl_context); + assert_eq!(procedure.data.dropped_at, None); + assert_eq!(procedure.data.retention_expires_at, None); let err = procedure.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::TableNotFound); + assert_eq!(procedure.data.dropped_at, None); + assert_eq!(procedure.data.retention_expires_at, None); +} + +#[tokio::test] +async fn test_soft_drop_prepare_assigns_stable_deadline_and_recovers_it() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + ddl_context.soft_drop_retention = Some(Duration::from_millis(100)); + let table_id = 1024; + let task = test_create_table_task("foo", table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + let mut procedure = DropTableProcedure::new( + new_drop_table_task("foo", table_id, false), + ddl_context.clone(), + ); + + procedure.on_prepare().await.unwrap(); + let dropped_at = procedure.data.dropped_at.unwrap(); + let deadline = procedure.data.retention_expires_at.unwrap(); + assert_eq!(deadline, dropped_at + 100); + + let mut changed_context = ddl_context; + changed_context.soft_drop_retention = Some(Duration::from_secs(10)); + let recovered = + DropTableProcedure::from_json(&procedure.dump().unwrap(), changed_context).unwrap(); + assert_eq!(recovered.data.dropped_at, Some(dropped_at)); + assert_eq!(recovered.data.retention_expires_at, Some(deadline)); +} + +#[tokio::test] +async fn test_soft_drop_prepare_rejects_deadline_overflow_before_metadata_delete() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + ddl_context.soft_drop_retention = Some(Duration::from_millis(i64::MAX as u64)); + let table_id = 1024; + let task = test_create_table_task("foo", table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + let mut procedure = DropTableProcedure::new( + new_drop_table_task("foo", table_id, false), + ddl_context.clone(), + ); + + let error = procedure.on_prepare().await.unwrap_err(); + assert!(error.to_string().contains("retention deadline"), "{error}"); + assert!( + ddl_context + .table_metadata_manager + .table_name_manager() + .get(TableNameKey::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + "foo", + )) + .await + .unwrap() + .is_some() + ); } #[tokio::test] @@ -248,6 +602,7 @@ async fn test_soft_drop_closes_regions_and_keeps_tombstone() { let detector_controller = Arc::new(RecordingRegionFailureDetectorController::default()); let mut ddl_context = new_ddl_context(node_manager); ddl_context.soft_drop_enabled = true; + ddl_context.soft_drop_retention = Some(Duration::from_millis(100)); ddl_context.region_failure_detector_controller = detector_controller.clone(); let table_id = 1024; let table_name = "foo"; @@ -293,16 +648,16 @@ async fn test_soft_drop_closes_regions_and_keeps_tombstone() { let Some(region_request::Body::Close(req)) = request.body else { unreachable!(); }; - requests.push((peer.id, req.region_id)); + requests.push((peer.id, req.region_id, req.flush_on_close)); } requests.sort_unstable(); assert_eq!( requests, vec![ - (1, RegionId::new(table_id, 1).as_u64()), - (1, RegionId::new(table_id, 2).as_u64()), - (2, RegionId::new(table_id, 1).as_u64()), - (2, RegionId::new(table_id, 2).as_u64()), + (1, RegionId::new(table_id, 1).as_u64(), true), + (1, RegionId::new(table_id, 2).as_u64(), false), + (2, RegionId::new(table_id, 1).as_u64(), false), + (2, RegionId::new(table_id, 2).as_u64(), true), ] ); assert!(rx.try_recv().is_err()); @@ -316,12 +671,35 @@ async fn test_soft_drop_closes_regions_and_keeps_tombstone() { .unwrap(); assert!(live_table.is_none()); + let dropped_at = procedure.data.dropped_at.unwrap(); let dropped_table = ddl_context .table_metadata_manager .get_dropped_table(&table_name) .await + .unwrap() .unwrap(); - assert!(dropped_table.is_some()); + assert_eq!(dropped_table.dropped_at, Some(dropped_at)); + assert_eq!(dropped_table.retention_expires_at, Some(dropped_at + 100)); + let dropped_tables = ddl_context + .table_metadata_manager + .list_dropped_tables() + .await + .unwrap(); + assert_eq!(dropped_tables[0].dropped_at, Some(dropped_at)); + assert_eq!( + dropped_tables[0].retention_expires_at, + Some(dropped_at + 100) + ); + + assert!( + ddl_context + .table_metadata_manager + .kv_backend() + .get(retention_expires_at_marker_key(table_id).as_bytes(),) + .await + .unwrap() + .is_some() + ); assert_eq!( detector_controller.deregistered().await, @@ -332,6 +710,69 @@ async fn test_soft_drop_closes_regions_and_keeps_tombstone() { ); } +#[tokio::test] +async fn test_soft_drop_timestamp_is_stable_across_retry_and_recovery() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let mut ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); + ddl_context.soft_drop_enabled = true; + let table_id = 1024; + let task = test_create_table_task("foo", table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + let procedure = DropTableProcedure::new( + new_drop_table_task("foo", table_id, false), + ddl_context.clone(), + ); + assert_eq!(procedure.data.dropped_at, None); + let mut procedure = procedure; + procedure.on_prepare().await.unwrap(); + let dropped_at = procedure.data.dropped_at.unwrap(); + let persisted = procedure.dump().unwrap(); + let recovered = DropTableProcedure::from_json(&persisted, ddl_context.clone()).unwrap(); + assert_eq!(recovered.data.dropped_at, Some(dropped_at)); + let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + data["dropped_at"] = 1_234_567_890_i64.into(); + let mut recovered = + DropTableProcedure::from_json(&data.to_string(), ddl_context.clone()).unwrap(); + + recovered.on_delete_metadata().await.unwrap(); + recovered.on_delete_metadata().await.unwrap(); + let persisted = recovered.dump().unwrap(); + let mut recovered_again = + DropTableProcedure::from_json(&persisted, ddl_context.clone()).unwrap(); + recovered_again.on_delete_metadata().await.unwrap(); + + let marker = kv_backend + .get(dropped_at_marker_key(table_id).as_bytes()) + .await + .unwrap() + .unwrap(); + assert_eq!(marker.value, b"1234567890"); + let deadline_marker = kv_backend + .get(retention_expires_at_marker_key(table_id).as_bytes()) + .await + .unwrap() + .unwrap(); + assert_eq!( + deadline_marker.value, + recovered_again + .data + .retention_expires_at + .unwrap() + .to_string() + .as_bytes() + ); + assert_eq!(recovered_again.data.dropped_at, Some(1_234_567_890)); +} + #[tokio::test] async fn test_hard_drop_keeps_delete_tombstone_flow() { let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); @@ -432,7 +873,7 @@ async fn test_create_table_succeeds_while_tombstone_exists() { } #[tokio::test] -async fn test_drop_recreated_table_fails_when_previous_tombstone_exists() { +async fn test_hard_drop_recreated_table_fails_when_soft_tombstone_exists() { let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); let mut ddl_context = new_ddl_context(node_manager); ddl_context.soft_drop_enabled = true; @@ -453,6 +894,7 @@ async fn test_drop_recreated_table_fails_when_previous_tombstone_exists() { let drop_task = new_drop_table_task(table_name, original_table_id, false); let mut drop_procedure = DropTableProcedure::new(drop_task, ddl_context.clone()); execute_procedure_until_done(&mut drop_procedure).await; + ddl_context.soft_drop_enabled = false; ddl_context .table_metadata_manager @@ -522,7 +964,8 @@ async fn test_undrop_table_restores_metadata_and_reopens_regions() { let datanode_handler = DatanodeWatcher::new(tx); let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); let detector_controller = Arc::new(RecordingRegionFailureDetectorController::default()); - let mut ddl_context = new_ddl_context(node_manager); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let mut ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); ddl_context.soft_drop_enabled = true; ddl_context.region_failure_detector_controller = detector_controller.clone(); let table_id = 1024; @@ -550,6 +993,16 @@ async fn test_undrop_table_restores_metadata_and_reopens_regions() { ddl_context.clone(), ); execute_procedure_until_done(&mut drop_procedure).await; + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .unwrap() + .dropped_at + .is_some() + ); detector_controller.clear().await; while rx.try_recv().is_ok() {} @@ -578,6 +1031,20 @@ async fn test_undrop_table_restores_metadata_and_reopens_regions() { .unwrap() .is_none() ); + assert!( + kv_backend + .get(dropped_at_marker_key(table_id).as_bytes()) + .await + .unwrap() + .is_none() + ); + assert!( + kv_backend + .get(retention_expires_at_marker_key(table_id).as_bytes()) + .await + .unwrap() + .is_none() + ); let mut opened_regions = HashSet::new(); for _ in 0..2 { @@ -734,10 +1201,48 @@ async fn test_soft_drop_metric_logical_table_fails() { let mut procedure = DropTableProcedure::new( new_drop_table_task("foo", logical_table_id, false), - ddl_context, + ddl_context.clone(), ); let err = procedure.on_prepare().await.unwrap_err(); + assert_eq!(err.status_code(), StatusCode::Unsupported); + let persisted_soft_drop = procedure.dump().unwrap(); + ddl_context.soft_drop_enabled = false; + let mut recovered = + DropTableProcedure::from_json(&persisted_soft_drop, ddl_context.clone()).unwrap(); + let err = recovered.on_prepare().await.unwrap_err(); + assert_eq!(err.status_code(), StatusCode::Unsupported); +} +#[tokio::test] +async fn test_soft_drop_file_engine_table_fails() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + let table_id = 1024; + let table_name = "foo"; + let mut task = test_create_table_task(table_name, table_id); + task.table_info.meta.engine = FILE_ENGINE.to_string(); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + + let mut procedure = DropTableProcedure::new( + new_drop_table_task(table_name, table_id, false), + ddl_context.clone(), + ); + let err = procedure.on_prepare().await.unwrap_err(); + assert_eq!(err.status_code(), StatusCode::Unsupported); + let persisted_soft_drop = procedure.dump().unwrap(); + ddl_context.soft_drop_enabled = false; + let mut recovered = + DropTableProcedure::from_json(&persisted_soft_drop, ddl_context.clone()).unwrap(); + let err = recovered.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::Unsupported); } @@ -872,6 +1377,174 @@ async fn test_undrop_table_fails_when_live_name_is_created_after_prepare() { assert_eq!(live_table.table_id(), live_table_id); } +#[tokio::test] +async fn test_undrop_table_closes_opened_regions_when_restore_metadata_races_with_create() { + let (tx, mut rx) = mpsc::channel(8); + let datanode_handler = DatanodeWatcher::new(tx); + let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); + let detector_controller = Arc::new(RecordingRegionFailureDetectorController::default()); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + ddl_context.region_failure_detector_controller = detector_controller.clone(); + let dropped_table_id = 1024; + let live_table_id = 1025; + let table_name = "foo"; + let task = test_create_table_task(table_name, dropped_table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(RegionId::new(dropped_table_id, 1)), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![Peer::empty(2)], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }]), + HashMap::new(), + ) + .await + .unwrap(); + let mut drop_procedure = DropTableProcedure::new( + new_drop_table_task(table_name, dropped_table_id, false), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut drop_procedure).await; + while rx.try_recv().is_ok() {} + detector_controller.clear().await; + + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(dropped_table_id), ddl_context.clone()); + let ctx = new_test_procedure_context(); + procedure.execute(&ctx).await.unwrap(); + procedure.execute(&ctx).await.unwrap(); + + let mut opened_regions = HashSet::new(); + for _ in 0..2 { + let (peer, request) = rx.try_recv().unwrap(); + let Some(region_request::Body::Open(req)) = request.body else { + unreachable!(); + }; + opened_regions.insert((peer.id, req.region_id)); + } + assert_eq!( + opened_regions, + HashSet::from([ + (1, RegionId::new(dropped_table_id, 1).as_u64()), + (2, RegionId::new(dropped_table_id, 1).as_u64()), + ]) + ); + assert_eq!( + detector_controller.registered().await, + vec![(1, RegionId::new(dropped_table_id, 1))] + ); + + ddl_context + .table_metadata_manager + .create_table_metadata( + test_create_table_task(table_name, live_table_id).table_info, + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + + let err = procedure.execute(&ctx).await.unwrap_err(); + assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); + + let mut closed_regions = HashSet::new(); + for _ in 0..2 { + let (peer, request) = rx.try_recv().unwrap(); + let Some(region_request::Body::Close(req)) = request.body else { + unreachable!(); + }; + closed_regions.insert((peer.id, req.region_id)); + } + assert_eq!( + closed_regions, + HashSet::from([ + (1, RegionId::new(dropped_table_id, 1).as_u64()), + (2, RegionId::new(dropped_table_id, 1).as_u64()), + ]) + ); + assert!(rx.try_recv().is_err()); + assert_eq!( + detector_controller.deregistered().await, + vec![(1, RegionId::new(dropped_table_id, 1))] + ); + + let live_table = ddl_context + .table_metadata_manager + .table_name_manager() + .get(TableNameKey::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + table_name, + )) + .await + .unwrap() + .unwrap(); + assert_eq!(live_table.table_id(), live_table_id); +} + +#[tokio::test] +async fn test_undrop_table_lock_key_includes_original_table_name_before_prepare() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + let table_id = 1024; + let table_name = "foo"; + let task = test_create_table_task(table_name, table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + let mut drop_procedure = DropTableProcedure::new( + new_drop_table_task(table_name, table_id, false), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut drop_procedure).await; + + let original_table_name = ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .unwrap() + .table_name; + let procedure = UndropTableProcedure::new_with_original_table_name( + new_undrop_table_task(table_id), + ddl_context, + Some(original_table_name), + ); + + let keys = procedure + .lock_key() + .keys_to_lock() + .cloned() + .collect::>(); + assert!( + keys.iter().any(|key| matches!( + key, + StringKey::Exclusive(key) if key == "__table_name_lock/greptime.public.foo" + )), + "undrop lock keys should include the original table name: {keys:?}" + ); + assert!( + keys.iter().any(|key| matches!( + key, + StringKey::Exclusive(key) if key == "__table_lock/1024" + )), + "undrop lock keys should include the dropped table id: {keys:?}" + ); +} + #[tokio::test] async fn test_undrop_table_replayed_restore_metadata_is_idempotent() { let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); @@ -899,6 +1572,7 @@ async fn test_undrop_table_replayed_restore_metadata_is_idempotent() { UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); + procedure.execute(&ctx).await.unwrap(); let restore_metadata_data = procedure.dump().unwrap(); procedure.execute(&ctx).await.unwrap(); @@ -921,12 +1595,58 @@ async fn test_undrop_table_replayed_restore_metadata_is_idempotent() { } #[tokio::test] -async fn test_purge_dropped_table_drops_regions_and_deletes_tombstone() { +async fn test_legacy_undrop_replays_restore_metadata() { + for partial_restore in [false, true] { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let mut context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); + context.soft_drop_enabled = true; + let table_id = 1024; + let mut procedure = undrop_at_restore_metadata(&context, table_id).await; + let persisted = legacy_undrop_snapshot(&procedure); + + if partial_restore { + let live_key = + TableNameKey::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "foo").to_bytes(); + let tombstone_key = [b"__tombstone/".as_slice(), live_key.as_slice()].concat(); + let value = kv_backend.get(&tombstone_key).await.unwrap().unwrap().value; + kv_backend + .put(PutRequest::new().with_key(live_key).with_value(value)) + .await + .unwrap(); + kv_backend + .batch_delete(BatchDeleteRequest::new().with_keys(vec![tombstone_key])) + .await + .unwrap(); + } else { + procedure + .execute(&new_test_procedure_context()) + .await + .unwrap(); + } + + let mut replayed = UndropTableProcedure::from_json(&persisted, context.clone()).unwrap(); + replayed.recover().unwrap(); + execute_procedure_until_done(&mut replayed).await; + + let (table_info, table_route) = context + .table_metadata_manager + .get_full_table_info(table_id) + .await + .unwrap(); + assert!(table_info.is_some(), "partial_restore={partial_restore}"); + assert!(table_route.is_some(), "partial_restore={partial_restore}"); + } +} + +#[tokio::test] +async fn test_purge_dropped_table_cleans_regions_offline_and_deletes_tombstone() { let (tx, mut rx) = mpsc::channel(8); let datanode_handler = DatanodeWatcher::new(tx); let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); let detector_controller = Arc::new(RecordingRegionFailureDetectorController::default()); - let mut ddl_context = new_ddl_context(node_manager); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let mut ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); ddl_context.soft_drop_enabled = true; ddl_context.region_failure_detector_controller = detector_controller.clone(); let table_id = 1024; @@ -963,15 +1683,16 @@ async fn test_purge_dropped_table_drops_regions_and_deletes_tombstone() { execute_procedure_until_done(&mut procedure).await; let mut requests = Vec::new(); - for _ in 0..4 { + for _ in 0..1 { let (peer, request) = rx.try_recv().unwrap(); requests.push((peer.id, request.body.unwrap())); } requests.sort_unstable_by_key(|(peer_id, _)| *peer_id); - assert_matches!(requests[0].1, region_request::Body::Open(_)); - assert_matches!(requests[1].1, region_request::Body::Drop(_)); - assert_matches!(requests[2].1, region_request::Body::Open(_)); - assert_matches!(requests[3].1, region_request::Body::Close(_)); + let region_request::Body::CleanUp(req) = &requests[0].1 else { + unreachable!(); + }; + assert_eq!(requests[0].0, 1); + assert_eq!(req.region_id, RegionId::new(table_id, 1).as_u64()); assert!(rx.try_recv().is_err()); assert!( ddl_context @@ -981,6 +1702,20 @@ async fn test_purge_dropped_table_drops_regions_and_deletes_tombstone() { .unwrap() .is_none() ); + assert!( + kv_backend + .get(dropped_at_marker_key(table_id).as_bytes()) + .await + .unwrap() + .is_none() + ); + assert!( + kv_backend + .get(retention_expires_at_marker_key(table_id).as_bytes()) + .await + .unwrap() + .is_none() + ); assert_eq!( detector_controller.deregistered().await, @@ -988,6 +1723,461 @@ async fn test_purge_dropped_table_drops_regions_and_deletes_tombstone() { ); } +#[tokio::test] +async fn test_automatic_purge_rechecks_unexpired_tombstone() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + let table_id = 1024; + let table_name = test_create_table_task("foo", table_id).table_name(); + let table_route_value = TableRouteValue::physical(vec![]); + ddl_context + .table_metadata_manager + .create_table_metadata( + test_create_table_task("foo", table_id).table_info, + table_route_value.clone(), + HashMap::new(), + ) + .await + .unwrap(); + ddl_context + .table_metadata_manager + .delete_table_metadata_with_retention( + table_id, + &table_name, + &table_route_value, + &HashMap::new(), + Some(1), + Some(i64::MAX), + ) + .await + .unwrap(); + + let procedure = PurgeDroppedTableProcedure::new( + new_purge_dropped_table_task(table_id), + ddl_context.clone(), + ); + let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + data["check_expired"] = true.into(); + let mut procedure = + PurgeDroppedTableProcedure::from_json(&data.to_string(), ddl_context.clone()).unwrap(); + assert_eq!( + procedure.type_name(), + PurgeDroppedTableProcedure::EXPIRED_TYPE_NAME + ); + + execute_procedure_until_done(&mut procedure).await; + + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_some() + ); +} + +#[tokio::test] +async fn test_recovered_automatic_purge_rechecks_unexpired_tombstone() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + let table_id = 1024; + let table_name = test_create_table_task("foo", table_id).table_name(); + let table_route_value = TableRouteValue::physical(vec![]); + ddl_context + .table_metadata_manager + .create_table_metadata( + test_create_table_task("foo", table_id).table_info, + table_route_value.clone(), + HashMap::new(), + ) + .await + .unwrap(); + ddl_context + .table_metadata_manager + .delete_table_metadata_with_retention( + table_id, + &table_name, + &table_route_value, + &HashMap::new(), + Some(0), + Some(0), + ) + .await + .unwrap(); + + let procedure = PurgeDroppedTableProcedure::new( + new_purge_dropped_table_task(table_id), + ddl_context.clone(), + ); + let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + data["check_expired"] = true.into(); + let mut procedure = + PurgeDroppedTableProcedure::from_json(&data.to_string(), ddl_context.clone()).unwrap(); + procedure + .execute(&new_test_procedure_context()) + .await + .unwrap(); + let persisted = procedure.dump().unwrap(); + + ddl_context + .table_metadata_manager + .restore_table_metadata(table_id, &table_name, &table_route_value, &HashMap::new()) + .await + .unwrap(); + ddl_context + .table_metadata_manager + .delete_table_metadata_with_retention( + table_id, + &table_name, + &table_route_value, + &HashMap::new(), + Some(1), + Some(i64::MAX), + ) + .await + .unwrap(); + + let mut recovered = + PurgeDroppedTableProcedure::from_json(&persisted, ddl_context.clone()).unwrap(); + recovered.recover().unwrap(); + execute_procedure_until_done(&mut recovered).await; + + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_some() + ); +} + +#[tokio::test] +async fn test_recovered_automatic_purge_obeys_runtime_disable() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let ddl_context = new_ddl_context(node_manager); + let table_id = 1024; + let table_name = test_create_table_task("foo", table_id).table_name(); + let table_route_value = TableRouteValue::physical(vec![]); + ddl_context + .table_metadata_manager + .create_table_metadata( + test_create_table_task("foo", table_id).table_info, + table_route_value.clone(), + HashMap::new(), + ) + .await + .unwrap(); + ddl_context + .table_metadata_manager + .delete_table_metadata_with_retention( + table_id, + &table_name, + &table_route_value, + &HashMap::new(), + Some(0), + Some(0), + ) + .await + .unwrap(); + + let procedure = PurgeDroppedTableProcedure::new_if_expired( + new_purge_dropped_table_task(table_id), + ddl_context.clone(), + ); + let mut recovered = + PurgeDroppedTableProcedure::from_json(&procedure.dump().unwrap(), ddl_context.clone()) + .unwrap(); + recovered.recover().unwrap(); + execute_procedure_until_done(&mut recovered).await; + + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_some() + ); +} + +#[tokio::test] +async fn test_recovered_automatic_purge_finishes_after_cleanup_is_claimed() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let mut enabled_context = + new_ddl_context_with_kv_backend(node_manager.clone(), kv_backend.clone()); + enabled_context.soft_drop_enabled = true; + let table_id = 1024; + create_dropped_table(&enabled_context, table_id, Some(0), Some(0)).await; + + let mut procedure = PurgeDroppedTableProcedure::new_if_expired( + new_purge_dropped_table_task(table_id), + enabled_context.clone(), + ); + procedure + .execute(&new_test_procedure_context()) + .await + .unwrap(); + let persisted_before_cleanup = procedure.dump().unwrap(); + procedure + .execute(&new_test_procedure_context()) + .await + .unwrap(); + assert!( + kv_backend + .get(purging_marker_key(table_id).as_bytes()) + .await + .unwrap() + .is_some() + ); + + let disabled_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); + let mut recovered = + PurgeDroppedTableProcedure::from_json(&persisted_before_cleanup, disabled_context.clone()) + .unwrap(); + recovered.recover().unwrap(); + execute_procedure_until_done(&mut recovered).await; + + assert!( + disabled_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_none() + ); + assert!( + kv_backend + .get(purging_marker_key(table_id).as_bytes()) + .await + .unwrap() + .is_none() + ); +} + +#[tokio::test] +async fn test_recovered_automatic_purge_ignores_foreign_generation_claim() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let mut ddl_context = new_ddl_context_with_kv_backend(node_manager.clone(), kv_backend.clone()); + ddl_context.soft_drop_enabled = true; + let table_id = 1024; + create_dropped_table(&ddl_context, table_id, Some(0), Some(0)).await; + kv_backend + .put( + PutRequest::new() + .with_key(drop_generation_marker_key(table_id)) + .with_value("generation-a"), + ) + .await + .unwrap(); + + let mut procedure = PurgeDroppedTableProcedure::new_if_expired( + new_purge_dropped_table_task(table_id), + ddl_context.clone(), + ); + procedure + .execute(&new_test_procedure_context()) + .await + .unwrap(); + let persisted = procedure.dump().unwrap(); + kv_backend + .put( + PutRequest::new() + .with_key(purging_marker_key(table_id)) + .with_value("generation-b"), + ) + .await + .unwrap(); + + let mut recovered = + PurgeDroppedTableProcedure::from_json(&persisted, ddl_context.clone()).unwrap(); + execute_procedure_until_done(&mut recovered).await; + + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_some() + ); +} + +#[tokio::test] +async fn test_undrop_rejects_table_claimed_for_purge() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); + let table_id = 1024; + create_dropped_table(&ddl_context, table_id, Some(0), Some(0)).await; + kv_backend + .put(PutRequest::new().with_key(purging_marker_key(table_id))) + .await + .unwrap(); + + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + let error = procedure.on_prepare().await.unwrap_err(); + + assert_eq!(StatusCode::TableNotFound, error.status_code()); + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_some() + ); +} + +#[tokio::test] +async fn test_recovered_undrop_rejects_tombstone_deleted_by_purge() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let ddl_context = new_ddl_context(node_manager); + let table_id = 1024; + create_dropped_table(&ddl_context, table_id, Some(1), Some(i64::MAX)).await; + + let mut undrop = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + undrop.on_prepare().await.unwrap(); + let open_regions_state = undrop.dump().unwrap(); + undrop.execute(&new_test_procedure_context()).await.unwrap(); + let restore_metadata_state = undrop.dump().unwrap(); + + let mut purge = PurgeDroppedTableProcedure::new( + new_purge_dropped_table_task(table_id), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut purge).await; + + for persisted in [open_regions_state, restore_metadata_state] { + let mut recovered = + UndropTableProcedure::from_json(&persisted, ddl_context.clone()).unwrap(); + let error = recovered + .execute(&new_test_procedure_context()) + .await + .unwrap_err(); + assert_eq!(StatusCode::TableNotFound, error.status_code()); + } +} + +#[tokio::test] +async fn test_recovered_undrop_rejects_replacement_generation() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); + let table_id = 1024; + create_dropped_table(&ddl_context, table_id, Some(1), Some(i64::MAX)).await; + kv_backend + .put( + PutRequest::new() + .with_key(drop_generation_marker_key(table_id)) + .with_value("generation-a"), + ) + .await + .unwrap(); + + let mut undrop = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + undrop.on_prepare().await.unwrap(); + let persisted = undrop.dump().unwrap(); + kv_backend + .put( + PutRequest::new() + .with_key(drop_generation_marker_key(table_id)) + .with_value("generation-b"), + ) + .await + .unwrap(); + + let mut recovered = UndropTableProcedure::from_json(&persisted, ddl_context).unwrap(); + let error = recovered + .execute(&new_test_procedure_context()) + .await + .unwrap_err(); + assert_eq!(StatusCode::TableNotFound, error.status_code()); +} + +#[tokio::test] +async fn test_legacy_unmarked_tombstone_can_be_undropped_and_purged() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let ddl_context = new_ddl_context(node_manager); + let table_id = 1024; + let table_name = test_create_table_task("foo", table_id).table_name(); + let table_route_value = TableRouteValue::physical(vec![]); + ddl_context + .table_metadata_manager + .create_table_metadata( + test_create_table_task("foo", table_id).table_info, + table_route_value.clone(), + HashMap::new(), + ) + .await + .unwrap(); + ddl_context + .table_metadata_manager + .delete_table_metadata( + table_id, + &table_name, + &table_route_value, + &HashMap::new(), + None, + ) + .await + .unwrap(); + assert_eq!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .unwrap() + .dropped_at, + None + ); + let dropped_tables = ddl_context + .table_metadata_manager + .list_dropped_tables() + .await + .unwrap(); + assert_eq!(dropped_tables.len(), 1); + assert_eq!(dropped_tables[0].dropped_at, None); + + let mut undrop = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); + execute_procedure_until_done(&mut undrop).await; + + ddl_context + .table_metadata_manager + .delete_table_metadata( + table_id, + &table_name, + &table_route_value, + &HashMap::new(), + None, + ) + .await + .unwrap(); + let mut purge = PurgeDroppedTableProcedure::new( + new_purge_dropped_table_task(table_id), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut purge).await; + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await + .unwrap() + .is_none() + ); +} + #[tokio::test] async fn test_purge_dropped_table_by_id_selects_tombstone_when_live_table_exists() { let (tx, mut rx) = mpsc::channel(8); @@ -1050,124 +2240,19 @@ async fn test_purge_dropped_table_by_id_selects_tombstone_when_live_table_exists assert_eq!(live_table.table_id(), live_table_id); let (_, request) = rx.try_recv().unwrap(); - let Some(region_request::Body::Open(req)) = request.body else { - unreachable!(); - }; - assert_eq!(req.region_id, RegionId::new(dropped_table_id, 1).as_u64()); - - let (_, request) = rx.try_recv().unwrap(); - let Some(region_request::Body::Drop(req)) = request.body else { + let Some(region_request::Body::CleanUp(req)) = request.body else { unreachable!(); }; assert_eq!(req.region_id, RegionId::new(dropped_table_id, 1).as_u64()); assert!(rx.try_recv().is_err()); } -#[tokio::test] -async fn test_purge_dropped_table_replayed_open_regions_ignores_dropped_regions() { - let (tx, mut rx) = mpsc::channel(8); - let dropped_regions = Arc::new(StdMutex::new(HashSet::new())); - let datanode_handler = DatanodeWatcher::new(tx).with_handler({ - let dropped_regions = dropped_regions.clone(); - move |_peer, request| { - let Some(body) = request.body.as_ref() else { - return Ok(RegionResponse::new(0)); - }; - match body { - region_request::Body::Open(req) - if dropped_regions.lock().unwrap().contains(&req.region_id) => - { - Err::(BoxedError::new(MockRegionNotFoundError)) - .context(error::ExternalSnafu) - } - region_request::Body::Drop(req) => { - dropped_regions.lock().unwrap().insert(req.region_id); - Ok(RegionResponse::new(0)) - } - _ => Ok(RegionResponse::new(0)), - } - } - }); - let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); - let mut ddl_context = new_ddl_context(node_manager); - ddl_context.soft_drop_enabled = true; - let table_id = 1024; - let table_name = "foo"; - let task = test_create_table_task(table_name, table_id); - ddl_context - .table_metadata_manager - .create_table_metadata( - task.table_info.clone(), - TableRouteValue::physical(vec![RegionRoute { - region: Region::new_test(RegionId::new(table_id, 1)), - leader_peer: Some(Peer::empty(1)), - follower_peers: vec![], - leader_state: None, - leader_down_since: None, - write_route_policy: None, - }]), - HashMap::new(), - ) - .await - .unwrap(); - let mut drop_procedure = DropTableProcedure::new( - new_drop_table_task(table_name, table_id, false), - ddl_context.clone(), - ); - execute_procedure_until_done(&mut drop_procedure).await; - while rx.try_recv().is_ok() {} - - let mut procedure = PurgeDroppedTableProcedure::new( - new_purge_dropped_table_task(table_id), - ddl_context.clone(), - ); - let ctx = new_test_procedure_context(); - procedure.execute(&ctx).await.unwrap(); - let open_regions_data = procedure.dump().unwrap(); - procedure.execute(&ctx).await.unwrap(); - procedure.execute(&ctx).await.unwrap(); - - let mut replayed = - PurgeDroppedTableProcedure::from_json(&open_regions_data, ddl_context.clone()).unwrap(); - execute_procedure_until_done(&mut replayed).await; - - assert!( - ddl_context - .table_metadata_manager - .get_dropped_table(&drop_procedure.data.task.table_name()) - .await - .unwrap() - .is_none() - ); -} - -#[derive(Debug, snafu::Snafu)] -#[snafu(display("mock region not found"))] -struct MockRegionNotFoundError; - -impl StackError for MockRegionNotFoundError { - fn debug_fmt(&self, _: usize, _: &mut Vec) {} - - fn next(&self) -> Option<&dyn StackError> { - None - } -} - -impl ErrorExt for MockRegionNotFoundError { - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn status_code(&self) -> StatusCode { - StatusCode::RegionNotFound - } -} - #[tokio::test] async fn test_on_rollback() { let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); let kv_backend = Arc::new(MemoryKvBackend::new()); - let ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); + let mut ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend.clone()); + ddl_context.soft_drop_enabled = true; // Prepares physical table metadata. let mut create_physical_table_task = test_create_physical_table_task("phy_table"); let TableMetadata { @@ -1208,8 +2293,36 @@ async fn test_on_rollback() { procedure.on_prepare().await.unwrap(); assert!(procedure.rollback_supported()); procedure.on_delete_metadata().await.unwrap(); + assert!( + kv_backend + .get(dropped_at_marker_key(physical_table_id).as_bytes()) + .await + .unwrap() + .is_some() + ); + assert!( + kv_backend + .get(retention_expires_at_marker_key(physical_table_id).as_bytes()) + .await + .unwrap() + .is_some() + ); assert!(procedure.rollback_supported()); procedure.rollback(&ctx).await.unwrap(); + assert!( + kv_backend + .get(dropped_at_marker_key(physical_table_id).as_bytes()) + .await + .unwrap() + .is_none() + ); + assert!( + kv_backend + .get(retention_expires_at_marker_key(physical_table_id).as_bytes()) + .await + .unwrap() + .is_none() + ); // Rollback again assert!(procedure.rollback_supported()); procedure.rollback(&ctx).await.unwrap(); @@ -1218,6 +2331,7 @@ async fn test_on_rollback() { } // Drops the logical table + ddl_context.soft_drop_enabled = false; let task = new_drop_table_task("foo", table_ids[0], false); let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); procedure.on_prepare().await.unwrap(); @@ -1258,6 +2372,7 @@ async fn create_metric_logical_table_tombstone( &task.table_name(), &TableRouteValue::logical(physical_table_id), &HashMap::new(), + None, ) .await .unwrap(); diff --git a/src/common/meta/src/ddl/undrop_table.rs b/src/common/meta/src/ddl/undrop_table.rs index f08d606a827..872a11ce65f 100644 --- a/src/common/meta/src/ddl/undrop_table.rs +++ b/src/common/meta/src/ddl/undrop_table.rs @@ -18,13 +18,12 @@ use api::v1::region::{ OpenRequest as PbOpenRegionRequest, RegionRequest, RegionRequestHeader, region_request, }; use async_trait::async_trait; -use common_error::ext::ErrorExt; -use common_error::status_code::StatusCode; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ Context as ProcedureContext, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::tracing_context::TracingContext; +use common_telemetry::warn; use common_wal::options::WalOptions; use futures::future::join_all; use serde::{Deserialize, Serialize}; @@ -34,6 +33,7 @@ use strum::AsRefStr; use table::metadata::TableId; use table::table_name::TableName; +use crate::ddl::drop_table::executor::DropTableExecutor; use crate::ddl::utils::{ add_peer_context_if_needed, convert_region_routes_to_detecting_regions, is_metric_engine_logical_table, map_to_procedure_error, region_storage_path, @@ -43,7 +43,7 @@ use crate::error::{self, Result}; use crate::instruction::CacheIdent; use crate::key::table_name::TableNameKey; use crate::key::table_route::TableRouteValue; -use crate::lock_key::TableLock; +use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; use crate::rpc::ddl::UndropTableTask; use crate::rpc::router::{ RegionRoute, find_follower_regions, find_followers, find_leader_regions, find_leaders, @@ -58,10 +58,17 @@ impl UndropTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::UndropTable"; pub fn new(task: UndropTableTask, context: DdlContext) -> Self { - Self { - context, - data: UndropTableData::new(task), - } + Self::new_with_original_table_name(task, context, None) + } + + pub(crate) fn new_with_original_table_name( + task: UndropTableTask, + context: DdlContext, + table_name: Option, + ) -> Self { + let mut data = UndropTableData::new(task); + data.table_name = table_name; + Self { context, data } } pub fn from_json(json: &str, context: DdlContext) -> ProcedureResult { @@ -70,6 +77,7 @@ impl UndropTableProcedure { } pub(crate) async fn on_prepare(&mut self) -> Result { + self.ensure_not_purging().await?; let dropped_table = self .context .table_metadata_manager @@ -93,6 +101,10 @@ impl UndropTableProcedure { self.data.table_name = Some(dropped_table.table_name.clone()); self.data.table_route_value = Some(dropped_table.table_route_value.clone()); self.data.region_wal_options = dropped_table.region_wal_options; + self.data.dropped_at = dropped_table.dropped_at; + self.data.retention_expires_at = dropped_table.retention_expires_at; + self.data.drop_generation = dropped_table.drop_generation; + self.data.tombstone_identity_loaded = true; ensure!( !is_metric_engine_logical_table( &dropped_table.table_info_value.table_info, @@ -108,30 +120,203 @@ impl UndropTableProcedure { } async fn on_restore_metadata(&mut self) -> Result { + if !self.data.tombstone_identity_loaded + && self + .context + .table_metadata_manager + .get_dropped_table_by_id(self.data.task.table_id) + .await? + .is_some() + && !self.live_table_name_matches().await? + { + self.data.state = UndropTableState::Prepare; + return Ok(Status::executing(true)); + } + self.ensure_table_name_loaded().await?; + let table_name = self.data.table_name().clone(); let table_route_value = self.data.table_route_value(); - self.context - .table_metadata_manager - .restore_table_metadata( - self.data.task.table_id, - self.data.table_name(), - table_route_value, - &self.data.region_wal_options, - ) - .await - .map_err(|err| match err { - error::Error::TombstoneTargetAlreadyExists { .. } => { - error::TableAlreadyExistsSnafu { - table_name: self.data.table_name().to_string(), - } - .build() - } - err => err, - })?; + let restore_result = async { + self.ensure_tombstone_available_for_restore().await?; + self.context + .table_metadata_manager + .restore_table_metadata( + self.data.task.table_id, + &table_name, + table_route_value, + &self.data.region_wal_options, + ) + .await?; + self.ensure_live_table_metadata_restored().await + } + .await; + if let Err(err) = restore_result { + let should_cleanup_opened_regions = !err.is_retry_later(); + let err = self.map_restore_metadata_error(err); + if should_cleanup_opened_regions + && let Err(cleanup_err) = self.cleanup_opened_regions_after_restore_failure().await + { + warn!( + cleanup_err; + "Failed to close opened regions after undrop metadata restore failure, table_id: {}", + self.data.task.table_id + ); + } + return Err(err); + } self.data.state = UndropTableState::InvalidateTableCache; Ok(Status::executing(true)) } + async fn ensure_table_name_loaded(&mut self) -> Result<()> { + if self.data.table_name.is_some() { + return Ok(()); + } + + let dropped_table = self + .context + .table_metadata_manager + .get_dropped_table_by_id(self.data.task.table_id) + .await? + .with_context(|| error::TableNotFoundSnafu { + table_name: self.data.task.table_id.to_string(), + })?; + self.data.table_name = Some(dropped_table.table_name); + Ok(()) + } + + async fn ensure_not_purging(&self) -> Result<()> { + ensure!( + !self + .context + .table_metadata_manager + .is_dropped_table_purging(self.data.task.table_id) + .await?, + error::TableNotFoundSnafu { + table_name: self.data.task.table_id.to_string(), + } + ); + Ok(()) + } + + async fn ensure_tombstone_available(&self) -> Result<()> { + self.ensure_not_purging().await?; + let dropped_table = self + .context + .table_metadata_manager + .get_dropped_table_by_id(self.data.task.table_id) + .await? + .with_context(|| error::TableNotFoundSnafu { + table_name: self.data.task.table_id.to_string(), + })?; + ensure!( + dropped_table.table_name == *self.data.table_name() + && dropped_table.dropped_at == self.data.dropped_at + && dropped_table.retention_expires_at == self.data.retention_expires_at + && dropped_table.drop_generation == self.data.drop_generation, + error::TableNotFoundSnafu { + table_name: self.data.task.table_id.to_string(), + } + ); + Ok(()) + } + + async fn ensure_tombstone_available_for_restore(&self) -> Result<()> { + if self + .context + .table_metadata_manager + .get_dropped_table_by_id(self.data.task.table_id) + .await? + .is_some() + { + if !self.data.tombstone_identity_loaded && self.live_table_name_matches().await? { + return Ok(()); + } + return self.ensure_tombstone_available().await; + } + + ensure!( + self.live_table_name_matches().await?, + error::TableNotFoundSnafu { + table_name: self.data.task.table_id.to_string(), + } + ); + Ok(()) + } + + async fn live_table_name_matches(&self) -> Result { + let live_table = self + .context + .table_metadata_manager + .table_name_manager() + .get(TableNameKey::from(self.data.table_name())) + .await?; + Ok(live_table.is_some_and(|table| table.table_id() == self.data.task.table_id)) + } + + async fn ensure_live_table_metadata_restored(&self) -> Result<()> { + let (table_info, table_route) = self + .context + .table_metadata_manager + .get_full_table_info(self.data.task.table_id) + .await?; + ensure!( + table_info.is_some() && table_route.is_some(), + error::TableNotFoundSnafu { + table_name: self.data.task.table_id.to_string(), + } + ); + Ok(()) + } + + fn map_restore_metadata_error(&self, err: error::Error) -> error::Error { + match err { + error::Error::TombstoneTargetAlreadyExists { .. } => { + let table_name = self + .data + .table_name + .as_ref() + .map(ToString::to_string) + .unwrap_or_else(|| self.data.task.table_id.to_string()); + error::TableAlreadyExistsSnafu { table_name }.build() + } + err => err, + } + } + + async fn cleanup_opened_regions_after_restore_failure(&self) -> Result<()> { + let Some(table_route_value) = self.data.table_route_value.as_ref() else { + return Ok(()); + }; + let TableRouteValue::Physical(route) = table_route_value else { + return Ok(()); + }; + let region_routes = route.region_routes.clone(); + let close_result = if let Some(table_name) = self.data.table_name.as_ref() { + let executor = + DropTableExecutor::new(table_name.clone(), self.data.task.table_id, false); + + executor + .on_close_regions( + &self.context.node_manager, + &self.context.leader_region_registry, + ®ion_routes, + false, + ) + .await + } else { + Ok(()) + }; + self.context + .deregister_failure_detectors(convert_region_routes_to_detecting_regions( + ®ion_routes, + )) + .await; + close_result?; + Ok(()) + } + async fn on_open_regions(&mut self) -> Result { + self.ensure_tombstone_available().await?; self.ensure_live_table_not_exists().await?; let TableRouteValue::Physical(route) = self.data.table_route_value() else { self.data.state = UndropTableState::RestoreMetadata; @@ -195,6 +380,10 @@ impl Procedure for UndropTableProcedure { } fn recover(&mut self) -> ProcedureResult<()> { + if self.data.state == UndropTableState::OpenRegions && !self.data.tombstone_identity_loaded + { + self.data.state = UndropTableState::Prepare; + } Ok(()) } @@ -213,7 +402,22 @@ impl Procedure for UndropTableProcedure { } fn lock_key(&self) -> LockKey { - LockKey::new(vec![TableLock::Write(self.data.task.table_id).into()]) + let mut lock_key = Vec::new(); + if let Some(table_name) = &self.data.table_name { + lock_key.push(CatalogLock::Read(&table_name.catalog_name).into()); + lock_key + .push(SchemaLock::read(&table_name.catalog_name, &table_name.schema_name).into()); + lock_key.push( + TableNameLock::new( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + ) + .into(), + ); + } + lock_key.push(TableLock::Write(self.data.task.table_id).into()); + LockKey::new(lock_key) } } @@ -232,27 +436,6 @@ pub(crate) async fn open_regions( table_info, region_routes, region_wal_options, - false, - ) - .await -} - -pub(crate) async fn open_regions_ignore_region_not_found( - context: &DdlContext, - table_id: TableId, - table_name: &TableName, - table_info: &table::metadata::TableInfo, - region_routes: &[RegionRoute], - region_wal_options: &HashMap, -) -> Result<()> { - open_regions_inner( - context, - table_id, - table_name, - table_info, - region_routes, - region_wal_options, - true, ) .await } @@ -264,7 +447,6 @@ async fn open_regions_inner( table_info: &table::metadata::TableInfo, region_routes: &[RegionRoute], region_wal_options: &HashMap, - ignore_region_not_found: bool, ) -> Result<()> { let template = build_template_from_raw_table_info(table_info)?; let builder = CreateRequestBuilder::new(template, None); @@ -303,9 +485,7 @@ async fn open_regions_inner( let datanode = datanode.clone(); let requester = requester.clone(); tasks.push(async move { - if let Err(err) = requester.handle(request).await - && !(ignore_region_not_found && err.status_code() == StatusCode::RegionNotFound) - { + if let Err(err) = requester.handle(request).await { return Err(add_peer_context_if_needed(datanode)(err)); } Ok(()) @@ -332,6 +512,14 @@ pub struct UndropTableData { table_route_value: Option, #[serde(default)] region_wal_options: HashMap, + #[serde(default)] + dropped_at: Option, + #[serde(default)] + retention_expires_at: Option, + #[serde(default)] + drop_generation: Option, + #[serde(default)] + tombstone_identity_loaded: bool, } impl UndropTableData { @@ -343,6 +531,10 @@ impl UndropTableData { table_info: None, table_route_value: None, region_wal_options: HashMap::new(), + dropped_at: None, + retention_expires_at: None, + drop_generation: None, + tombstone_identity_loaded: false, } } @@ -366,3 +558,151 @@ enum UndropTableState { OpenRegions, InvalidateTableCache, } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use api::region::RegionResponse; + use api::v1::region::region_request; + use async_trait::async_trait; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_error::ext::ErrorExt; + use common_error::status_code::StatusCode; + use store_api::storage::RegionId; + use table::table_name::TableName; + use tokio::sync::{Mutex, mpsc}; + + use super::*; + use crate::ddl::test_util::datanode_handler::DatanodeWatcher; + use crate::ddl::{DetectingRegion, RegionFailureDetectorController}; + use crate::peer::Peer; + use crate::rpc::router::{Region, RegionRoute}; + use crate::test_util::{MockDatanodeManager, new_ddl_context}; + + #[derive(Default)] + struct RecordingRegionFailureDetectorController { + deregistered: Mutex>, + } + + #[async_trait] + impl RegionFailureDetectorController for RecordingRegionFailureDetectorController { + async fn register_failure_detectors(&self, _detecting_regions: Vec) {} + + async fn reset_failure_detectors(&self, _detecting_regions: Vec) {} + + async fn deregister_failure_detectors(&self, detecting_regions: Vec) { + self.deregistered.lock().await.extend(detecting_regions); + } + } + + #[test] + fn test_map_restore_metadata_error_without_table_name() { + let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + let procedure = UndropTableProcedure::new(UndropTableTask { table_id: 42 }, context); + + let err = procedure.map_restore_metadata_error( + error::TombstoneTargetAlreadyExistsSnafu { + key: "table-name-key".to_string(), + } + .build(), + ); + + assert_eq!(StatusCode::TableAlreadyExists, err.status_code()); + } + + #[test] + fn test_recovered_legacy_snapshot_reloads_tombstone_identity() { + let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + let mut procedure = + UndropTableProcedure::new(UndropTableTask { table_id: 42 }, context.clone()); + procedure.data.state = UndropTableState::OpenRegions; + let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); + let data = data.as_object_mut().unwrap(); + data.remove("dropped_at"); + data.remove("retention_expires_at"); + data.remove("drop_generation"); + data.remove("tombstone_identity_loaded"); + + let mut recovered = + UndropTableProcedure::from_json(&serde_json::to_string(data).unwrap(), context) + .unwrap(); + recovered.recover().unwrap(); + + assert_eq!(UndropTableState::Prepare, recovered.data.state); + } + + #[tokio::test] + async fn test_cleanup_opened_regions_without_table_name_deregisters_detectors() { + let detector_controller = Arc::new(RecordingRegionFailureDetectorController::default()); + let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + context.region_failure_detector_controller = detector_controller.clone(); + + let table_id = 1024; + let region_id = RegionId::new(table_id, 1); + let mut procedure = UndropTableProcedure::new(UndropTableTask { table_id }, context); + procedure.data.table_route_value = Some(TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(region_id), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }])); + + procedure + .cleanup_opened_regions_after_restore_failure() + .await + .unwrap(); + + assert_eq!( + detector_controller.deregistered.lock().await.as_slice(), + &[(1, region_id)] + ); + } + + #[tokio::test] + async fn test_cleanup_opened_regions_deregisters_detectors_when_close_fails() { + let (tx, _rx) = mpsc::channel(8); + let datanode_handler = DatanodeWatcher::new(tx).with_handler(|_, request| { + if matches!(request.body, Some(region_request::Body::Close(_))) { + return error::UnexpectedSnafu { + err_msg: "mock close error".to_string(), + } + .fail(); + } + Ok(RegionResponse::new(0)) + }); + let detector_controller = Arc::new(RecordingRegionFailureDetectorController::default()); + let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(datanode_handler))); + context.region_failure_detector_controller = detector_controller.clone(); + + let table_id = 1024; + let region_id = RegionId::new(table_id, 1); + let mut procedure = UndropTableProcedure::new(UndropTableTask { table_id }, context); + procedure.data.table_name = Some(TableName::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + "foo", + )); + procedure.data.table_route_value = Some(TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(region_id), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }])); + + let err = procedure + .cleanup_opened_regions_after_restore_failure() + .await + .unwrap_err(); + + assert_eq!(StatusCode::Unexpected, err.status_code()); + assert_eq!( + detector_controller.deregistered.lock().await.as_slice(), + &[(1, region_id)] + ); + } +} diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index bfd858a1963..455fc098b86 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -272,6 +272,18 @@ impl DdlManager { .context(RegisterProcedureLoaderSnafu { type_name })?; } + let type_name = PurgeDroppedTableProcedure::EXPIRED_TYPE_NAME; + let context = self.create_context(); + self.procedure_manager + .register_loader( + type_name, + Box::new(move |json: &str| { + PurgeDroppedTableProcedure::from_json(json, context.clone()) + .map(|procedure| Box::new(procedure) as _) + }), + ) + .context(RegisterProcedureLoaderSnafu { type_name })?; + self.repartition_procedure_factory .register_loaders(&self.ddl_context, &self.procedure_manager) .context(RegisterRepartitionProcedureLoaderSnafu)?; @@ -480,7 +492,19 @@ impl DdlManager { undrop_table_task: UndropTableTask, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = UndropTableProcedure::new(undrop_table_task, context); + let original_table_name = context + .table_metadata_manager + .get_dropped_table_by_id(undrop_table_task.table_id) + .await? + .with_context(|| TableNotFoundSnafu { + table_name: undrop_table_task.table_id.to_string(), + })? + .table_name; + let procedure = UndropTableProcedure::new_with_original_table_name( + undrop_table_task, + context, + Some(original_table_name), + ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); self.execute_procedure_and_wait(procedure_with_id).await @@ -499,6 +523,20 @@ impl DdlManager { self.execute_procedure_and_wait(procedure_with_id).await } + /// Submits and executes a purge task that first rechecks the tombstone deadline. + #[tracing::instrument(skip_all)] + pub async fn submit_expired_purge_dropped_table_task( + &self, + purge_dropped_table_task: PurgeDroppedTableTask, + ) -> Result<(ProcedureId, Option)> { + let context = self.create_context(); + let procedure = + PurgeDroppedTableProcedure::new_if_expired(purge_dropped_table_task, context); + let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + + self.execute_procedure_and_wait(procedure_with_id).await + } + /// Submits and executes a create database task. #[tracing::instrument(skip_all)] pub async fn submit_create_database( @@ -1204,7 +1242,8 @@ mod tests { use std::sync::Arc; use std::time::Duration; - use common_error::ext::BoxedError; + use common_error::ext::{BoxedError, ErrorExt}; + use common_error::status_code::StatusCode; use common_procedure::local::LocalManager; use common_procedure::test_util::InMemoryPoisonStore; use common_procedure::{BoxedProcedure, ProcedureManagerRef}; @@ -1228,6 +1267,7 @@ mod tests { use crate::peer::Peer; use crate::region_keeper::MemoryRegionKeeper; use crate::region_registry::LeaderRegionRegistry; + use crate::rpc::ddl::UndropTableTask; use crate::sequence::SequenceBuilder; use crate::state_store::KvStateStore; use crate::wal_provider::WalProvider; @@ -1309,6 +1349,7 @@ mod tests { leader_region_registry: Arc::new(LeaderRegionRegistry::default()), region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl), soft_drop_enabled: false, + soft_drop_retention: None, }, procedure_manager.clone(), Arc::new(DummyRepartitionProcedureFactory), @@ -1326,4 +1367,56 @@ mod tests { assert!(procedure_manager.contains_loader(loader)); } } + + #[tokio::test] + async fn test_submit_undrop_missing_tombstone_returns_table_not_found_directly() { + let kv_backend = Arc::new(MemoryKvBackend::new()); + let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone())); + let table_metadata_allocator = Arc::new(TableMetadataAllocator::new( + Arc::new(SequenceBuilder::new("test", kv_backend.clone()).build()), + Arc::new(WalProvider::default()), + )); + let flow_metadata_manager = Arc::new(FlowMetadataManager::new(kv_backend.clone())); + let flow_metadata_allocator = Arc::new(FlowMetadataAllocator::with_noop_peer_allocator( + Arc::new(SequenceBuilder::new("flow-test", kv_backend.clone()).build()), + )); + + let state_store = Arc::new(KvStateStore::new(kv_backend.clone())); + let poison_manager = Arc::new(InMemoryPoisonStore::default()); + let procedure_manager = Arc::new(LocalManager::new( + Default::default(), + state_store, + poison_manager, + None, + None, + )); + + let ddl_manager = DdlManager::try_new( + DdlContext { + node_manager: Arc::new(DummyDatanodeManager), + cache_invalidator: Arc::new(DummyCacheInvalidator), + table_metadata_manager, + table_metadata_allocator, + flow_metadata_manager, + flow_metadata_allocator, + memory_region_keeper: Arc::new(MemoryRegionKeeper::default()), + leader_region_registry: Arc::new(LeaderRegionRegistry::default()), + region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl), + soft_drop_enabled: true, + soft_drop_retention: Some(std::time::Duration::from_secs(1)), + }, + procedure_manager, + Arc::new(DummyRepartitionProcedureFactory), + true, + ) + .unwrap(); + + let err = ddl_manager + .submit_undrop_table_task(UndropTableTask { table_id: 1024 }) + .await + .unwrap_err(); + + assert_eq!(err.status_code(), StatusCode::TableNotFound); + assert!(matches!(err, crate::error::Error::TableNotFound { .. })); + } } diff --git a/src/common/meta/src/key.rs b/src/common/meta/src/key.rs index 56fa35b9431..71e9ce4f4b0 100644 --- a/src/common/meta/src/key.rs +++ b/src/common/meta/src/key.rs @@ -164,8 +164,9 @@ use crate::key::topic_region::TopicRegionValue; use crate::key::txn_helper::TxnOpGetResponseSet; use crate::kv_backend::KvBackendRef; use crate::kv_backend::txn::{Txn, TxnOp}; +use crate::rpc::KeyValue; use crate::rpc::router::{LeaderState, RegionRoute, region_distribution}; -use crate::rpc::store::BatchDeleteRequest; +use crate::rpc::store::{BatchDeleteRequest, PutRequest}; use crate::state_store::PoisonValue; use crate::wal_provider::RegionWalOptions; @@ -440,6 +441,12 @@ pub struct DroppedTableName { pub table_id: TableId, /// Original fully qualified table name. pub table_name: TableName, + /// Unix timestamp in milliseconds when this table was soft-dropped. + pub dropped_at: Option, + /// Fixed automatic-GC deadline in Unix milliseconds. + pub retention_expires_at: Option, + /// Unique identity of this soft-drop generation. + pub drop_generation: Option, } #[derive(Debug, Clone)] @@ -454,6 +461,40 @@ pub struct DroppedTableMetadata { pub table_route_value: TableRouteValue, /// Per-region WAL options recovered from tombstoned datanode metadata. pub region_wal_options: HashMap, + /// Unix timestamp in milliseconds when this table was soft-dropped. + pub dropped_at: Option, + /// Fixed automatic-GC deadline in Unix milliseconds. + pub retention_expires_at: Option, + /// Unique identity of this soft-drop generation. + pub drop_generation: Option, +} + +/// Lifecycle markers stored with a logically deleted table. +pub struct DroppedTableLifecycle<'a> { + pub dropped_at: Option, + pub retention_expires_at: Option, + pub drop_generation: Option<&'a str>, +} + +const DROPPED_AT_KEY_PREFIX: &str = "__dropped_at"; +const RETENTION_EXPIRES_AT_KEY_PREFIX: &str = "__retention_expires_at"; +const DROP_GENERATION_KEY_PREFIX: &str = "__drop_generation"; +const PURGING_KEY_PREFIX: &str = "__purging"; + +fn dropped_at_key(table_id: TableId) -> Vec { + format!("{DROPPED_AT_KEY_PREFIX}/{table_id}").into_bytes() +} + +fn retention_expires_at_key(table_id: TableId) -> Vec { + format!("{RETENTION_EXPIRES_AT_KEY_PREFIX}/{table_id}").into_bytes() +} + +fn drop_generation_key(table_id: TableId) -> Vec { + format!("{DROP_GENERATION_KEY_PREFIX}/{table_id}").into_bytes() +} + +fn purging_key(table_id: TableId) -> Vec { + format!("{PURGING_KEY_PREFIX}/{table_id}").into_bytes() } impl Deref for DeserializedValueWithBytes { @@ -998,10 +1039,78 @@ impl TableMetadataManager { table_name: &TableName, table_route_value: &TableRouteValue, region_wal_options: &HashMap, + dropped_at: Option, + ) -> Result<()> { + self.delete_table_metadata_with_retention( + table_id, + table_name, + table_route_value, + region_wal_options, + dropped_at, + None, + ) + .await + } + + /// Deletes metadata logically and stores fixed soft-drop lifecycle markers. + pub async fn delete_table_metadata_with_retention( + &self, + table_id: TableId, + table_name: &TableName, + table_route_value: &TableRouteValue, + region_wal_options: &HashMap, + dropped_at: Option, + retention_expires_at: Option, + ) -> Result<()> { + self.delete_table_metadata_with_retention_and_generation( + table_id, + table_name, + table_route_value, + region_wal_options, + DroppedTableLifecycle { + dropped_at, + retention_expires_at, + drop_generation: None, + }, + ) + .await + } + + /// Deletes metadata logically and stores fixed soft-drop lifecycle markers. + pub async fn delete_table_metadata_with_retention_and_generation( + &self, + table_id: TableId, + table_name: &TableName, + table_route_value: &TableRouteValue, + region_wal_options: &HashMap, + lifecycle: DroppedTableLifecycle<'_>, ) -> Result<()> { let keys = self.table_metadata_keys(table_id, table_name, table_route_value, region_wal_options)?; - self.tombstone_manager.create(keys).await.map(|_| ()) + if let Some(dropped_at) = lifecycle.dropped_at { + let mut markers = vec![( + dropped_at_key(table_id), + dropped_at.to_string().into_bytes(), + )]; + if let Some(retention_expires_at) = lifecycle.retention_expires_at { + markers.push(( + retention_expires_at_key(table_id), + retention_expires_at.to_string().into_bytes(), + )); + } + if let Some(drop_generation) = lifecycle.drop_generation { + markers.push(( + drop_generation_key(table_id), + drop_generation.as_bytes().to_vec(), + )); + } + self.tombstone_manager + .create_with_markers(keys, markers) + .await + .map(|_| ()) + } else { + self.tombstone_manager.create(keys).await.map(|_| ()) + } } /// Lists dropped tables from tombstoned table-name entries. @@ -1016,8 +1125,39 @@ impl TableMetadataManager { dropped_tables.push(DroppedTableName { table_id, table_name, + dropped_at: None, + retention_expires_at: None, + drop_generation: None, }); } + if dropped_tables.is_empty() { + return Ok(dropped_tables); + } + + let marker_keys = dropped_tables + .iter() + .flat_map(|table| { + vec![ + dropped_at_key(table.table_id), + retention_expires_at_key(table.table_id), + drop_generation_key(table.table_id), + ] + }) + .collect::>(); + let marker_values = self.tombstone_manager.batch_get(&marker_keys).await?; + for table in &mut dropped_tables { + table.dropped_at = Self::parse_dropped_at( + table.table_id, + marker_values.get(&dropped_at_key(table.table_id)), + )?; + table.retention_expires_at = Self::parse_retention_expires_at( + table.table_id, + marker_values.get(&retention_expires_at_key(table.table_id)), + )?; + table.drop_generation = Self::parse_drop_generation( + marker_values.get(&drop_generation_key(table.table_id)), + ); + } Ok(dropped_tables) } @@ -1049,6 +1189,38 @@ impl TableMetadataManager { self.get_dropped_table_metadata(table_id, None).await } + /// Returns whether an automatic purge has durably claimed the dropped table. + pub async fn is_dropped_table_purging(&self, table_id: TableId) -> Result { + self.dropped_table_purge_claim(table_id) + .await + .map(|claim| claim.is_some()) + } + + /// Returns the soft-drop generation claimed by an automatic purge. + pub async fn dropped_table_purge_claim(&self, table_id: TableId) -> Result> { + self.tombstone_manager + .get(&purging_key(table_id)) + .await + .map(|marker| marker.map(|marker| String::from_utf8_lossy(&marker.value).into_owned())) + } + + /// Durably claims a dropped table before automatic purge starts cleaning its regions. + /// The caller MUST hold the table lock and revalidate the tombstone before calling this. + pub async fn mark_dropped_table_purging( + &self, + table_id: TableId, + drop_generation: Option<&str>, + ) -> Result<()> { + self.kv_backend + .put( + PutRequest::new() + .with_key(self.tombstone_manager.to_tombstone(&purging_key(table_id))) + .with_value(drop_generation.unwrap_or_default()), + ) + .await?; + Ok(()) + } + /// Deletes metadata tombstone for table **permanently**. /// The caller MUST ensure it has the exclusive access to `TableNameKey`. pub async fn delete_table_metadata_tombstone( @@ -1061,7 +1233,15 @@ impl TableMetadataManager { let table_metadata_keys = self.table_metadata_keys(table_id, table_name, table_route_value, region_wal_options)?; self.tombstone_manager - .delete(table_metadata_keys) + .delete_with_markers( + table_metadata_keys, + vec![ + dropped_at_key(table_id), + retention_expires_at_key(table_id), + drop_generation_key(table_id), + purging_key(table_id), + ], + ) .await .map(|_| ()) } @@ -1077,7 +1257,17 @@ impl TableMetadataManager { ) -> Result<()> { let keys = self.table_metadata_keys(table_id, table_name, table_route_value, region_wal_options)?; - self.tombstone_manager.restore(keys).await.map(|_| ()) + self.tombstone_manager + .restore_with_markers( + keys, + vec![ + dropped_at_key(table_id), + retention_expires_at_key(table_id), + drop_generation_key(table_id), + ], + ) + .await + .map(|_| ()) } /// Deletes metadata for table **permanently**. @@ -1138,6 +1328,23 @@ impl TableMetadataManager { let region_wal_options = self .dropped_region_wal_options(table_id, &table_route_value) .await?; + let dropped_at_key = dropped_at_key(table_id); + let retention_expires_at_key = retention_expires_at_key(table_id); + let drop_generation_key = drop_generation_key(table_id); + let marker_values = self + .tombstone_manager + .batch_get(&[ + dropped_at_key.clone(), + retention_expires_at_key.clone(), + drop_generation_key.clone(), + ]) + .await?; + let dropped_at = Self::parse_dropped_at(table_id, marker_values.get(&dropped_at_key))?; + let retention_expires_at = Self::parse_retention_expires_at( + table_id, + marker_values.get(&retention_expires_at_key), + )?; + let drop_generation = Self::parse_drop_generation(marker_values.get(&drop_generation_key)); Ok(Some(DroppedTableMetadata { table_id, @@ -1145,9 +1352,41 @@ impl TableMetadataManager { table_info_value, table_route_value, region_wal_options, + dropped_at, + retention_expires_at, + drop_generation, })) } + fn parse_dropped_at(table_id: TableId, kv: Option<&KeyValue>) -> Result> { + Self::parse_timestamp_marker(table_id, "dropped timestamp", kv) + } + + fn parse_retention_expires_at(table_id: TableId, kv: Option<&KeyValue>) -> Result> { + Self::parse_timestamp_marker(table_id, "retention deadline", kv) + } + + fn parse_drop_generation(kv: Option<&KeyValue>) -> Option { + kv.map(|kv| String::from_utf8_lossy(&kv.value).into_owned()) + } + + fn parse_timestamp_marker( + table_id: TableId, + description: &str, + kv: Option<&KeyValue>, + ) -> Result> { + let Some(kv) = kv else { + return Ok(None); + }; + let value = String::from_utf8_lossy(&kv.value); + value.parse().map(Some).map_err(|err| { + error::UnexpectedSnafu { + err_msg: format!("Invalid {description} '{value}' for table id {table_id}: {err}"), + } + .build() + }) + } + /// Rebuilds region WAL options from tombstoned datanode-table entries. async fn dropped_region_wal_options( &self, @@ -1633,6 +1872,7 @@ impl_optional_metadata_value! { mod tests { use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use bytes::Bytes; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; @@ -1662,9 +1902,10 @@ mod tests { use crate::kv_backend::KvBackend; use crate::kv_backend::memory::MemoryKvBackend; use crate::kv_backend::read_only::ReadOnlyKvBackend; + use crate::kv_backend::test_util::MockKvBackend; use crate::peer::Peer; use crate::rpc::router::{LeaderState, Region, RegionRoute, region_distribution}; - use crate::rpc::store::{PutRequest, RangeRequest}; + use crate::rpc::store::{BatchGetResponse, PutRequest, RangeRequest, RangeResponse}; use crate::wal_provider::{RegionWalOptions, WalProvider}; #[test] @@ -1837,6 +2078,7 @@ mod tests { &table_name, &table_route_value, ®ion_wal_options, + None, ) .await .unwrap(); @@ -2188,6 +2430,7 @@ mod tests { &table_name, table_route_value, ®ion_wal_options, + None, ) .await .unwrap(); @@ -2198,6 +2441,7 @@ mod tests { &table_name, table_route_value, ®ion_wal_options, + None, ) .await .unwrap(); @@ -2927,7 +3171,7 @@ mod tests { let table_name = TableName::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, table_name); let table_route_value = TableRouteValue::physical(region_routes.clone()); table_metadata_manager - .delete_table_metadata(table_id, &table_name, &table_route_value, &options) + .delete_table_metadata(table_id, &table_name, &table_route_value, &options, None) .await .unwrap(); table_metadata_manager @@ -2996,6 +3240,92 @@ mod tests { ); } + #[tokio::test] + async fn test_list_dropped_tables_batches_timestamp_lookup() { + let mem_kv = Arc::new(MemoryKvBackend::default()); + let table_metadata_manager = TableMetadataManager::new(mem_kv.clone()); + for (table_id, table_name, dropped_at) in + [(1025, "legacy", None), (1026, "marked", Some(1234))] + { + let task = test_create_table_task(table_name, table_id); + let table_name = task.table_name(); + let table_route = TableRouteValue::physical(vec![]); + table_metadata_manager + .create_table_metadata(task.table_info, table_route.clone(), HashMap::new()) + .await + .unwrap(); + table_metadata_manager + .delete_table_metadata( + table_id, + &table_name, + &table_route, + &HashMap::new(), + dropped_at, + ) + .await + .unwrap(); + } + let all_tombstones = mem_kv + .range(RangeRequest::new().with_prefix("__tombstone/")) + .await + .unwrap() + .kvs; + let range_calls = Arc::new(AtomicUsize::new(0)); + let batch_get_calls = Arc::new(AtomicUsize::new(0)); + let mock = MockKvBackend { + range_fn: Some({ + let all_tombstones = all_tombstones.clone(); + let range_calls = range_calls.clone(); + Arc::new(move |req| { + range_calls.fetch_add(1, Ordering::Relaxed); + let kvs = all_tombstones + .iter() + .filter(|kv| { + if req.range_end.is_empty() { + kv.key == req.key + } else { + kv.key >= req.key && kv.key < req.range_end + } + }) + .cloned() + .collect(); + Ok(RangeResponse { kvs, more: false }) + }) + }), + batch_get_fn: Some({ + let batch_get_calls = batch_get_calls.clone(); + Arc::new(move |req| { + batch_get_calls.fetch_add(1, Ordering::Relaxed); + let kvs = all_tombstones + .iter() + .filter(|kv| req.keys.contains(&kv.key)) + .cloned() + .collect(); + Ok(BatchGetResponse { kvs }) + }) + }), + put_fn: None, + batch_put_fn: None, + delete_range_fn: None, + batch_delete_fn: None, + txn: None, + max_txn_ops: None, + }; + let table_metadata_manager = TableMetadataManager::new(Arc::new(mock)); + + let dropped_tables = table_metadata_manager.list_dropped_tables().await.unwrap(); + + assert_eq!( + dropped_tables + .iter() + .map(|table| (table.table_id, table.dropped_at)) + .collect::>(), + vec![(1025, None), (1026, Some(1234))] + ); + assert_eq!(range_calls.load(Ordering::Relaxed), 1); + assert_eq!(batch_get_calls.load(Ordering::Relaxed), 1); + } + #[tokio::test] async fn test_dropped_table_lookup_survives_live_name_recreation() { let dropped_table_id = 1025; diff --git a/src/common/meta/src/key/tombstone.rs b/src/common/meta/src/key/tombstone.rs index 01fa4405776..284737dad0b 100644 --- a/src/common/meta/src/key/tombstone.rs +++ b/src/common/meta/src/key/tombstone.rs @@ -190,6 +190,7 @@ impl TombstoneManager { keys: &[Vec], dest_keys: &[Vec], require_dest_not_exists: bool, + extra_ops: Vec, ) -> Result { ensure!( keys.len() == dest_keys.len(), @@ -213,6 +214,9 @@ impl TombstoneManager { .into_iter() .map(|kv| (kv.key, kv.value)) .collect::>(); + if results.is_empty() && extra_ops.is_empty() { + return Ok(0); + } const MAX_RETRIES: usize = 8; for _ in 0..MAX_RETRIES { @@ -228,6 +232,10 @@ impl TombstoneManager { (txn, key.clone()) }) .unzip(); + let mut txns = txns; + if !extra_ops.is_empty() { + txns.push(Txn::new().and_then(extra_ops.clone())); + } let mut resp = self.kv_backend.txn(Txn::merge_all(txns)).await?; if resp.succeeded { return Ok(keys.len()); @@ -272,14 +280,30 @@ impl TombstoneManager { self.kv_backend.max_txn_ops() } + async fn execute_extra_ops(&self, extra_ops: Vec) -> Result<()> { + let max_txn_ops = self.max_txn_ops(); + ensure!( + max_txn_ops > 0, + error::UnexpectedSnafu { + err_msg: "max_txn_ops must be greater than 0".to_string(), + } + ); + for chunk in extra_ops.chunks(max_txn_ops) { + self.move_values_inner(&[], &[], false, chunk.to_vec()) + .await?; + } + Ok(()) + } + /// Moves values to `dest_key`. /// /// Returns the number of keys that were moved. - async fn move_values( + async fn move_values_with_extra( &self, keys: Vec>, dest_keys: Vec>, require_dest_not_exists: bool, + mut extra_ops: Vec, ) -> Result { ensure!( keys.len() == dest_keys.len(), @@ -292,6 +316,9 @@ impl TombstoneManager { } ); if keys.is_empty() { + if !extra_ops.is_empty() { + self.execute_extra_ops(extra_ops).await?; + } return Ok(0); } let txn_ops_per_key = if require_dest_not_exists { @@ -308,8 +335,15 @@ impl TombstoneManager { ) } ); - let chunk_size = max_txn_ops / txn_ops_per_key; - if keys.len() > chunk_size { + let merge_extra_ops = + !extra_ops.is_empty() && max_txn_ops.saturating_sub(txn_ops_per_key) >= extra_ops.len(); + let reserved_ops = if merge_extra_ops { extra_ops.len() } else { 0 }; + let chunk_size = (max_txn_ops - reserved_ops) / txn_ops_per_key; + let total_ops = keys + .len() + .saturating_mul(txn_ops_per_key) + .saturating_add(extra_ops.len()); + if keys.len() > chunk_size || total_ops > max_txn_ops { debug!( "Moving values with multiple chunks, keys len: {}, chunk_size: {}", keys.len(), @@ -318,18 +352,39 @@ impl TombstoneManager { let mut moved_keys = 0; let keys_chunks = keys.chunks(chunk_size).collect::>(); let dest_keys_chunks = dest_keys.chunks(chunk_size).collect::>(); - for (keys, dest_keys) in keys_chunks.into_iter().zip(dest_keys_chunks) { + let num_chunks = keys_chunks.len(); + for (index, (keys, dest_keys)) in + keys_chunks.into_iter().zip(dest_keys_chunks).enumerate() + { + let chunk_extra_ops = if merge_extra_ops && index + 1 == num_chunks { + std::mem::take(&mut extra_ops) + } else { + vec![] + }; moved_keys += self - .move_values_inner(keys, dest_keys, require_dest_not_exists) + .move_values_inner(keys, dest_keys, require_dest_not_exists, chunk_extra_ops) .await?; } + if !extra_ops.is_empty() { + self.execute_extra_ops(extra_ops).await?; + } Ok(moved_keys) } else { - self.move_values_inner(&keys, &dest_keys, require_dest_not_exists) + self.move_values_inner(&keys, &dest_keys, require_dest_not_exists, extra_ops) .await } } + async fn move_values( + &self, + keys: Vec>, + dest_keys: Vec>, + require_dest_not_exists: bool, + ) -> Result { + self.move_values_with_extra(keys, dest_keys, require_dest_not_exists, vec![]) + .await + } + /// Creates tombstones for keys. /// /// Preforms to: @@ -349,6 +404,39 @@ impl TombstoneManager { self.move_values(keys, dest_keys, false).await } + /// Creates tombstones and stores an additional marker in the tombstone namespace. + pub async fn create_with_marker( + &self, + keys: Vec>, + marker_key: Vec, + marker_value: Vec, + ) -> Result { + self.create_with_markers(keys, vec![(marker_key, marker_value)]) + .await + } + + /// Creates tombstones and stores additional markers in the tombstone namespace. + pub async fn create_with_markers( + &self, + keys: Vec>, + markers: Vec<(Vec, Vec)>, + ) -> Result { + let (keys, dest_keys): (Vec<_>, Vec<_>) = keys + .into_iter() + .map(|key| { + let tombstone_key = self.to_tombstone(&key); + (key, tombstone_key) + }) + .unzip(); + let marker_ops = markers + .into_iter() + .map(|(key, value)| TxnOp::Put(self.to_tombstone(&key), value)) + .collect(); + + self.move_values_with_extra(keys, dest_keys, false, marker_ops) + .await + } + /// Restores tombstones for keys. /// /// Preforms to: @@ -368,6 +456,46 @@ impl TombstoneManager { self.move_values(keys, dest_keys, true).await } + /// Restores tombstones and deletes an associated marker. + pub async fn restore_with_marker( + &self, + keys: Vec>, + marker_key: Vec, + ) -> Result { + let (keys, dest_keys): (Vec<_>, Vec<_>) = keys + .into_iter() + .map(|key| { + let tombstone_key = self.to_tombstone(&key); + (tombstone_key, key) + }) + .unzip(); + let marker_key = self.to_tombstone(&marker_key); + + self.move_values_with_extra(keys, dest_keys, true, vec![TxnOp::Delete(marker_key)]) + .await + } + + /// Restores tombstones and deletes associated markers. + pub async fn restore_with_markers( + &self, + keys: Vec>, + marker_keys: Vec>, + ) -> Result { + let (keys, dest_keys): (Vec<_>, Vec<_>) = keys + .into_iter() + .map(|key| { + let tombstone_key = self.to_tombstone(&key); + (tombstone_key, key) + }) + .unzip(); + let marker_ops = marker_keys + .into_iter() + .map(|key| TxnOp::Delete(self.to_tombstone(&key))) + .collect(); + self.move_values_with_extra(keys, dest_keys, true, marker_ops) + .await + } + /// Deletes tombstones values for the specified `keys`. /// /// Returns the number of keys that were deleted. @@ -385,6 +513,26 @@ impl TombstoneManager { Ok(num_keys) } + + /// Deletes tombstones and an associated marker. + pub async fn delete_with_marker( + &self, + mut keys: Vec>, + marker_key: Vec, + ) -> Result { + keys.push(marker_key); + self.delete(keys).await + } + + /// Deletes tombstones and associated markers. + pub async fn delete_with_markers( + &self, + mut keys: Vec>, + marker_keys: Vec>, + ) -> Result { + keys.extend(marker_keys); + self.delete(keys).await + } } #[cfg(test)] @@ -392,10 +540,12 @@ mod tests { use std::any::Any; use std::collections::HashMap; - use std::sync::Arc; + use std::sync::{Arc, Mutex}; use crate::error::{Error, Result}; - use crate::key::tombstone::TombstoneManager; + use crate::key::tombstone::{ + MOVE_VALUE_TXN_OPS_PER_KEY, RESTORE_VALUE_TXN_OPS_PER_KEY, TombstoneManager, + }; use crate::kv_backend::memory::MemoryKvBackend; use crate::kv_backend::txn::{Txn, TxnRequest, TxnResponse}; use crate::kv_backend::{KvBackend, TxnService}; @@ -408,6 +558,7 @@ mod tests { struct TxnOpLimitKvBackend { inner: Arc>, max_txn_ops: usize, + txn_op_counts: Mutex>, } #[async_trait::async_trait] @@ -426,6 +577,7 @@ mod tests { "txn ops {txn_ops} exceeds limit {}", self.max_txn_ops ); + self.txn_op_counts.lock().unwrap().push(txn_ops); self.inner.txn(txn).await } @@ -746,6 +898,7 @@ mod tests { let kv_backend = Arc::new(TxnOpLimitKvBackend { inner: inner.clone(), max_txn_ops: 6, + txn_op_counts: Mutex::new(Vec::new()), }); let tombstone_manager = TombstoneManager::new(kv_backend); let kvs = HashMap::from([ @@ -778,6 +931,7 @@ mod tests { let kv_backend = Arc::new(TxnOpLimitKvBackend { inner: inner.clone(), max_txn_ops: 5, + txn_op_counts: Mutex::new(Vec::new()), }); let tombstone_manager = TombstoneManager::new(kv_backend); let key = b"foo".to_vec(); @@ -805,6 +959,7 @@ mod tests { let kv_backend = Arc::new(TxnOpLimitKvBackend { inner: inner.clone(), max_txn_ops: 4, + txn_op_counts: Mutex::new(Vec::new()), }); let tombstone_manager = TombstoneManager::new(kv_backend); let kvs = HashMap::from([ @@ -831,6 +986,195 @@ mod tests { assert_eq!(kvs.len(), moved); } + #[tokio::test] + async fn test_create_with_marker_uses_final_transaction_at_exact_limit() { + let inner = Arc::new(MemoryKvBackend::default()); + let kv_backend = Arc::new(TxnOpLimitKvBackend { + inner: inner.clone(), + max_txn_ops: MOVE_VALUE_TXN_OPS_PER_KEY, + txn_op_counts: Mutex::new(Vec::new()), + }); + let tombstone_manager = TombstoneManager::new(kv_backend.clone()); + for (key, value) in [(b"bar".as_slice(), b"baz".as_slice()), (b"foo", b"hi")] { + inner + .put(PutRequest::new().with_key(key).with_value(value)) + .await + .unwrap(); + } + + tombstone_manager + .create_with_marker( + vec![b"bar".to_vec(), b"foo".to_vec()], + b"__dropped_at/42".to_vec(), + b"1234".to_vec(), + ) + .await + .unwrap(); + tombstone_manager + .create_with_marker( + vec![b"bar".to_vec(), b"foo".to_vec()], + b"__dropped_at/42".to_vec(), + b"1234".to_vec(), + ) + .await + .unwrap(); + + assert_eq!( + kv_backend.txn_op_counts.lock().unwrap().as_slice(), + &[4, 4, 1, 1] + ); + assert_eq!( + inner + .get(&tombstone_manager.to_tombstone(b"__dropped_at/42")) + .await + .unwrap() + .unwrap() + .value, + b"1234" + ); + } + + #[tokio::test] + async fn test_create_with_markers_chunks_extra_ops_after_key_move() { + let inner = Arc::new(MemoryKvBackend::default()); + let kv_backend = Arc::new(TxnOpLimitKvBackend { + inner: inner.clone(), + max_txn_ops: MOVE_VALUE_TXN_OPS_PER_KEY + 1, + txn_op_counts: Mutex::new(Vec::new()), + }); + let tombstone_manager = TombstoneManager::new(kv_backend.clone()); + inner + .put(PutRequest::new().with_key(b"foo").with_value(b"bar")) + .await + .unwrap(); + + tombstone_manager + .create_with_markers( + vec![b"foo".to_vec()], + (0..6) + .map(|index| (format!("marker-{index}").into_bytes(), vec![index])) + .collect(), + ) + .await + .unwrap(); + + assert_eq!( + kv_backend.txn_op_counts.lock().unwrap().as_slice(), + &[MOVE_VALUE_TXN_OPS_PER_KEY, 5, 1] + ); + } + + #[tokio::test] + async fn test_create_with_markers_chunks_marker_only_transactions() { + let inner = Arc::new(MemoryKvBackend::default()); + let kv_backend = Arc::new(TxnOpLimitKvBackend { + inner, + max_txn_ops: 2, + txn_op_counts: Mutex::new(Vec::new()), + }); + let tombstone_manager = TombstoneManager::new(kv_backend.clone()); + let markers = (0..5) + .map(|index| (format!("marker-{index}").into_bytes(), vec![index])) + .collect(); + + tombstone_manager + .create_with_markers(vec![], markers) + .await + .unwrap(); + + assert_eq!( + kv_backend.txn_op_counts.lock().unwrap().as_slice(), + &[2, 2, 1] + ); + } + + #[tokio::test] + async fn test_restore_with_marker_uses_final_transaction_at_exact_limit() { + let inner = Arc::new(MemoryKvBackend::default()); + let kv_backend = Arc::new(TxnOpLimitKvBackend { + inner: inner.clone(), + max_txn_ops: RESTORE_VALUE_TXN_OPS_PER_KEY, + txn_op_counts: Mutex::new(Vec::new()), + }); + let tombstone_manager = TombstoneManager::new(kv_backend.clone()); + for (key, value) in [(b"bar".as_slice(), b"baz".as_slice()), (b"foo", b"hi")] { + inner + .put( + PutRequest::new() + .with_key(tombstone_manager.to_tombstone(key)) + .with_value(value), + ) + .await + .unwrap(); + } + inner + .put( + PutRequest::new() + .with_key(tombstone_manager.to_tombstone(b"__dropped_at/42")) + .with_value(b"1234"), + ) + .await + .unwrap(); + + tombstone_manager + .restore_with_marker( + vec![b"bar".to_vec(), b"foo".to_vec()], + b"__dropped_at/42".to_vec(), + ) + .await + .unwrap(); + tombstone_manager + .restore_with_marker( + vec![b"bar".to_vec(), b"foo".to_vec()], + b"__dropped_at/42".to_vec(), + ) + .await + .unwrap(); + + assert_eq!( + kv_backend.txn_op_counts.lock().unwrap().as_slice(), + &[6, 6, 1, 1] + ); + assert!( + inner + .get(&tombstone_manager.to_tombstone(b"__dropped_at/42")) + .await + .unwrap() + .is_none() + ); + } + + #[tokio::test] + async fn test_delete_with_marker_removes_marker_with_limited_backend() { + let inner = Arc::new(MemoryKvBackend::default()); + let kv_backend = Arc::new(TxnOpLimitKvBackend { + inner: inner.clone(), + max_txn_ops: MOVE_VALUE_TXN_OPS_PER_KEY, + txn_op_counts: Mutex::new(Vec::new()), + }); + let tombstone_manager = TombstoneManager::new(kv_backend); + for key in [b"bar".as_slice(), b"foo", b"__dropped_at/42"] { + inner + .put( + PutRequest::new() + .with_key(tombstone_manager.to_tombstone(key)) + .with_value(b"value"), + ) + .await + .unwrap(); + } + + tombstone_manager + .delete_with_marker( + vec![b"bar".to_vec(), b"foo".to_vec()], + b"__dropped_at/42".to_vec(), + ) + .await + .unwrap(); + + assert!(inner.is_empty()); + } + #[tokio::test] async fn test_move_values_with_non_exists_values() { let kv_backend = Arc::new(MemoryKvBackend::default()); diff --git a/src/common/meta/src/test_util.rs b/src/common/meta/src/test_util.rs index 2300679b337..24c4b7b759c 100644 --- a/src/common/meta/src/test_util.rs +++ b/src/common/meta/src/test_util.rs @@ -207,6 +207,7 @@ pub fn new_ddl_context_with_kv_backend( flow_metadata_manager, region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl), soft_drop_enabled: false, + soft_drop_retention: Some(std::time::Duration::from_secs(7 * 24 * 60 * 60)), } } diff --git a/src/common/query/src/lib.rs b/src/common/query/src/lib.rs index bdba49eaf70..e70c4d97e7b 100644 --- a/src/common/query/src/lib.rs +++ b/src/common/query/src/lib.rs @@ -17,6 +17,7 @@ pub mod error; pub mod logical_plan; pub mod native_histogram; pub mod prelude; +pub mod prometheus; pub mod request; pub mod stream; #[cfg(any(test, feature = "testing"))] diff --git a/src/common/query/src/prometheus.rs b/src/common/query/src/prometheus.rs new file mode 100644 index 00000000000..709aa5d7693 --- /dev/null +++ b/src/common/query/src/prometheus.rs @@ -0,0 +1,37 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +/// Canonical Prometheus stale-marker NaN bit pattern. +pub const PROMETHEUS_STALE_NAN_BITS: u64 = 0x7ff0_0000_0000_0002; + +/// Returns whether `value` is the canonical Prometheus stale-marker NaN. +#[inline] +pub fn is_prometheus_stale_nan(value: f64) -> bool { + value.to_bits() == PROMETHEUS_STALE_NAN_BITS +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn recognizes_only_the_canonical_stale_marker() { + assert!(is_prometheus_stale_nan(f64::from_bits( + 0x7ff0_0000_0000_0002 + ))); + assert!(!is_prometheus_stale_nan(f64::from_bits( + 0x7ff8_0000_0000_0000 + ))); + } +} diff --git a/src/common/recordbatch/Cargo.toml b/src/common/recordbatch/Cargo.toml index efc6b6f60e9..59929ba3477 100644 --- a/src/common/recordbatch/Cargo.toml +++ b/src/common/recordbatch/Cargo.toml @@ -34,3 +34,7 @@ tokio.workspace = true [[bench]] name = "iter_record_batch_rows" harness = false + +[[bench]] +name = "memory_accounting" +harness = false diff --git a/src/common/recordbatch/benches/memory_accounting.rs b/src/common/recordbatch/benches/memory_accounting.rs new file mode 100644 index 00000000000..05fb63555a4 --- /dev/null +++ b/src/common/recordbatch/benches/memory_accounting.rs @@ -0,0 +1,117 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::hint::black_box; +use std::sync::Arc; + +use common_recordbatch::RecordBatch; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datatypes::data_type::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema}; +use datatypes::vectors::{Int32Vector, StringVector, UInt64Vector, VectorRef}; + +const LARGE_ROWS: usize = 8_192; + +fn batch(columns: Vec<(&str, ConcreteDataType, bool, VectorRef)>) -> RecordBatch { + let schema = Arc::new(Schema::new( + columns + .iter() + .map(|(name, data_type, nullable, _)| { + ColumnSchema::new(*name, data_type.clone(), *nullable) + }) + .collect(), + )); + let vectors: Vec = columns + .into_iter() + .map(|(_, _, _, vector)| vector) + .collect(); + RecordBatch::new(schema, vectors).unwrap() +} + +fn primitive_string_batch(rows: usize) -> RecordBatch { + let strings = (0..rows) + .map(|row| (row % 5 != 0).then(|| format!("value-{row:08}-with-payload"))) + .collect::>(); + batch(vec![ + ( + "i32", + ConcreteDataType::int32_datatype(), + false, + Arc::new(Int32Vector::from_slice( + (0..rows).map(|value| value as i32).collect::>(), + )), + ), + ( + "u64", + ConcreteDataType::uint64_datatype(), + false, + Arc::new(UInt64Vector::from_slice( + (0..rows).map(|value| value as u64).collect::>(), + )), + ), + ( + "string", + ConcreteDataType::string_datatype(), + true, + Arc::new(StringVector::from(strings)), + ), + ]) +} + +fn bench_one_call(c: &mut Criterion) { + let mut group = c.benchmark_group("memory_accounting/one_call"); + for rows in [128, LARGE_ROWS] { + let batch = primitive_string_batch(rows); + group.bench_function(BenchmarkId::new("buffer", rows), |b| { + b.iter(|| black_box(black_box(&batch).buffer_memory_size())) + }); + group.bench_function(BenchmarkId::new("logical", rows), |b| { + b.iter(|| black_box(black_box(&batch).logical_slice_memory_size())) + }); + } + group.finish(); +} + +fn bench_shared_buffer_slices(c: &mut Criterion) { + let backing = primitive_string_batch(LARGE_ROWS); + let mut group = c.benchmark_group("memory_accounting/shared_buffer_slices"); + for slice_len in [1, 16] { + let slices = (0..512) + .map(|index| backing.slice(index * slice_len, slice_len).unwrap()) + .collect::>(); + group.throughput(Throughput::Elements(slices.len() as u64)); + group.bench_function(BenchmarkId::new("buffer", slice_len), |b| { + b.iter(|| { + let total = black_box(&slices) + .iter() + .map(RecordBatch::buffer_memory_size) + .fold(0usize, usize::saturating_add); + black_box(total) + }) + }); + group.bench_function(BenchmarkId::new("logical", slice_len), |b| { + b.iter(|| { + let total = black_box(&slices) + .iter() + .map(RecordBatch::logical_slice_memory_size) + .fold(0usize, usize::saturating_add); + black_box(total) + }) + }); + } + group.finish(); +} + +criterion_group!(benches, bench_one_call, bench_shared_buffer_slices); +criterion_main!(benches); diff --git a/src/common/recordbatch/src/lib.rs b/src/common/recordbatch/src/lib.rs index 3db3d9c08b6..46375d8fd77 100644 --- a/src/common/recordbatch/src/lib.rs +++ b/src/common/recordbatch/src/lib.rs @@ -769,7 +769,7 @@ impl MemoryTrackedStream { batch: RecordBatch, cx: &mut Context<'_>, ) -> Poll>> { - let additional = batch.buffer_memory_size(); + let additional = batch.logical_slice_memory_size(); let tracker = self.ready_tracker_mut(); if let Err(error) = tracker.try_track(additional) { @@ -860,6 +860,35 @@ mod tests { as usize } + #[tokio::test] + async fn test_memory_tracked_stream_charges_logical_slice_size() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "payload", + ConcreteDataType::string_datatype(), + false, + )])); + let payloads: Vec<_> = (0..1024) + .map(|value| format!("payload-{value:04}")) + .collect(); + let batch = RecordBatch::new(schema, vec![Arc::new(StringVector::from(payloads)) as _]) + .unwrap() + .slice(512, 1) + .unwrap(); + let expected_bytes = aligned_tracked_bytes(batch.logical_slice_memory_size()); + assert!(expected_bytes < aligned_tracked_bytes(batch.buffer_memory_size())); + let tracker = QueryMemoryTracker::builder(MB, OnExhaustedPolicy::Fail).build(); + let mut stream = MemoryTrackedStream::new( + RecordBatches::try_new(batch.schema.clone(), vec![batch]) + .unwrap() + .as_stream(), + tracker.clone(), + ); + + stream.next().await.unwrap().unwrap(); + + assert_eq!(tracker.current(), expected_bytes); + } + #[test] fn test_recordbatches_try_from_columns() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( @@ -1052,7 +1081,7 @@ mod tests { }) .build(); let batch = large_string_batch(700 * 1024); - let expected_bytes = aligned_tracked_bytes(batch.buffer_memory_size()); + let expected_bytes = aligned_tracked_bytes(batch.logical_slice_memory_size()); let mut stream1 = MemoryTrackedStream::new( RecordBatches::try_new(batch.schema.clone(), vec![batch.clone()]) diff --git a/src/common/recordbatch/src/recordbatch.rs b/src/common/recordbatch/src/recordbatch.rs index 7280772d05d..77a36836d6a 100644 --- a/src/common/recordbatch/src/recordbatch.rs +++ b/src/common/recordbatch/src/recordbatch.rs @@ -269,6 +269,28 @@ impl RecordBatch { .sum() } + /// Returns the logical memory size of this batch's array slices. + /// + /// This sums Arrow's logical visible slice buffers rather than the capacity of their shared + /// backing buffers. View out-of-line payloads and nested custom payloads are not separately + /// traversed or accounted. It is not an exact measure of live physical memory. If Arrow cannot + /// calculate a slice's size, the full buffer size is used conservatively. + /// + /// Mito's current scan paths do not produce top-level View arrays. If they do in the future, + /// their out-of-line payload accounting must be reassessed here. + pub fn logical_slice_memory_size(&self) -> usize { + self.df_record_batch + .columns() + .iter() + .fold(0, |total, array| { + let array_size = array + .to_data() + .get_slice_memory_size() + .unwrap_or_else(|_| array.get_buffer_memory_size()); + total.saturating_add(array_size) + }) + } + /// Iterate the values as strings in the column at index `i`. /// /// Note that if the underlying array is not a valid GreptimeDB vector, an empty iterator is @@ -382,13 +404,14 @@ fn maybe_align_json_array_with_schema( mod tests { use std::sync::Arc; - use datatypes::arrow::array::{AsArray, BinaryArray, UInt32Array}; + use datatypes::arrow::array::{ + AsArray, BinaryArray, StringArray, StringViewArray, UInt32Array, + }; use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema, UInt32Type}; - use datatypes::arrow_array::StringArray; use datatypes::data_type::ConcreteDataType; use datatypes::extension::json::{JsonExtensionType, JsonMetadata}; use datatypes::schema::{ColumnSchema, Schema}; - use datatypes::vectors::{StringVector, UInt32Vector}; + use datatypes::vectors::{BinaryVector, StringVector, UInt32Vector}; use super::*; @@ -475,6 +498,97 @@ mod tests { assert!(recordbatch.slice(1, 5).is_err()); } + #[test] + fn test_logical_slice_memory_size_for_visible_primitive_string_binary_slices() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("numbers", ConcreteDataType::uint32_datatype(), false), + ColumnSchema::new("strings", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("binary", ConcreteDataType::binary_datatype(), true), + ])); + let numbers: Vec<_> = (0..1024).collect(); + let strings = (0..1024) + .map(|value| (value % 3 != 0).then(|| format!("value-{value}"))) + .collect::>(); + let binary = (0_u32..1024) + .map(|value| (value % 3 != 0).then(|| value.to_le_bytes().to_vec())) + .collect::>(); + let columns: Vec = vec![ + Arc::new(UInt32Vector::from_slice(numbers)), + Arc::new(StringVector::from(strings)), + Arc::new(BinaryVector::from(binary)), + ]; + let batch = RecordBatch::new(schema, columns).unwrap(); + let slice = batch.slice(511, 3).unwrap(); + + assert!(slice.columns().iter().any(|column| column.null_count() > 0)); + assert!(slice.logical_slice_memory_size() < slice.buffer_memory_size()); + assert_eq!( + slice.logical_slice_memory_size(), + slice + .columns() + .iter() + .map(|column| column.to_data().get_slice_memory_size().unwrap()) + .sum::() + ); + } + + #[test] + fn test_logical_slice_memory_size_for_many_shared_buffer_slices() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "strings", + ConcreteDataType::string_datatype(), + false, + )])); + let strings = (0..1024) + .map(|value| format!("shared-value-{value}")) + .collect::>(); + let backing = RecordBatch::new( + schema, + vec![Arc::new(StringVector::from(strings)) as VectorRef], + ) + .unwrap(); + let slices = (0..128) + .map(|index| backing.slice(index * 4, 2).unwrap()) + .collect::>(); + let logical_total = slices + .iter() + .map(RecordBatch::logical_slice_memory_size) + .sum::(); + let buffer_total = slices + .iter() + .map(RecordBatch::buffer_memory_size) + .sum::(); + + assert!(logical_total < buffer_total); + assert!(slices.iter().all(|slice| { + slice.logical_slice_memory_size() + == slice.column(0).to_data().get_slice_memory_size().unwrap() + })); + } + + #[test] + fn test_logical_slice_memory_size_uses_arrow_slice_scope_for_views() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "strings", + ConcreteDataType::utf8_view_datatype(), + false, + )])); + let columns: Vec = + vec![Arc::new(StringVector::from(StringViewArray::from(vec![ + "unrelated backing payload", + "visible string view payload", + ])))]; + let batch = RecordBatch::new(schema, columns) + .unwrap() + .slice(1, 1) + .unwrap(); + + assert_eq!( + batch.column(0).to_data().get_slice_memory_size().unwrap(), + batch.logical_slice_memory_size() + ); + } + #[test] fn test_merge_record_batch() { let column_schemas = vec![ diff --git a/src/common/telemetry/src/logging.rs b/src/common/telemetry/src/logging.rs index 953dbbc2df9..56513531632 100644 --- a/src/common/telemetry/src/logging.rs +++ b/src/common/telemetry/src/logging.rs @@ -237,7 +237,7 @@ enum TraceState { } /// The logging options that used to initialize the logger. -#[derive(Clone, Debug, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(default)] pub struct LoggingOptions { /// The directory to store log files. If not set, logs will be written to stdout. @@ -341,20 +341,6 @@ pub enum LogFormat { Text, } -impl PartialEq for LoggingOptions { - fn eq(&self, other: &Self) -> bool { - self.dir == other.dir - && self.level == other.level - && self.enable_otlp_tracing == other.enable_otlp_tracing - && self.otlp_endpoint == other.otlp_endpoint - && self.tracing_sample_ratio == other.tracing_sample_ratio - && self.append_stdout == other.append_stdout - && self.enable_per_region_metrics == other.enable_per_region_metrics - } -} - -impl Eq for LoggingOptions {} - #[derive(Clone, Debug)] struct TraceContext { app_name: String, @@ -874,4 +860,27 @@ mod tests { let protocol: OtlpExportProtocol = serde_json::from_str(http_json).unwrap(); assert_eq!(protocol, OtlpExportProtocol::Http); } + + #[test] + fn test_logging_options_partial_eq_all_fields() { + let base = LoggingOptions::default(); + + let mut log_format = base.clone(); + log_format.log_format = LogFormat::Json; + assert_ne!(base, log_format); + + let mut max_log_files = base.clone(); + max_log_files.max_log_files += 1; + assert_ne!(base, max_log_files); + + let mut otlp_export_protocol = base.clone(); + otlp_export_protocol.otlp_export_protocol = Some(OtlpExportProtocol::Http); + assert_ne!(base, otlp_export_protocol); + + let mut otlp_headers = base.clone(); + otlp_headers + .otlp_headers + .insert("key".to_string(), "value".to_string()); + assert_ne!(base, otlp_headers); + } } diff --git a/src/common/time/src/date.rs b/src/common/time/src/date.rs index bda9105f559..38667abbcac 100644 --- a/src/common/time/src/date.rs +++ b/src/common/time/src/date.rs @@ -198,6 +198,10 @@ impl Date { pub fn negative(&self) -> Self { Self(-self.0) } + + pub fn checked_negative(&self) -> Option { + self.0.checked_neg().map(Self) + } } #[cfg(test)] diff --git a/src/common/time/src/duration.rs b/src/common/time/src/duration.rs index a1e582ca82c..b34ced336bd 100644 --- a/src/common/time/src/duration.rs +++ b/src/common/time/src/duration.rs @@ -97,6 +97,11 @@ impl Duration { self.value = -self.value; self } + + pub fn checked_negative(mut self) -> Option { + self.value = self.value.checked_neg()?; + Some(self) + } } /// Convert i64 to Duration Type. diff --git a/src/common/time/src/interval.rs b/src/common/time/src/interval.rs index 4cb8c45c55f..21bdc656d88 100644 --- a/src/common/time/src/interval.rs +++ b/src/common/time/src/interval.rs @@ -80,6 +80,10 @@ impl IntervalYearMonth { Self::new(-self.months) } + pub fn checked_negative(&self) -> Option { + self.months.checked_neg().map(Self::new) + } + pub fn to_iso8601_string(&self) -> String { IntervalFormat::from(*self).to_iso8601_string() } @@ -157,6 +161,13 @@ impl IntervalDayTime { Self::new(-self.days, -self.milliseconds) } + pub fn checked_negative(&self) -> Option { + Some(Self::new( + self.days.checked_neg()?, + self.milliseconds.checked_neg()?, + )) + } + pub fn to_iso8601_string(&self) -> String { IntervalFormat::from(*self).to_iso8601_string() } @@ -278,6 +289,14 @@ impl IntervalMonthDayNano { Self::new(-self.months, -self.days, -self.nanoseconds) } + pub fn checked_negative(&self) -> Option { + Some(Self::new( + self.months.checked_neg()?, + self.days.checked_neg()?, + self.nanoseconds.checked_neg()?, + )) + } + pub fn to_iso8601_string(&self) -> String { IntervalFormat::from(*self).to_iso8601_string() } diff --git a/src/common/time/src/time.rs b/src/common/time/src/time.rs index 00b123f6607..a715e2efec0 100644 --- a/src/common/time/src/time.rs +++ b/src/common/time/src/time.rs @@ -145,6 +145,11 @@ impl Time { self.value = -self.value; self } + + pub fn checked_negative(mut self) -> Option { + self.value = self.value.checked_neg()?; + Some(self) + } } impl From for Time { diff --git a/src/common/time/src/timestamp.rs b/src/common/time/src/timestamp.rs index 120dc8e1e97..8ad38e280fe 100644 --- a/src/common/time/src/timestamp.rs +++ b/src/common/time/src/timestamp.rs @@ -484,6 +484,11 @@ impl Timestamp { self.value = -self.value; self } + + pub fn checked_negative(mut self) -> Option { + self.value = self.value.checked_neg()?; + Some(self) + } } impl Timestamp { diff --git a/src/common/wal/src/config/kafka/datanode.rs b/src/common/wal/src/config/kafka/datanode.rs index 374a7229d4a..4003b52e4bc 100644 --- a/src/common/wal/src/config/kafka/datanode.rs +++ b/src/common/wal/src/config/kafka/datanode.rs @@ -60,7 +60,7 @@ impl Default for DatanodeKafkaConfig { consumer_wait_timeout: Duration::from_millis(100), kafka_topic: KafkaTopicConfig::default(), auto_create_topics: true, - create_index: true, + create_index: false, dump_index_interval: Duration::from_secs(60), topic_latest_offset_fetch_interval: DEFAULT_TOPIC_LATEST_OFFSET_FETCH_INTERVAL, overwrite_entry_start_id: false, diff --git a/src/datanode/src/alive_keeper.rs b/src/datanode/src/alive_keeper.rs index dbf99fdb284..dbfb699036c 100644 --- a/src/datanode/src/alive_keeper.rs +++ b/src/datanode/src/alive_keeper.rs @@ -153,7 +153,7 @@ impl RegionAliveKeeper { async fn close_staled_region(&self, region_id: RegionId) { info!("Closing staled region: {region_id}"); - let request = RegionRequest::Close(RegionCloseRequest {}); + let request = RegionRequest::Close(RegionCloseRequest::default()); if let Err(e) = self.region_server.handle_request(region_id, request).await && e.status_code() != StatusCode::RegionNotFound { diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index 0e2e1da959c..6c176547edd 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -287,6 +287,18 @@ pub enum Error { location: Location, }, + #[snafu(display( + "GC configuration mismatch: metasrv.gc.enable={}, datanode.region_engine.mito.gc.enable={}", + metasrv_gc_enabled, + datanode_gc_enabled, + ))] + GcConfigMismatch { + metasrv_gc_enabled: bool, + datanode_gc_enabled: bool, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Unsupported output type, expected: {}", expected))] UnsupportedOutput { expected: String, @@ -456,6 +468,7 @@ impl ErrorExt for Error { | ColumnNoneDefaultValue { .. } | MissingRequiredField { .. } | RegionEngineNotFound { .. } + | GcConfigMismatch { .. } | ParseAddr { .. } | TomlFormat { .. } | BuildDatanode { .. } => StatusCode::InvalidArguments, diff --git a/src/datanode/src/heartbeat.rs b/src/datanode/src/heartbeat.rs index bfdaf32a1c6..0e0c51624e3 100644 --- a/src/datanode/src/heartbeat.rs +++ b/src/datanode/src/heartbeat.rs @@ -121,6 +121,7 @@ impl HeartbeatTask { pub async fn create_streams( meta_client: &MetaClient, + local_gc_enabled: bool, running: Arc, handler_executor: HeartbeatResponseHandlerExecutorRef, mailbox: MailboxRef, @@ -129,6 +130,13 @@ impl HeartbeatTask { ) -> Result<(HeartbeatSender, HeartbeatConfig)> { let client_id = meta_client.id(); let (tx, mut rx, config) = meta_client.heartbeat().await.context(MetaClientInitSnafu)?; + if config.gc_enabled != local_gc_enabled { + return error::GcConfigMismatchSnafu { + metasrv_gc_enabled: config.gc_enabled, + datanode_gc_enabled: local_gc_enabled, + } + .fail(); + } let mut last_received_lease = Instant::now(); @@ -220,9 +228,17 @@ impl HeartbeatTask { let mailbox = Arc::new(HeartbeatMailbox::new(outgoing_tx)); let quit_signal = Arc::new(Notify::new()); + let local_gc_enabled = self + .region_server + .mito_engine() + .context(RegionEngineNotFoundSnafu { name: "mito" })? + .mito_config() + .gc + .enable; let (mut tx, config) = Self::create_streams( &meta_client, + local_gc_enabled, running.clone(), handler_executor.clone(), mailbox.clone(), @@ -368,6 +384,7 @@ impl HeartbeatTask { error!(e; "Failed to send heartbeat to metasrv"); match Self::create_streams( &meta_client, + local_gc_enabled, running.clone(), handler_executor.clone(), mailbox.clone(), diff --git a/src/datanode/src/heartbeat/handler.rs b/src/datanode/src/heartbeat/handler.rs index 79e0baaef36..71797592c3b 100644 --- a/src/datanode/src/heartbeat/handler.rs +++ b/src/datanode/src/heartbeat/handler.rs @@ -522,7 +522,10 @@ mod tests { .unwrap(); region_server - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); let mut heartbeat_env = HeartbeatResponseTestEnv::new(); diff --git a/src/datanode/src/heartbeat/handler/close_region.rs b/src/datanode/src/heartbeat/handler/close_region.rs index 581c35f3d8f..930c3f11525 100644 --- a/src/datanode/src/heartbeat/handler/close_region.rs +++ b/src/datanode/src/heartbeat/handler/close_region.rs @@ -40,8 +40,10 @@ impl InstructionHandler for CloseRegionsHandler { .collect::>(); let futs = region_ids.iter().map(|region_id| { - ctx.region_server - .handle_request(*region_id, RegionRequest::Close(RegionCloseRequest {})) + ctx.region_server.handle_request( + *region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) }); let results = join_all(futs).await; diff --git a/src/datanode/src/heartbeat/handler/open_region.rs b/src/datanode/src/heartbeat/handler/open_region.rs index 6d57d26e2a6..89ce857ce4d 100644 --- a/src/datanode/src/heartbeat/handler/open_region.rs +++ b/src/datanode/src/heartbeat/handler/open_region.rs @@ -169,11 +169,17 @@ mod tests { .await .unwrap(); region_server - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); region_server - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 25c2c4684a5..5a6ab700854 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -1214,6 +1214,24 @@ impl RegionServerInner { })? .clone(), }, + RegionChange::OfflineCleanup(attribute) => match current_region_status { + Some(status) => match status.clone() { + RegionEngineWithStatus::Registering(_) + | RegionEngineWithStatus::Deregistering(_) + | RegionEngineWithStatus::Ready(_) => { + return error::RegionBusySnafu { region_id }.fail(); + } + }, + None => self + .engines + .read() + .unwrap() + .get(attribute.engine()) + .with_context(|| RegionEngineNotFoundSnafu { + name: attribute.engine(), + })? + .clone(), + }, RegionChange::Deregisters => match current_region_status { Some(status) => match status.clone() { RegionEngineWithStatus::Registering(_) => { @@ -1559,6 +1577,10 @@ impl RegionServerInner { let attribute = parse_region_attribute(&open.engine, &open.options)?; RegionChange::Register(attribute) } + RegionRequest::CleanUp(clean_up) => { + let attribute = parse_region_attribute(&clean_up.engine, &clean_up.options)?; + RegionChange::OfflineCleanup(attribute) + } RegionRequest::Close(_) | RegionRequest::Drop(_) => RegionChange::Deregisters, RegionRequest::Put(_) | RegionRequest::Delete(_) | RegionRequest::BulkInserts(_) => { RegionChange::Ingest @@ -1678,7 +1700,7 @@ impl RegionServerInner { region_change: RegionChange, ) { match region_change { - RegionChange::None | RegionChange::Ingest => {} + RegionChange::None | RegionChange::Ingest | RegionChange::OfflineCleanup(_) => {} RegionChange::Register(_) => { self.region_map.remove(®ion_id); } @@ -1698,7 +1720,7 @@ impl RegionServerInner { ) -> Result<()> { let engine_type = engine.name(); match region_change { - RegionChange::None | RegionChange::Ingest => {} + RegionChange::None | RegionChange::Ingest | RegionChange::OfflineCleanup(_) => {} RegionChange::Register(attribute) => { info!( "Region {region_id} is registered to engine {}", @@ -1852,7 +1874,10 @@ impl RegionServerInner { for (region_id, engine) in regions { let closed = engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await; match closed { Ok(_) => debug!("Region {region_id} is closed"), @@ -1880,6 +1905,7 @@ impl RegionServerInner { enum RegionChange { None, Register(RegionAttribute), + OfflineCleanup(RegionAttribute), Deregisters, Catchup, Ingest, @@ -1944,8 +1970,8 @@ mod tests { use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::region_engine::RegionEngine; use store_api::region_request::{ - PathType, RegionCompactRequest, RegionDeleteRequest, RegionDropRequest, RegionOpenRequest, - RegionPutRequest, RegionTruncateRequest, + PathType, RegionCleanUpRequest, RegionCompactRequest, RegionDeleteRequest, + RegionDropRequest, RegionOpenRequest, RegionPutRequest, RegionTruncateRequest, }; use store_api::storage::RegionId; @@ -2116,6 +2142,77 @@ mod tests { assert!(pinned.as_ref().get_ref().metrics().is_none()); } + #[tokio::test] + async fn test_offline_cleanup_does_not_register_region() { + let mut mock_region_server = mock_region_server(); + let (engine, mut receiver) = MockRegionEngine::new(MITO_ENGINE_NAME); + mock_region_server.register_engine(engine); + + let region_id = RegionId::new(1, 1); + let response = mock_region_server + .handle_request( + region_id, + RegionRequest::CleanUp(RegionCleanUpRequest { + engine: MITO_ENGINE_NAME.to_string(), + table_dir: String::new(), + path_type: PathType::Bare, + options: HashMap::new(), + }), + ) + .await + .unwrap(); + + assert_eq!(response.affected_rows, 0); + let (handled_region_id, handled_request) = receiver.try_recv().unwrap(); + assert_eq!(handled_region_id, region_id); + assert_matches!(handled_request, RegionRequest::CleanUp(_)); + assert!( + mock_region_server + .inner + .region_map + .get(®ion_id) + .is_none() + ); + } + + #[tokio::test] + async fn test_offline_cleanup_rejects_registered_region() { + let mut mock_region_server = mock_region_server(); + let (engine, mut receiver) = MockRegionEngine::new(MITO_ENGINE_NAME); + mock_region_server.register_engine(engine.clone()); + + let region_id = RegionId::new(1, 1); + mock_region_server + .inner + .region_map + .insert(region_id, RegionEngineWithStatus::Ready(engine)); + + let err = mock_region_server + .handle_request( + region_id, + RegionRequest::CleanUp(RegionCleanUpRequest { + engine: MITO_ENGINE_NAME.to_string(), + table_dir: String::new(), + path_type: PathType::Bare, + options: HashMap::new(), + }), + ) + .await + .unwrap_err(); + + assert_eq!(err.status_code(), StatusCode::RegionBusy); + assert!(receiver.try_recv().is_err()); + assert!(matches!( + mock_region_server + .inner + .region_map + .get(®ion_id) + .unwrap() + .clone(), + RegionEngineWithStatus::Ready(_) + )); + } + #[tokio::test] async fn test_region_registering() { common_telemetry::init_default_ut_logging(); @@ -2218,7 +2315,10 @@ mod tests { ); let response = mock_region_server - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); assert_eq!(response.affected_rows, 0); @@ -2524,6 +2624,43 @@ mod tests { assert_matches!(current_engine, CurrentEngine::Engine(_)); }), }, + // RegionChange::OfflineCleanup + CurrentEngineTest { + region_id, + current_region_status: None, + region_change: RegionChange::OfflineCleanup(RegionAttribute::Mito), + assert: Box::new(|result| { + let current_engine = result.unwrap(); + assert_matches!(current_engine, CurrentEngine::Engine(_)); + }), + }, + CurrentEngineTest { + region_id, + current_region_status: Some(RegionEngineWithStatus::Registering(engine.clone())), + region_change: RegionChange::OfflineCleanup(RegionAttribute::Mito), + assert: Box::new(|result| { + let err = result.unwrap_err(); + assert_eq!(err.status_code(), StatusCode::RegionBusy); + }), + }, + CurrentEngineTest { + region_id, + current_region_status: Some(RegionEngineWithStatus::Deregistering(engine.clone())), + region_change: RegionChange::OfflineCleanup(RegionAttribute::Mito), + assert: Box::new(|result| { + let err = result.unwrap_err(); + assert_eq!(err.status_code(), StatusCode::RegionBusy); + }), + }, + CurrentEngineTest { + region_id, + current_region_status: Some(RegionEngineWithStatus::Ready(engine.clone())), + region_change: RegionChange::OfflineCleanup(RegionAttribute::Mito), + assert: Box::new(|result| { + let err = result.unwrap_err(); + assert_eq!(err.status_code(), StatusCode::RegionBusy); + }), + }, ]; for test in tests { diff --git a/src/datatypes/src/json/value.rs b/src/datatypes/src/json/value.rs index 40217f75141..b727bd8b5ab 100644 --- a/src/datatypes/src/json/value.rs +++ b/src/datatypes/src/json/value.rs @@ -317,13 +317,6 @@ impl JsonValue { } } - pub(crate) fn new_with(json_variant: JsonVariant, json_type: Arc) -> Self { - Self { - json_type: OnceLock::from(json_type), - json_variant, - } - } - pub(crate) fn data_type(&self) -> ConcreteDataType { ConcreteDataType::json2(self.json_type().clone()) } @@ -383,10 +376,6 @@ impl JsonValue { self.json_variant } - pub(crate) fn variant(&self) -> &JsonVariant { - &self.json_variant - } - pub(crate) fn into_value(self) -> Value { fn helper(v: JsonVariant) -> Value { match v { @@ -828,6 +817,15 @@ impl<'a> JsonValueRef<'a> { ().into() } + /// Creates a JSON value reference with its precomputed native type. + /// The native type must describe `json_variant` exactly. + pub fn new_with_type(json_variant: JsonVariantRef<'a>, json_type: JsonNativeType) -> Self { + Self { + json_type: OnceLock::from(Arc::new(json_type)), + json_variant, + } + } + pub(crate) fn data_type(&self) -> ConcreteDataType { ConcreteDataType::json2(self.json_type().as_ref().clone()) } diff --git a/src/datatypes/src/value.rs b/src/datatypes/src/value.rs index 5d703480d3f..9d50ac78af4 100644 --- a/src/datatypes/src/value.rs +++ b/src/datatypes/src/value.rs @@ -543,20 +543,20 @@ impl Value { None } } - Value::Int8(x) => Some(Value::Int8(-*x)), - Value::Int16(x) => Some(Value::Int16(-*x)), - Value::Int32(x) => Some(Value::Int32(-*x)), - Value::Int64(x) => Some(Value::Int64(-*x)), + Value::Int8(x) => x.checked_neg().map(Value::Int8), + Value::Int16(x) => x.checked_neg().map(Value::Int16), + Value::Int32(x) => x.checked_neg().map(Value::Int32), + Value::Int64(x) => x.checked_neg().map(Value::Int64), Value::Float32(x) => Some(Value::Float32(-*x)), Value::Float64(x) => Some(Value::Float64(-*x)), Value::Decimal128(x) => Some(Value::Decimal128(x.negative())), - Value::Date(x) => Some(Value::Date(x.negative())), - Value::Timestamp(x) => Some(Value::Timestamp(x.negative())), - Value::Time(x) => Some(Value::Time(x.negative())), - Value::Duration(x) => Some(Value::Duration(x.negative())), - Value::IntervalYearMonth(x) => Some(Value::IntervalYearMonth(x.negative())), - Value::IntervalDayTime(x) => Some(Value::IntervalDayTime(x.negative())), - Value::IntervalMonthDayNano(x) => Some(Value::IntervalMonthDayNano(x.negative())), + Value::Date(x) => x.checked_negative().map(Value::Date), + Value::Timestamp(x) => x.checked_negative().map(Value::Timestamp), + Value::Time(x) => x.checked_negative().map(Value::Time), + Value::Duration(x) => x.checked_negative().map(Value::Duration), + Value::IntervalYearMonth(x) => x.checked_negative().map(Value::IntervalYearMonth), + Value::IntervalDayTime(x) => x.checked_negative().map(Value::IntervalDayTime), + Value::IntervalMonthDayNano(x) => x.checked_negative().map(Value::IntervalMonthDayNano), Value::Binary(_) | Value::String(_) @@ -1744,6 +1744,49 @@ pub(crate) mod tests { use crate::types::json_type::{JsonNativeType, JsonObjectType}; use crate::vectors::ListVectorBuilder; + #[test] + fn test_try_negative_overflow() { + // Negating a MIN value overflows, so try_negative returns None instead + // of panicking (consistent with the unsigned arms). + assert_eq!(Value::Int8(i8::MIN).try_negative(), None); + assert_eq!(Value::Int16(i16::MIN).try_negative(), None); + assert_eq!(Value::Int32(i32::MIN).try_negative(), None); + assert_eq!(Value::Int64(i64::MIN).try_negative(), None); + assert_eq!( + Value::Timestamp(Timestamp::new_nanosecond(i64::MIN)).try_negative(), + None + ); + assert_eq!(Value::Date(Date::new(i32::MIN)).try_negative(), None); + assert_eq!( + Value::Time(Time::new_nanosecond(i64::MIN)).try_negative(), + None + ); + assert_eq!( + Value::Duration(Duration::new_nanosecond(i64::MIN)).try_negative(), + None + ); + assert_eq!( + Value::IntervalYearMonth(IntervalYearMonth::new(i32::MIN)).try_negative(), + None + ); + assert_eq!( + Value::IntervalDayTime(IntervalDayTime::new(i32::MIN, i32::MIN)).try_negative(), + None + ); + assert_eq!( + Value::IntervalMonthDayNano(IntervalMonthDayNano::new(i32::MIN, i32::MIN, i64::MIN)) + .try_negative(), + None + ); + + // Non-MIN values still negate. + assert_eq!(Value::Int64(5).try_negative(), Some(Value::Int64(-5))); + assert_eq!( + Value::Timestamp(Timestamp::new_nanosecond(5)).try_negative(), + Some(Value::Timestamp(Timestamp::new_nanosecond(-5))) + ); + } + pub(crate) fn build_struct_type() -> StructType { StructType::new(Arc::new(vec![ StructField::new("id".to_string(), ConcreteDataType::int32_datatype(), false), diff --git a/src/datatypes/src/vectors/json/builder.rs b/src/datatypes/src/vectors/json/builder.rs index 41aff8c29c6..ebf11ab7961 100644 --- a/src/datatypes/src/vectors/json/builder.rs +++ b/src/datatypes/src/vectors/json/builder.rs @@ -19,17 +19,17 @@ use arrow_schema::DataType; use crate::data_type::ConcreteDataType; use crate::error::{Result, TryFromValueSnafu, UnexpectedSnafu, UnsupportedOperationSnafu}; -use crate::json::value::{JsonNumber, JsonValue, JsonVariant}; +use crate::json::value::{JsonNumber, JsonVariant, encode_json_variant}; use crate::prelude::{ValueRef, Vector, VectorRef}; use crate::types::StructType; use crate::types::json_type::{JsonNativeType, is_include}; -use crate::value::{ListValueRef, StructValueRef}; +use crate::value::{ListValue, StructValue, StructValueRef, Value}; use crate::vectors::{MutableVector, StructVectorBuilder}; #[derive(Clone)] pub(crate) struct JsonVectorBuilder { merged_type: JsonNativeType, - values: Vec, + values: Vec, } impl JsonVectorBuilder { @@ -56,23 +56,22 @@ impl JsonVectorBuilder { let mut builder = StructVectorBuilder::with_type_and_capacity(struct_type.clone(), self.values.len()); - for value in self.values.iter_mut() { - if value.is_null() { + for value in std::mem::take(&mut self.values) { + if matches!(&value, JsonVariant::Null) { builder.push_null(); continue; } - value.try_align(&self.merged_type)?; - let value_ref = json_variant_to_struct_value_ref(value.variant(), struct_type.clone())?; - builder.push_struct_value_ref(value_ref)?; + let value = json_variant_into_struct_value(value, struct_type.clone())?; + builder.push_struct_value_ref(StructValueRef::Ref(&value))?; } Ok(builder.to_vector()) } } -fn json_variant_to_struct_value_ref( - value: &JsonVariant, +fn json_variant_into_struct_value( + value: JsonVariant, struct_type: StructType, -) -> Result> { +) -> Result { let JsonVariant::Object(object) = value else { return TryFromValueSnafu { reason: format!("expected json object value, got {value:?}"), @@ -80,63 +79,87 @@ fn json_variant_to_struct_value_ref( .fail(); }; - let values = struct_type - .fields() - .iter() - .map(|field| { - object - .get(field.name()) - .map(|v| json_variant_to_value_ref(v, field.data_type())) - .unwrap_or(Ok(ValueRef::Null)) - }) - .collect::>>()?; + let mut entries = object.into_iter(); + let mut entry = entries.next(); + let mut values = Vec::with_capacity(struct_type.fields().len()); + for field in struct_type.fields().iter() { + let value = match entry.take() { + Some((name, value)) if name == field.name() => { + entry = entries.next(); + json_variant_into_value(value, field.data_type())? + } + Some((name, _)) if name.as_str() < field.name() => { + return TryFromValueSnafu { + reason: format!("field {name} is missing from merged JSON type"), + } + .fail(); + } + next => { + entry = next; + Value::Null + } + }; + values.push(value); + } + if let Some((name, _)) = entry { + return TryFromValueSnafu { + reason: format!("field {name} is missing from merged JSON type"), + } + .fail(); + } - Ok(StructValueRef::RefList { - val: values, - fields: struct_type, - }) + Ok(StructValue::new(values, struct_type)) } -fn json_variant_to_value_ref<'a>( - value: &'a JsonVariant, - expected_type: &ConcreteDataType, -) -> Result> { - let value = match value { - JsonVariant::Null => ValueRef::Null, - JsonVariant::Bool(x) => ValueRef::Boolean(*x), - JsonVariant::Number(x) => match x { - JsonNumber::PosInt(i) => ValueRef::UInt64(*i), - JsonNumber::NegInt(i) => ValueRef::Int64(*i), - JsonNumber::Float(f) => ValueRef::Float64(*f), - }, - JsonVariant::String(x) => ValueRef::String(x), - JsonVariant::Array(array) => { - let item_type = match expected_type { - ConcreteDataType::List(list_type) => list_type.item_type().clone(), - _ => ConcreteDataType::null_datatype(), - }; - let values = array - .iter() - .map(|v| json_variant_to_value_ref(v, &item_type)) - .collect::>>()?; - ValueRef::List(ListValueRef::RefList { - val: values, - item_datatype: Arc::new(item_type), - }) +fn json_variant_into_value(value: JsonVariant, expected_type: &ConcreteDataType) -> Result { + let value = match (value, expected_type) { + (JsonVariant::Null, _) | (_, ConcreteDataType::Null(_)) => Value::Null, + (JsonVariant::Bool(x), ConcreteDataType::Boolean(_)) => Value::Boolean(x), + (JsonVariant::Number(JsonNumber::PosInt(x)), ConcreteDataType::UInt64(_)) => { + Value::UInt64(x) } - JsonVariant::Object(_) => { - let ConcreteDataType::Struct(struct_type) = expected_type else { + (JsonVariant::Number(x), ConcreteDataType::Int64(_)) => { + let x = match x { + JsonNumber::PosInt(x) => i64::try_from(x).ok(), + JsonNumber::NegInt(x) => Some(x), + JsonNumber::Float(_) => None, + }; + let Some(x) = x else { return TryFromValueSnafu { - reason: format!("expected struct type, got {expected_type}"), + reason: format!("unable to convert {x:?} to Int64"), } .fail(); }; - ValueRef::Struct(json_variant_to_struct_value_ref( - value, - struct_type.clone(), - )?) + Value::Int64(x) + } + (JsonVariant::Number(JsonNumber::PosInt(x)), ConcreteDataType::Float64(_)) => { + Value::Float64((x as f64).into()) + } + (JsonVariant::Number(JsonNumber::NegInt(x)), ConcreteDataType::Float64(_)) => { + Value::Float64((x as f64).into()) + } + (JsonVariant::Number(JsonNumber::Float(x)), ConcreteDataType::Float64(_)) => { + Value::Float64(x) + } + (JsonVariant::String(x), ConcreteDataType::String(_)) => Value::String(x.into()), + (JsonVariant::Array(array), ConcreteDataType::List(list_type)) => { + let item_type = list_type.item_type().clone(); + let values = array + .into_iter() + .map(|v| json_variant_into_value(v, &item_type)) + .collect::>>()?; + Value::List(ListValue::new(values, Arc::new(item_type))) + } + (value @ JsonVariant::Object(_), ConcreteDataType::Struct(struct_type)) => { + Value::Struct(json_variant_into_struct_value(value, struct_type.clone())?) + } + (value, ConcreteDataType::Binary(_)) => Value::from(encode_json_variant(value)?), + (value, expected_type) => { + return TryFromValueSnafu { + reason: format!("unable to convert json value {value:?} to {expected_type}"), + } + .fail(); } - JsonVariant::Variant(x) => ValueRef::Binary(x), }; Ok(value) } @@ -185,13 +208,12 @@ impl MutableVector for JsonVectorBuilder { self.merged_type.merge(json_type); } - let value = JsonValue::new_with(JsonVariant::from(value.variant()), value.json_type()); - self.values.push(value); + self.values.push(JsonVariant::from(value.variant())); Ok(()) } fn push_null(&mut self) { - self.values.push(JsonValue::null()) + self.values.push(JsonVariant::Null) } fn extend_slice_of(&mut self, _: &dyn Vector, _: usize, _: usize) -> Result<()> { @@ -227,7 +249,7 @@ mod tests { } // Object inputs should merge into a superset schema, preserve null rows, - // and align conflicting nested values into Variant payloads. + // and project conflicting nested values into Variant payloads. let mut builder = JsonVectorBuilder::new(JsonNativeType::Object(Default::default()), 3); let first = parse_json_value(r#"{"id":1,"payload":{"name":"foo"}}"#); let second = parse_json_value(r#"{"id":2,"extra":true,"payload":"raw"}"#); @@ -332,7 +354,7 @@ mod tests { } #[test] - fn test_json_variant_to_struct_value_ref() -> Result<()> { + fn test_json_variant_into_struct_value() -> Result<()> { let item_type = ConcreteDataType::struct_datatype(StructType::new(Arc::new(vec![StructField::new( "id".to_string(), @@ -370,8 +392,10 @@ mod tests { JsonVariant::from([("name", JsonVariant::from("foo"))]), ), ]); - let value_ref = json_variant_to_struct_value_ref(&variant, struct_type.clone())?; - let value = value_ref.to_value(); + let value = Value::Struct(json_variant_into_struct_value( + variant, + struct_type.clone(), + )?); assert_eq!( value, diff --git a/src/flow/src/adapter.rs b/src/flow/src/adapter.rs index 1e663582a28..e8dd10ad7d8 100644 --- a/src/flow/src/adapter.rs +++ b/src/flow/src/adapter.rs @@ -76,6 +76,23 @@ pub(crate) mod table_source; use crate::FrontendInvoker; use crate::error::Error; +fn expire_after_secs_to_millis(expire_after_secs: i64) -> Result { + ensure!( + expire_after_secs >= 0, + InvalidQuerySnafu { + reason: format!("EXPIRE AFTER must be non-negative, got {expire_after_secs} seconds"), + } + ); + + expire_after_secs + .checked_mul(1_000) + .with_context(|| InvalidQuerySnafu { + reason: format!( + "EXPIRE AFTER value {expire_after_secs} seconds cannot be represented in milliseconds" + ), + }) +} + // `GREPTIME_TIMESTAMP` is not used to distinguish when table is created automatically by flow pub const AUTO_CREATED_PLACEHOLDER_TS_COL: &str = "__ts_placeholder"; @@ -740,7 +757,7 @@ impl StreamingEngine { source_table_ids, create_if_not_exists, or_replace, - expire_after, + expire_after: expire_after_secs, eval_interval: _, comment, sql, @@ -748,6 +765,9 @@ impl StreamingEngine { query_ctx, .. } = args; + let expire_after = expire_after_secs + .map(expire_after_secs_to_millis) + .transpose()?; let mut node_ctx = self.node_context.write().await; // assign global id to source and sink table diff --git a/src/flow/src/adapter/refill.rs b/src/flow/src/adapter/refill.rs index 6d66505e89e..3672c0427e1 100644 --- a/src/flow/src/adapter/refill.rs +++ b/src/flow/src/adapter/refill.rs @@ -85,7 +85,10 @@ impl StreamingEngine { } } - let expire_after = info.expire_after(); + let expire_after = info + .expire_after() + .map(super::expire_after_secs_to_millis) + .transpose()?; // TODO(discord9): better way to get last point let now = self.tick_manager.tick(); let plan = self diff --git a/src/flow/src/adapter/tests.rs b/src/flow/src/adapter/tests.rs index 9ed7e676c29..df6a58c46ce 100644 --- a/src/flow/src/adapter/tests.rs +++ b/src/flow/src/adapter/tests.rs @@ -61,3 +61,25 @@ pub fn new_test_table_info_with_name>( /// /// containing several default table info and schema fn mock_harness_flow_node_manager() {} + +#[test] +fn test_expire_after_secs_to_millis() { + assert_eq!(expire_after_secs_to_millis(300).unwrap(), 300_000); + assert_eq!(expire_after_secs_to_millis(0).unwrap(), 0); + + let max_secs = i64::MAX / 1_000; + assert_eq!( + expire_after_secs_to_millis(max_secs).unwrap(), + max_secs * 1_000 + ); +} + +#[test] +fn test_expire_after_secs_to_millis_invalid() { + let invalid_values = [i64::MAX / 1_000 + 1, -1]; + + for invalid_secs in invalid_values { + let error = expire_after_secs_to_millis(invalid_secs).unwrap_err(); + assert!(matches!(error, Error::InvalidQuery { .. })); + } +} diff --git a/src/flow/src/engine.rs b/src/flow/src/engine.rs index 386610ad716..3a139a898a9 100644 --- a/src/flow/src/engine.rs +++ b/src/flow/src/engine.rs @@ -34,6 +34,7 @@ pub struct CreateFlowArgs { pub source_table_ids: Vec, pub create_if_not_exists: bool, pub or_replace: bool, + /// Duration in seconds. pub expire_after: Option, pub eval_interval: Option, pub comment: Option, diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 918185cb8f6..c7f7e71d34c 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -237,6 +237,7 @@ mod tests { is_handshake.then_some(api::v1::meta::HeartbeatConfig { heartbeat_interval_ms, retry_interval_ms: heartbeat_interval_ms, + gc_enabled: false, }); is_handshake = false; let response = HeartbeatResponse { diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 7242d601c1b..5beb8eb827a 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -611,7 +611,11 @@ impl Instance { let checker_ref = self.plugins.get::(); checker_ref .as_ref() - .check_permission(query_ctx.current_user(), PermissionReq::SqlStatement(&stmt)) + .check_permission_with_context( + query_ctx.current_user(), + PermissionReq::SqlStatement(&stmt), + Some(&query_ctx.current_schema()), + ) .context(PermissionSnafu)?; check_permission(self.plugins.clone(), &stmt, &query_ctx)?; let catalog_name = query_ctx.current_catalog().to_string(); @@ -672,9 +676,10 @@ impl Instance { let mut results = Vec::with_capacity(stmts.len()); for stmt in stmts { if let Err(e) = checker - .check_permission( + .check_permission_with_context( query_ctx.current_user(), PermissionReq::SqlStatement(&stmt), + Some(&query_ctx.current_schema()), ) .context(PermissionSnafu) { @@ -848,7 +853,11 @@ impl Instance { self.plugins .get::() .as_ref() - .check_permission(query_ctx.current_user(), PermissionReq::SqlStatement(&stmt)) + .check_permission_with_context( + query_ctx.current_user(), + PermissionReq::SqlStatement(&stmt), + Some(&query_ctx.current_schema()), + ) .context(PermissionSnafu)?; let plan = self diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index 0ca1a2cf209..808a2bd6ad8 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -71,7 +71,11 @@ impl GrpcQueryHandler for Instance { self.plugins .get::() .as_ref() - .check_permission(ctx.current_user(), PermissionReq::GrpcRequest(&request)) + .check_permission_with_context( + ctx.current_user(), + PermissionReq::GrpcRequest(&request), + Some(&ctx.current_schema()), + ) .context(PermissionSnafu)?; let output = match request { @@ -356,7 +360,14 @@ impl Instance { plugins .get::() .as_ref() - .check_permission(ctx.current_user(), PermissionReq::BulkInsert) + .check_permission( + ctx.current_user(), + PermissionReq::BulkInsert { + catalog: &table_name.catalog_name, + schema: &table_name.schema_name, + table: &table_name.table_name, + }, + ) .context(PermissionSnafu)?; // Resolve table reference diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index 16e3b996064..c970ca144db 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -12,20 +12,18 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod trace_ingest; +pub mod trace_semconv; +pub mod trace_types; + use std::sync::Arc; -use api::helper::ColumnDataTypeWrapper; -use api::v1::alter_table_expr::Kind; -use api::v1::{ - AlterTableExpr, ColumnDataType, ModifyColumnType, ModifyColumnTypes, RowInsertRequests, -}; use async_trait::async_trait; use auth::{PermissionChecker, PermissionCheckerRef, PermissionReq}; use client::Output; -use common_error::ext::{BoxedError, ErrorExt}; -use common_error::status_code::StatusCode; +use common_error::ext::BoxedError; use common_query::prelude::GREPTIME_PHYSICAL_TABLE; -use common_telemetry::{tracing, warn}; +use common_telemetry::tracing; use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; @@ -34,66 +32,20 @@ use servers::error::{self, AuthSnafu, Result as ServerResult}; use servers::http::prom_store::PHYSICAL_TABLE_PARAM; use servers::interceptor::{OpenTelemetryProtocolInterceptor, OpenTelemetryProtocolInterceptorRef}; use servers::otlp; -use servers::otlp::coerce::{coerce_value_data, trace_value_datatype}; -use servers::otlp::trace::TraceAuxData; -use servers::otlp::trace::span::{TraceSpan, TraceSpanGroup}; use servers::query_handler::{ OpenTelemetryProtocolHandler, PipelineHandlerRef, TraceIngestOutcome, }; use session::context::QueryContextRef; -use snafu::{IntoError, ResultExt}; +use snafu::ResultExt; use table::requests::{ OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM, SEMANTIC_PER_TABLE_INDEX_KEY, - SEMANTIC_PIPELINE, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SEMANTIC_TRACE_CONVENTIONS, - SEMANTIC_VALUE_MIXED, SEMANTIC_VALUE_UNKNOWN, SIGNAL_TYPE_LOG, SIGNAL_TYPE_METRIC, - SIGNAL_TYPE_TRACE, SOURCE_OPENTELEMETRY, TABLE_DATA_MODEL_TRACE_V1, + SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_LOG, SIGNAL_TYPE_METRIC, + SOURCE_OPENTELEMETRY, }; +use self::trace_ingest::trace_conventions; use crate::instance::Instance; -use crate::instance::otlp::trace_semconv::trace_semconv_fixed_type; -use crate::instance::otlp::trace_types::{ - PendingTraceColumnRewrite, choose_trace_reconcile_decision, enrich_trace_reconcile_error, - is_trace_reconcile_candidate_type, push_observed_trace_type, validate_trace_column_rewrites, -}; -use crate::metrics::{ - OTLP_LOGS_ROWS, OTLP_METRICS_ROWS, OTLP_TRACES_FAILURE_COUNT, OTLP_TRACES_ROWS, -}; - -pub mod trace_semconv; -pub mod trace_types; - -const TRACE_FAILURE_MESSAGE_LIMIT: usize = 4; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum ChunkFailureReaction { - RetryPerSpan, - DiscardChunk, - Propagate, -} - -impl ChunkFailureReaction { - fn as_metric_label(self) -> &'static str { - match self { - Self::RetryPerSpan => "retry_per_span", - Self::DiscardChunk => "discard_chunk", - Self::Propagate => "propagate_failure", - } - } -} - -struct TraceChunkIngestContext<'a> { - pipeline_handler: PipelineHandlerRef, - pipeline: &'a PipelineWay, - pipeline_params: &'a GreptimePipelineParams, - table_name: &'a str, - is_trace_v1_model: bool, -} - -struct TraceIngestState { - aux_data: TraceAuxData, - outcome: TraceIngestOutcome, - failure_messages: Vec, -} +use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS}; #[async_trait] impl OpenTelemetryProtocolHandler for Instance { @@ -265,868 +217,3 @@ impl OpenTelemetryProtocolHandler for Instance { Ok(outputs) } } - -impl Instance { - /// Ingest OTLP trace spans with chunk-level writes and span-level fallback on - /// deterministic chunk failures. - #[allow(clippy::too_many_arguments)] - async fn ingest_trace_spans( - &self, - pipeline_handler: PipelineHandlerRef, - pipeline: &PipelineWay, - pipeline_params: &GreptimePipelineParams, - table_name: String, - groups: Vec, - conventions: &str, - ctx: QueryContextRef, - ) -> ServerResult { - let is_trace_v1_model = matches!(pipeline, PipelineWay::OtlpTraceDirectV1); - - // Only the main span table gets the identity; the derived `_services` / - // `_operations` lookup tables keep the unstamped `ctx`. - let main_ctx = { - let mut c = (*ctx).clone(); - c.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_TRACE); - c.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY); - if is_trace_v1_model { - c.set_extension(SEMANTIC_PIPELINE, TABLE_DATA_MODEL_TRACE_V1); - c.set_extension(SEMANTIC_TRACE_CONVENTIONS, conventions); - } - Arc::new(c) - }; - - let ingest_ctx = TraceChunkIngestContext { - pipeline_handler, - pipeline, - pipeline_params, - table_name: &table_name, - is_trace_v1_model, - }; - let mut ingest_state = TraceIngestState { - aux_data: TraceAuxData::default(), - outcome: TraceIngestOutcome::default(), - failure_messages: Vec::new(), - }; - - for group in groups { - let chunks = chunk_owned(group.spans, self.trace_ingest_chunk_size); - for chunk in chunks { - self.ingest_trace_chunk(&ingest_ctx, chunk, main_ctx.clone(), &mut ingest_state) - .await?; - } - } - - OTLP_TRACES_ROWS.inc_by(ingest_state.outcome.accepted_spans as u64); - - if !ingest_state.aux_data.is_empty() { - // Auxiliary trace tables are derived from spans whose main-table - // writes are already confirmed, so they never create new accepted - // spans and they do not affect rejected span counts. - let (aux_requests, _) = otlp::trace::to_grpc_insert_requests_for_aux_tables( - std::mem::take(&mut ingest_state.aux_data), - ingest_ctx.pipeline, - ingest_ctx.table_name, - )?; - - if !aux_requests.inserts.is_empty() { - match self - .insert_trace_requests(aux_requests, ingest_ctx.is_trace_v1_model, ctx) - .await - { - Ok(output) => { - Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); - } - Err(err) => { - Self::push_trace_failure_message( - &mut ingest_state.failure_messages, - "aux_table_update_failed", - format!( - "Auxiliary trace tables were not fully updated ({})", - err.status_code().as_ref() - ), - ); - } - } - } - } - - ingest_state.outcome.error_message = Self::finish_trace_failure_message( - ingest_state.outcome.accepted_spans, - ingest_state.outcome.rejected_spans, - ingest_state.failure_messages, - ); - - Ok(ingest_state.outcome) - } - - /// Ingest one owned trace chunk so successful spans can be moved into the - /// accepted set without extra cloning. - async fn ingest_trace_chunk( - &self, - ingest_ctx: &TraceChunkIngestContext<'_>, - chunk: Vec, - ctx: QueryContextRef, - ingest_state: &mut TraceIngestState, - ) -> ServerResult<()> { - // Try the fast path first so healthy batches keep their original - // throughput and write amplification stays low. - let (requests, chunk_rows) = otlp::trace::to_grpc_insert_requests_from_spans( - &chunk, - ingest_ctx.pipeline, - ingest_ctx.pipeline_params, - ingest_ctx.table_name, - &ctx, - ingest_ctx.pipeline_handler.clone(), - )?; - - match self - .insert_trace_requests(requests, ingest_ctx.is_trace_v1_model, ctx.clone()) - .await - { - Ok(output) => { - Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); - ingest_state.outcome.accepted_spans += chunk_rows; - for span in &chunk { - ingest_state.aux_data.observe_span(span); - } - } - Err(err) => match Self::classify_trace_chunk_failure(err.status_code()) { - ChunkFailureReaction::RetryPerSpan => { - Self::push_trace_failure_message( - &mut ingest_state.failure_messages, - ChunkFailureReaction::RetryPerSpan.as_metric_label(), - format!("Chunk fallback triggered by {}", err.status_code().as_ref()), - ); - // Only deterministic failures are retried span by span. - // This includes schemaless table or column creation paths for - // trace ingestion. Ambiguous failures are handled below - // without retrying because the chunk may already have been - // ingested. - self.ingest_trace_chunk_span_by_span( - ingest_ctx, - chunk, - ctx.clone(), - ingest_state, - ) - .await?; - } - ChunkFailureReaction::DiscardChunk => { - ingest_state.outcome.rejected_spans += chunk.len(); - Self::push_trace_failure_message( - &mut ingest_state.failure_messages, - ChunkFailureReaction::DiscardChunk.as_metric_label(), - format!( - "Discarded {} spans after ambiguous chunk failure ({})", - chunk.len(), - err.status_code().as_ref() - ), - ); - // TODO(shuiyisong): Add an idempotent retry-safe recovery path for - // ambiguous chunk failures such as timeout-like errors. - } - // Retryable or ambiguous failures must fail the request instead of - // becoming partial success. This path is not retry-safe because the - // chunk may already have been committed before the error surfaced. - ChunkFailureReaction::Propagate => { - Self::push_trace_failure_message( - &mut ingest_state.failure_messages, - ChunkFailureReaction::Propagate.as_metric_label(), - format!( - "Propagating retryable chunk failure ({})", - err.status_code().as_ref() - ), - ); - return Err(err); - } - }, - } - - Ok(()) - } - - /// Retry spans one by one only after a deterministic chunk failure. - async fn ingest_trace_chunk_span_by_span( - &self, - ingest_ctx: &TraceChunkIngestContext<'_>, - chunk: Vec, - ctx: QueryContextRef, - ingest_state: &mut TraceIngestState, - ) -> ServerResult<()> { - for span in chunk { - let (requests, rows) = otlp::trace::to_grpc_insert_requests_from_spans( - std::slice::from_ref(&span), - ingest_ctx.pipeline, - ingest_ctx.pipeline_params, - ingest_ctx.table_name, - &ctx, - ingest_ctx.pipeline_handler.clone(), - )?; - - match self - .insert_trace_requests(requests, ingest_ctx.is_trace_v1_model, ctx.clone()) - .await - { - Ok(output) => { - Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); - ingest_state.outcome.accepted_spans += rows; - ingest_state.aux_data.observe_span(&span); - } - Err(err) => { - if Self::should_propagate_trace_span_failure(err.status_code()) { - Self::push_trace_failure_message( - &mut ingest_state.failure_messages, - ChunkFailureReaction::Propagate.as_metric_label(), - format!( - "Propagating retryable span failure for {}:{} ({})", - span.trace_id, - span.span_id, - err.status_code().as_ref() - ), - ); - return Err(err); - } - - ingest_state.outcome.rejected_spans += 1; - Self::push_trace_failure_message( - &mut ingest_state.failure_messages, - "span_rejected", - format!( - "Rejected span {}:{} ({})", - span.trace_id, - span.span_id, - err.status_code().as_ref() - ), - ); - } - } - } - - Ok(()) - } - - /// Reconcile and insert one trace request batch. - async fn insert_trace_requests( - &self, - mut requests: RowInsertRequests, - is_trace_v1_model: bool, - ctx: QueryContextRef, - ) -> ServerResult { - if is_trace_v1_model { - self.reconcile_trace_column_types(&mut requests, &ctx) - .await?; - self.handle_trace_inserts(requests, ctx) - .await - .map_err(BoxedError::new) - .context(error::ExecuteGrpcQuerySnafu) - } else { - self.handle_log_inserts(requests, ctx) - .await - .map_err(BoxedError::new) - .context(error::ExecuteGrpcQuerySnafu) - } - } - - fn classify_trace_chunk_failure(status: StatusCode) -> ChunkFailureReaction { - match status { - StatusCode::InvalidArguments - | StatusCode::InvalidSyntax - | StatusCode::Unsupported - | StatusCode::TableNotFound - | StatusCode::TableColumnNotFound => ChunkFailureReaction::RetryPerSpan, - StatusCode::DatabaseNotFound => ChunkFailureReaction::DiscardChunk, - StatusCode::Cancelled | StatusCode::DeadlineExceeded => ChunkFailureReaction::Propagate, - StatusCode::StorageUnavailable - | StatusCode::RuntimeResourcesExhausted - | StatusCode::Internal - | StatusCode::RegionNotReady - | StatusCode::TableUnavailable - | StatusCode::RegionBusy => ChunkFailureReaction::Propagate, - _ => ChunkFailureReaction::DiscardChunk, - } - } - - fn should_propagate_trace_span_failure(status: StatusCode) -> bool { - matches!( - Self::classify_trace_chunk_failure(status), - ChunkFailureReaction::Propagate - ) - } - - fn add_trace_write_cost(outcome: &mut TraceIngestOutcome, cost: usize) { - outcome.write_cost += cost; - } - - fn push_trace_failure_message(messages: &mut Vec, label: &str, message: String) { - OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).inc(); - - if messages.len() < TRACE_FAILURE_MESSAGE_LIMIT { - messages.push(message); - } else if messages.len() == TRACE_FAILURE_MESSAGE_LIMIT { - tracing::debug!( - label, - limit = TRACE_FAILURE_MESSAGE_LIMIT, - "Trace ingest failure message limit reached; suppressing additional failure details" - ); - } - } - - fn finish_trace_failure_message( - accepted_spans: usize, - rejected_spans: usize, - messages: Vec, - ) -> Option { - if rejected_spans == 0 && messages.is_empty() { - return None; - } - - let mut summary = format!( - "Accepted {} spans, rejected {} spans", - accepted_spans, rejected_spans - ); - - if !messages.is_empty() { - summary.push_str(": "); - summary.push_str(&messages.join("; ")); - } - - Some(summary) - } - - /// Widen existing trace table columns to Float64 before request rewrite. - async fn alter_trace_table_columns_to_float64( - &self, - ctx: &QueryContextRef, - table_name: &str, - column_names: &[String], - ) -> ServerResult<()> { - let catalog_name = ctx.current_catalog().to_string(); - let schema_name = ctx.current_schema(); - let alter_expr = AlterTableExpr { - catalog_name: catalog_name.clone(), - schema_name: schema_name.clone(), - table_name: table_name.to_string(), - kind: Some(Kind::ModifyColumnTypes(ModifyColumnTypes { - modify_column_types: column_names - .iter() - .map(|column_name| ModifyColumnType { - column_name: column_name.clone(), - target_type: ColumnDataType::Float64 as i32, - target_type_extension: None, - }) - .collect(), - })), - }; - - if let Err(err) = self - .statement_executor - .alter_table_inner(alter_expr, ctx.clone()) - .await - { - let table = self - .catalog_manager - .table(&catalog_name, &schema_name, table_name, None) - .await - .map_err(servers::error::Error::from)?; - let alter_already_applied = table - .map(|table| { - let table_schema = table.schema(); - column_names.iter().all(|column_name| { - table_schema - .column_schema_by_name(column_name) - .and_then(|table_col| { - ColumnDataTypeWrapper::try_from(table_col.data_type.clone()) - .ok() - .map(|wrapper| wrapper.datatype()) - }) - == Some(ColumnDataType::Float64) - }) - }) - .unwrap_or(false); - - if alter_already_applied { - return Ok(()); - } - - warn!( - table_name, - columns = ?column_names, - error = %err, - "failed to widen trace columns before insert" - ); - - return Err(wrap_trace_alter_failure(err)); - } - - Ok(()) - } - - /// Coerce request column types and values to match the existing table schema - /// for compatible type pairs. Existing table schema wins when present; - /// otherwise the full request batch decides a stable target type. - async fn reconcile_trace_column_types( - &self, - requests: &mut RowInsertRequests, - ctx: &QueryContextRef, - ) -> ServerResult<()> { - let catalog = ctx.current_catalog(); - let schema = ctx.current_schema(); - - for req in &mut requests.inserts { - let table = self - .catalog_manager - .table(catalog, &schema, &req.table_name, None) - .await?; - - let Some(rows) = req.rows.as_mut() else { - continue; - }; - - let table_schema = table.map(|table| table.schema()); - let mut pending_rewrites = Vec::new(); - let mut pending_alter_columns = Vec::new(); - - for (col_idx, col_schema) in rows.schema.iter().enumerate() { - let Some(current_type) = ColumnDataType::try_from(col_schema.datatype).ok() else { - continue; - }; - - let mut observed_types = Vec::new(); - push_observed_trace_type(&mut observed_types, current_type); - - // Scan the full request first so the final type decision is not affected - // by row order inside the batch. - for row in &rows.rows { - let Some(value) = row - .values - .get(col_idx) - .and_then(|value| value.value_data.as_ref()) - else { - continue; - }; - - let Some(value_type) = trace_value_datatype(value) else { - continue; - }; - push_observed_trace_type(&mut observed_types, value_type); - } - - let existing_type = table_schema - .as_ref() - .and_then(|schema| schema.column_schema_by_name(&col_schema.column_name)) - .and_then(|table_col| { - ColumnDataTypeWrapper::try_from(table_col.data_type.clone()) - .ok() - .map(|wrapper| wrapper.datatype()) - }); - let fixed_type = trace_semconv_fixed_type(&col_schema.column_name); - - if !observed_types - .iter() - .copied() - .any(is_trace_reconcile_candidate_type) - && existing_type - .map(|datatype| !is_trace_reconcile_candidate_type(datatype)) - .unwrap_or(true) - && fixed_type.is_none() - { - continue; - } - - // Decide the final type once per column, then rewrite all affected cells - // together in one row pass below. - let Some(decision) = choose_trace_reconcile_decision( - &col_schema.column_name, - &observed_types, - existing_type, - ) - .map_err(|_| { - enrich_trace_reconcile_error( - &req.table_name, - &col_schema.column_name, - &observed_types, - existing_type, - fixed_type, - ) - })? - else { - continue; - }; - let target_type = decision.target_type(); - - if !decision.requires_alter() - && observed_types - .iter() - .all(|observed| *observed == target_type) - && col_schema.datatype == target_type as i32 - { - continue; - } - - if decision.requires_alter() - && !pending_alter_columns.contains(&col_schema.column_name) - { - pending_alter_columns.push(col_schema.column_name.clone()); - } - - pending_rewrites.push(PendingTraceColumnRewrite { - col_idx, - target_type, - column_name: col_schema.column_name.clone(), - }); - } - - if pending_rewrites.is_empty() { - continue; - } - - validate_trace_column_rewrites(&rows.rows, &pending_rewrites, &req.table_name)?; - - if !pending_alter_columns.is_empty() { - self.alter_trace_table_columns_to_float64( - ctx, - &req.table_name, - &pending_alter_columns, - ) - .await?; - } - - // Update schema metadata before mutating row values so both stay in sync. - for pending_rewrite in &pending_rewrites { - rows.schema[pending_rewrite.col_idx].datatype = pending_rewrite.target_type as i32; - } - - // Apply all pending column rewrites in one row pass. - for row in &mut rows.rows { - for pending_rewrite in &pending_rewrites { - let Some(value) = row.values.get_mut(pending_rewrite.col_idx) else { - continue; - }; - let Some(request_type) = - value.value_data.as_ref().and_then(trace_value_datatype) - else { - continue; - }; - if request_type == pending_rewrite.target_type { - continue; - } - - value.value_data = coerce_value_data( - &value.value_data, - pending_rewrite.target_type, - request_type, - ) - .map_err(|_| { - error::InvalidParameterSnafu { - reason: format!( - "failed to coerce trace column '{}' in table '{}' from {:?} to {:?}", - pending_rewrite.column_name, - req.table_name, - request_type, - pending_rewrite.target_type - ), - } - .build() - })?; - } - } - } - - Ok(()) - } -} - -fn chunk_owned(items: Vec, chunk_size: usize) -> Vec> { - if items.is_empty() { - return Vec::new(); - } - - if chunk_size == 0 { - return vec![items]; - } - - let mut chunks = Vec::with_capacity(items.len().div_ceil(chunk_size)); - let mut iter = items.into_iter(); - while iter.len() > 0 { - chunks.push(iter.by_ref().take(chunk_size).collect()); - } - chunks -} - -/// Preserve the original alter failure status so chunk retry behavior stays correct. -fn wrap_trace_alter_failure(err: E) -> servers::error::Error -where - E: ErrorExt + Send + Sync + 'static, -{ - error::ExecuteGrpcQuerySnafu.into_error(BoxedError::new(err)) -} - -/// Derives `trace.conventions` from the request's resource/scope `schema_url`s. -/// A single distinct non-empty value is concrete; multiple distinct values are -/// `mixed`; none is `unknown`. `schema_url` is row-level in OTLP, so the -/// table-level value is best-effort per the RFC conflict rule. -fn trace_conventions(request: &ExportTraceServiceRequest) -> String { - let mut seen: Option<&str> = None; - let mut mixed = false; - - for resource_spans in &request.resource_spans { - let urls = std::iter::once(resource_spans.schema_url.as_str()).chain( - resource_spans - .scope_spans - .iter() - .map(|s| s.schema_url.as_str()), - ); - for url in urls { - if url.is_empty() { - continue; - } - match seen { - None => seen = Some(url), - Some(prev) if prev == url => {} - Some(_) => { - mixed = true; - break; - } - } - } - if mixed { - break; - } - } - - if mixed { - SEMANTIC_VALUE_MIXED.to_string() - } else { - seen.map(str::to_string) - .unwrap_or_else(|| SEMANTIC_VALUE_UNKNOWN.to_string()) - } -} - -#[cfg(test)] -mod tests { - use common_error::ext::ErrorExt; - use common_error::status_code::StatusCode; - use servers::query_handler::TraceIngestOutcome; - - use super::{ChunkFailureReaction, Instance, chunk_owned, wrap_trace_alter_failure}; - use crate::metrics::OTLP_TRACES_FAILURE_COUNT; - - #[test] - fn test_chunk_owned() { - let chunks = chunk_owned(vec![1, 2, 3], 2); - assert_eq!(chunks.iter().map(Vec::len).collect::>(), vec![2, 1]); - - let chunks = chunk_owned(vec![1, 2, 3], 0); - assert_eq!(chunks.len(), 1); - assert_eq!(chunks[0].len(), 3); - - assert!(chunk_owned::(Vec::new(), 0).is_empty()); - } - - #[test] - fn test_classify_trace_chunk_failure() { - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::InvalidArguments), - ChunkFailureReaction::RetryPerSpan - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::InvalidSyntax), - ChunkFailureReaction::RetryPerSpan - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::Unsupported), - ChunkFailureReaction::RetryPerSpan - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::TableColumnNotFound), - ChunkFailureReaction::RetryPerSpan - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::TableNotFound), - ChunkFailureReaction::RetryPerSpan - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::DatabaseNotFound), - ChunkFailureReaction::DiscardChunk - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::DeadlineExceeded), - ChunkFailureReaction::Propagate - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::Cancelled), - ChunkFailureReaction::Propagate - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::StorageUnavailable), - ChunkFailureReaction::Propagate - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::Internal), - ChunkFailureReaction::Propagate - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::RegionNotReady), - ChunkFailureReaction::Propagate - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::TableUnavailable), - ChunkFailureReaction::Propagate - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::RegionBusy), - ChunkFailureReaction::Propagate - ); - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::RuntimeResourcesExhausted), - ChunkFailureReaction::Propagate - ); - } - - #[test] - fn test_classify_trace_span_failure() { - assert!(Instance::should_propagate_trace_span_failure( - StatusCode::DeadlineExceeded - )); - assert!(Instance::should_propagate_trace_span_failure( - StatusCode::StorageUnavailable - )); - assert!(!Instance::should_propagate_trace_span_failure( - StatusCode::InvalidArguments - )); - } - - #[test] - fn test_add_trace_write_cost() { - let mut outcome = TraceIngestOutcome::default(); - Instance::add_trace_write_cost(&mut outcome, 3); - Instance::add_trace_write_cost(&mut outcome, 5); - assert_eq!(outcome.write_cost, 8); - } - - #[test] - fn test_finish_trace_failure_message() { - let message = Instance::finish_trace_failure_message( - 3, - 2, - vec!["Rejected span trace:span (InvalidArguments)".to_string()], - ) - .unwrap(); - assert!(message.contains("Accepted 3 spans, rejected 2 spans")); - assert!(message.contains("Rejected span trace:span")); - - assert_eq!(Instance::finish_trace_failure_message(2, 0, vec![]), None); - } - - #[test] - fn test_finish_trace_failure_message_without_detail_messages() { - assert_eq!( - Instance::finish_trace_failure_message(0, 2, vec![]), - Some("Accepted 0 spans, rejected 2 spans".to_string()) - ); - } - - #[test] - fn test_push_trace_failure_message_increments_labeled_counter() { - let label = "retry_per_span_counter_test"; - let initial = OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(); - let mut messages = Vec::new(); - - Instance::push_trace_failure_message( - &mut messages, - label, - "Chunk fallback triggered by InvalidArguments".to_string(), - ); - - assert_eq!(messages.len(), 1); - assert_eq!( - OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(), - initial + 1 - ); - } - - #[test] - fn test_push_trace_failure_message_caps_recorded_messages() { - let label = "retry_per_span_limit_test"; - let mut messages = Vec::new(); - - for idx in 0..=4 { - Instance::push_trace_failure_message(&mut messages, label, format!("failure-{idx}")); - } - - assert_eq!(messages.len(), 4); - assert_eq!( - messages, - vec![ - "failure-0".to_string(), - "failure-1".to_string(), - "failure-2".to_string(), - "failure-3".to_string() - ] - ); - } - - #[test] - fn test_classify_trace_chunk_failure_defaults_to_discard() { - assert_eq!( - Instance::classify_trace_chunk_failure(StatusCode::Unknown), - ChunkFailureReaction::DiscardChunk - ); - } - - #[test] - fn test_wrap_trace_alter_failure_preserves_status_code() { - let err = wrap_trace_alter_failure( - servers::error::TableNotFoundSnafu { - catalog: "greptime".to_string(), - schema: "public".to_string(), - table: "trace_type_missing".to_string(), - } - .build(), - ); - - assert_eq!(err.status_code(), StatusCode::TableNotFound); - } - - use opentelemetry_proto::tonic::trace::v1::{ResourceSpans, ScopeSpans}; - - use super::{ExportTraceServiceRequest, trace_conventions}; - - fn resource_spans(resource_url: &str, scope_urls: &[&str]) -> ResourceSpans { - ResourceSpans { - schema_url: resource_url.to_string(), - scope_spans: scope_urls - .iter() - .map(|u| ScopeSpans { - schema_url: u.to_string(), - ..Default::default() - }) - .collect(), - ..Default::default() - } - } - - #[test] - fn test_trace_conventions() { - let unknown = ExportTraceServiceRequest::default(); - assert_eq!(trace_conventions(&unknown), "unknown"); - - let url = "https://opentelemetry.io/schemas/1.27.0"; - let single = ExportTraceServiceRequest { - resource_spans: vec![resource_spans("", &[url, url])], - }; - assert_eq!(trace_conventions(&single), url); - - let resource_level = ExportTraceServiceRequest { - resource_spans: vec![resource_spans(url, &[""])], - }; - assert_eq!(trace_conventions(&resource_level), url); - - let conflicting = ExportTraceServiceRequest { - resource_spans: vec![resource_spans( - "", - &[url, "https://opentelemetry.io/schemas/1.30.0"], - )], - }; - assert_eq!(trace_conventions(&conflicting), "mixed"); - } -} diff --git a/src/frontend/src/instance/otlp/README.md b/src/frontend/src/instance/otlp/README.md new file mode 100644 index 00000000000..6380d56b7ac --- /dev/null +++ b/src/frontend/src/instance/otlp/README.md @@ -0,0 +1,182 @@ +# OTLP trace ingestion + +This directory owns the frontend orchestration for OTLP trace ingestion. It +starts after the HTTP server has decoded an OTLP protobuf request and ends after +the main trace table and its derived lookup tables have been written. + +The responsibility is split across crates: + +- `servers` decodes OTLP requests, parses spans, and converts spans into row + insert requests; +- this module chunks requests, reconciles v1 schemas, writes rows, handles + fallback, and accounts for accepted and rejected spans; +- `operator` creates or extends trace tables and performs the inserts. + +## End-to-end path + +```mermaid +flowchart TD + A["POST /v1/otlp/v1/traces
OTLP protobuf"] --> B["HTTP handler
decode, select table and pipeline"] + B --> C["Instance::traces
permission and interceptor"] + C --> D["Derive trace conventions
parse resource/scope groups"] + D --> E["Chunk spans within each group"] + E --> F{"Trace pipeline"} + + F -->|greptime_trace_v0| G["v0 conversion
fixed columns and JSON attributes"] + F -->|greptime_trace_v1| H["v1 conversion
flattened dynamic attributes"] + H --> I["Collect request-wide schema observations"] + I --> J["Resolve types, prevalidate rewrites,
apply DDL, and recheck table schema"] + + G --> K["Write main-table chunks"] + J --> K + K --> L{"Chunk result"} + L -->|success| M["Accept spans and collect
service/operation data"] + L -->|deterministic failure| N["Retry spans individually"] + L -->|discard| O["Count rejected spans"] + L -->|propagate| P["Return a request error"] + N --> M + N --> O + + M --> Q["Write <table>_services and
<table>_operations"] + O --> R["Build TraceIngestOutcome"] + Q --> R + R --> S["HTTP full success,
partial success, or failure"] +``` + +### 1. HTTP ingress + +The route is registered as `POST /v1/otlp/v1/traces` in +[`servers/src/http.rs`](../../../../servers/src/http.rs). The handler in +[`servers/src/http/otlp.rs`](../../../../servers/src/http/otlp.rs): + +1. accepts protobuf OTLP payloads and rejects JSON content types; +2. selects the main table, defaulting to `opentelemetry_traces`; +3. resolves the requested pipeline; and +4. calls `OpenTelemetryProtocolHandler::traces`. + +The trace converter currently accepts the internal `greptime_trace_v0` and +`greptime_trace_v1` pipelines. Their `PipelineWay` variants are defined in +[`pipeline/src/manager.rs`](../../../../pipeline/src/manager.rs). + +### 2. Parse and chunk + +[`Instance::traces`](../otlp.rs) checks the OTLP permission, runs the protocol +interceptor, derives the semantic-convention version, and parses the request. +The parser in +[`servers/src/otlp/trace/span.rs`](../../../../servers/src/otlp/trace/span.rs) +produces one `TraceSpanGroup` per resource/scope pair. + +`ingest_trace_spans` in [`trace_ingest.rs`](trace_ingest.rs) splits the spans in +each group into owned chunks. `trace_ingest_chunk_size` defaults to 128; setting +it to 0 disables splitting. The option is defined in +[`service_config/otlp.rs`](../../service_config/otlp.rs). + +### 3. Convert the selected trace model + +| Model | Conversion and storage shape | Frontend behavior | +| --- | --- | --- | +| v0 | [`trace/v0.rs`](../../../../servers/src/otlp/trace/v0.rs) writes a fixed span schema and keeps span, scope, and resource attributes in JSON columns. | Converts and writes one chunk at a time through the legacy log insertion path. | +| v1 | [`trace/v1.rs`](../../../../servers/src/otlp/trace/v1.rs) flattens OTLP attributes into dynamic columns and records the value types seen in each chunk. | Scans chunks one at a time to build a request-wide schema plan, then materializes and writes one chunk at a time. | + +Consequently, v1 chunks from one request can have different schemas: attribute +keys are data-dependent. The request-wide plan reconciles compatible column +observations across all chunks. The v0 main-table schema is comparatively fixed. + +The v1 converter intentionally uses `write_column_unchecked` from +[`row_writer.rs`](../../../../servers/src/row_writer.rs) to preserve the original +dynamic values instead of making the first value in a chunk decide the column +type. `MultiTableData` later pads sparse rows to the completed chunk schema. The +frontend schema planner is the validation and reconciliation boundary. + +## Trace v1 schema reconciliation + +The v1 path performs these steps before the normal chunk writes: + +1. Convert each chunk into a temporary `RowInsertRequest` and a + `TraceBatchSchema` that records observed scalar types, sparse row presence, + and binary logical types. Drop the dense request after collecting its + lightweight schema observations. +2. Mark a chunk for span-only fallback if one column contains both raw OTLP + bytes and arrays or key-value lists encoded as JSONB. Both use protobuf + `BinaryValue`, but a Greptime column cannot represent both logical schemas. + The chunk bypasses request-wide schema planning and is reconciled span by + span. +3. Aggregate the compatible observations into `TraceRequestSchema` and + compare them with the existing table schema and the fixed types in + [`trace_semconv.rs`](trace_semconv.rs). +4. Select a stable target type. Existing table types are authoritative when the + incoming values can be coerced. New columns use all remaining request + observations. + The supported existing-column widening is `Int64` to `Float64`. +5. Use `prepare_trace_column_rewrites` in + [`trace_types.rs`](trace_types.rs) to precompute every coercion without + mutating rows. A chunk whose coercion cannot be prepared is split into + single-span planning units; only spans that still fail are removed from the + request-wide plan, so invalid data cannot contribute to persistent DDL. +6. Create the table or add missing columns with `ensure_trace_table_on_demand`, + and widen planned numeric columns. The implementation is in + [`operator/src/insert.rs`](../../../../operator/src/insert.rs). +7. Read and plan against the table schema again after DDL to account for + concurrent creates or alters. Then rematerialize, rewrite, and write one + chunk at a time. If the schema does not converge, affected chunks use the + per-request reconciliation path. + +This ordering keeps conversion failures atomic: rows are not partly rewritten, +and columns from an invalid chunk are not added to the table before that chunk +falls back. + +## Writes, fallback, and accounting + +Schema planning and request reconciliation happen before a main-table write. +`classify_trace_prewrite_failure` in [`trace_ingest.rs`](trace_ingest.rs) +determines the next action when that preparation fails: + +| Reaction | Behavior | +| --- | --- | +| `RetryPerSpan` | Prepare and write each span independently. Successful spans are accepted; deterministic per-span preparation failures are rejected. A v1 chunk with incompatible raw-binary and JSONB values starts here without attempting a chunk write. | +| `DiscardChunk` | Do not retry the chunk at span granularity. Count every span in the chunk as rejected. | +| `Propagate` | Return retryable or unclassified failures instead of reporting a partial success. Retry hints take precedence over status-code classification. | + +Span-specific invalid values are isolated before request-wide schema planning. +Therefore, a deterministic error that escapes request-wide preparation rejects +all remaining spans instead of repeating the same global failure per span. + +Once a main-table write is dispatched, any returned error is propagated. The +distributed insert may have committed writes on some peers before another peer +failed, so retrying or rejecting the whole chunk would risk duplicate data or +incorrect accounting. + +Accounting follows the main-table write: + +- a span is accepted only after its main-table insert is confirmed; +- only accepted spans contribute service and operation rows; +- auxiliary rows are deduplicated and written afterwards to + `_services` and `_operations`; +- auxiliary rows accumulated by earlier successful chunks are still flushed + before a later main-table error is returned; +- an auxiliary-table failure adds failure detail but does not change the + accepted or rejected span counts; and +- failure details are bounded before they are folded into `TraceIngestOutcome`. + +Finally, the HTTP handler returns: + +- full success when there are no rejected spans or failure details; +- OTLP partial success when spans were rejected or auxiliary/fallback details + were recorded; or +- HTTP `400 Bad Request` when no span was accepted and at least one was + rejected. + +A propagated error bypasses that outcome and fails the request. + +## Code map + +| Area | Entry point | +| --- | --- | +| HTTP route and response mapping | [`servers/src/http/otlp.rs`](../../../../servers/src/http/otlp.rs) | +| Frontend protocol handler | [`instance/otlp.rs`](../otlp.rs) | +| Chunking, v1 planning, writes, fallback, accounting | [`trace_ingest.rs`](trace_ingest.rs) | +| Reconciliation decisions and atomic value rewrites | [`trace_types.rs`](trace_types.rs) | +| Semantic-convention fixed types | [`trace_semconv.rs`](trace_semconv.rs) | +| OTLP span parsing | [`servers/src/otlp/trace/span.rs`](../../../../servers/src/otlp/trace/span.rs) | +| v0 and v1 row conversion | [`servers/src/otlp/trace`](../../../../servers/src/otlp/trace.rs) | +| Trace-aware table creation and inserts | [`operator/src/insert.rs`](../../../../operator/src/insert.rs) | diff --git a/src/frontend/src/instance/otlp/trace_ingest.rs b/src/frontend/src/instance/otlp/trace_ingest.rs new file mode 100644 index 00000000000..0866189b571 --- /dev/null +++ b/src/frontend/src/instance/otlp/trace_ingest.rs @@ -0,0 +1,1965 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::sync::Arc; + +use api::helper::ColumnDataTypeWrapper; +use api::v1::alter_table_expr::Kind; +use api::v1::{ + AlterTableExpr, ColumnDataType, ColumnSchema, ModifyColumnType, ModifyColumnTypes, + RowInsertRequest, RowInsertRequests, Rows, Value, +}; +use client::Output; +use common_error::ext::{BoxedError, ErrorExt}; +use common_error::status_code::StatusCode; +use common_telemetry::warn; +use datatypes::prelude::ConcreteDataType; +use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; +use pipeline::{GreptimePipelineParams, PipelineWay}; +use servers::error::{self, Result as ServerResult}; +use servers::otlp; +use servers::otlp::coerce::{coerce_value_data, trace_value_datatype}; +use servers::otlp::trace::TraceAuxData; +use servers::otlp::trace::span::{TraceSpan, TraceSpanGroup}; +use servers::otlp::trace::v1::{TraceBatchSchema, TraceBinaryType, TraceRetryColumns}; +use servers::query_handler::{PipelineHandlerRef, TraceIngestOutcome}; +use session::context::QueryContextRef; +use snafu::{IntoError, ResultExt}; +use table::requests::{ + SEMANTIC_PIPELINE, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SEMANTIC_TRACE_CONVENTIONS, + SEMANTIC_VALUE_MIXED, SEMANTIC_VALUE_UNKNOWN, SIGNAL_TYPE_TRACE, SOURCE_OPENTELEMETRY, + TABLE_DATA_MODEL_TRACE_V1, +}; + +use crate::instance::Instance; +use crate::instance::otlp::trace_semconv::trace_semconv_fixed_type; +use crate::instance::otlp::trace_types::{ + PendingTraceColumnRewrite, PreparedTraceColumnRewrites, TraceColumnRewriteError, + choose_trace_reconcile_decision, enrich_trace_reconcile_error, + is_trace_reconcile_candidate_type, prepare_trace_column_rewrites, push_observed_trace_type, +}; +use crate::metrics::{OTLP_TRACES_FAILURE_COUNT, OTLP_TRACES_ROWS}; + +const TRACE_FAILURE_MESSAGE_LIMIT: usize = 4; + +/// Determines how trace ingestion responds to a failure before a write is dispatched. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ChunkFailureReaction { + RetryPerSpan, + DiscardChunk, + Propagate, +} + +impl ChunkFailureReaction { + fn as_metric_label(self) -> &'static str { + match self { + Self::RetryPerSpan => "retry_per_span", + Self::DiscardChunk => "discard_chunk", + Self::Propagate => "propagate_failure", + } + } +} + +/// Shared dependencies and request metadata used while ingesting trace chunks. +struct TraceChunkIngestContext<'a> { + pipeline_handler: PipelineHandlerRef, + pipeline: &'a PipelineWay, + pipeline_params: &'a GreptimePipelineParams, + table_name: &'a str, + is_trace_v1_model: bool, +} + +/// Accumulates trace outcomes, auxiliary rows, and bounded failure details. +struct TraceIngestState { + aux_data: TraceAuxData, + outcome: TraceIngestOutcome, + failure_messages: Vec, +} + +/// How a v1 chunk should be reconciled when it is written. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TraceChunkSchemaState { + Prepared, + ReconcilePerChunk, + SpanFallbackOnly, +} + +impl TraceChunkSchemaState { + fn mark_for_reconcile(&mut self) { + if matches!(self, Self::Prepared) { + *self = Self::ReconcilePerChunk; + } + } +} + +/// Sparse representation retained for a v1 trace chunk. +struct TraceChunkRetry { + table_name: String, + schema: Vec, + rows: Vec, +} + +/// The projected retry values and identifying metadata for one span. +struct TraceRetryRow { + span_metadata: TraceSpanMetadata, + values: Vec, +} + +/// One retry cell and the logical type to restore before insertion. +struct TraceRetryValue { + column_index: usize, + value: Value, + schema_type: Option, +} + +/// The scalar or binary logical type recorded for a retry value. +enum TraceRetrySchemaType { + Scalar(ColumnDataType), + Binary(TraceBinaryType), +} + +/// Span identity and operation fields retained after request conversion. +struct TraceSpanMetadata { + trace_id: String, + span_id: String, + operation: Option<(String, String, String)>, +} + +impl From<&TraceSpan> for TraceSpanMetadata { + fn from(span: &TraceSpan) -> Self { + Self { + trace_id: span.trace_id.clone(), + span_id: span.span_id.clone(), + operation: span.service_name.as_ref().map(|service_name| { + ( + service_name.clone(), + span.span_name.clone(), + span.span_kind.clone(), + ) + }), + } + } +} + +impl TraceSpanMetadata { + fn add_to_aux_data(self, aux_data: &mut TraceAuxData) { + if let Some((service_name, span_name, span_kind)) = self.operation { + aux_data.services.insert(service_name.clone()); + aux_data + .operations + .insert((service_name, span_name, span_kind)); + } + } +} + +impl TraceChunkRetry { + fn try_new( + table_name: String, + rows: Rows, + span_metadata: Vec, + retry_columns: TraceRetryColumns, + ) -> ServerResult { + let Rows { schema, rows } = rows; + let row_count = rows.len(); + if row_count != span_metadata.len() { + return error::InternalSnafu { + err_msg: format!( + "trace row count {} does not match span metadata count {}", + row_count, + span_metadata.len() + ), + } + .fail(); + } + + let retry_columns_by_index = schema + .iter() + .map(|column| retry_columns.get(&column.column_name)) + .collect::>(); + let mut fixed_columns = Vec::with_capacity(schema.len()); + let mut sparse_columns = vec![Vec::new(); rows.len()]; + for (column_index, retry_column) in retry_columns_by_index.iter().enumerate() { + let Some(retry_column) = retry_column else { + fixed_columns.push(column_index); + continue; + }; + for present_row in &retry_column.present_rows { + if *present_row < sparse_columns.len() { + sparse_columns[*present_row].push(column_index); + } + } + } + + let mut retry_rows = Vec::with_capacity(rows.len()); + for (row_index, (mut row, span_metadata)) in rows.into_iter().zip(span_metadata).enumerate() + { + if row.values.len() > schema.len() { + return error::InternalSnafu { + err_msg: format!( + "trace row column count {} exceeds schema column count {}", + row.values.len(), + schema.len() + ), + } + .fail(); + } + let sparse_row = &mut sparse_columns[row_index]; + let mut selected_columns = Vec::with_capacity(fixed_columns.len() + sparse_row.len()); + selected_columns.extend_from_slice(&fixed_columns); + selected_columns.append(sparse_row); + selected_columns.sort_unstable(); + selected_columns.dedup(); + + let mut values = Vec::with_capacity(selected_columns.len()); + for column_index in selected_columns { + let column = &schema[column_index]; + let value = row + .values + .get_mut(column_index) + .map(std::mem::take) + .unwrap_or_default(); + let value_type = value.value_data.as_ref().and_then(trace_value_datatype); + let binary_type = if value_type == Some(ColumnDataType::Binary) { + retry_columns_by_index[column_index].and_then(|retry_column| { + retry_column + .binary_types + .binary_search_by_key(&row_index, |(row_index, _)| *row_index) + .ok() + .map(|index| retry_column.binary_types[index].1) + }) + } else { + None + }; + let schema_type = binary_type.map(TraceRetrySchemaType::Binary).or_else(|| { + value_type + .filter(|datatype| { + is_trace_reconcile_candidate_type(*datatype) + && ColumnDataType::try_from(column.datatype).ok() != Some(*datatype) + }) + .map(TraceRetrySchemaType::Scalar) + }); + values.push(TraceRetryValue { + column_index, + value, + schema_type, + }); + } + retry_rows.push(TraceRetryRow { + span_metadata, + values, + }); + } + + Ok(Self { + table_name, + schema, + rows: retry_rows, + }) + } + + fn to_request(&self) -> ServerResult { + let mut rows = Vec::with_capacity(self.rows.len()); + for retry_row in &self.rows { + let mut values = vec![Value::default(); self.schema.len()]; + for retry_value in &retry_row.values { + let Some(value) = values.get_mut(retry_value.column_index) else { + return error::InternalSnafu { + err_msg: format!( + "trace retry column index {} is out of bounds", + retry_value.column_index + ), + } + .fail(); + }; + *value = retry_value.value.clone(); + } + rows.push(api::v1::Row { values }); + } + + Ok(RowInsertRequest { + table_name: self.table_name.clone(), + rows: Some(Rows { + schema: self.schema.clone(), + rows, + }), + }) + } + + fn into_single_span_chunks(self) -> ServerResult> { + let Self { + table_name, + schema, + rows, + } = self; + rows.into_iter() + .map(|row| { + let (span_metadata, rows_for_insert) = row.into_rows(&schema)?; + Self::try_new( + table_name.clone(), + rows_for_insert, + vec![span_metadata], + TraceRetryColumns::default(), + ) + }) + .collect() + } + + fn add_to_aux_data(self, aux_data: &mut TraceAuxData) { + for row in self.rows { + row.span_metadata.add_to_aux_data(aux_data); + } + } +} + +impl TraceRetryRow { + fn into_rows(self, schema: &[ColumnSchema]) -> ServerResult<(TraceSpanMetadata, Rows)> { + let mut projected_schema = Vec::with_capacity(self.values.len()); + let mut projected_values = Vec::with_capacity(self.values.len()); + for retry_value in self.values { + let Some(mut column) = schema.get(retry_value.column_index).cloned() else { + return error::InternalSnafu { + err_msg: format!( + "trace fallback column index {} is out of bounds", + retry_value.column_index + ), + } + .fail(); + }; + match retry_value.schema_type { + Some(TraceRetrySchemaType::Scalar(datatype)) => { + column.datatype = datatype as i32; + column.datatype_extension = None; + } + Some(TraceRetrySchemaType::Binary(binary_type)) => { + binary_type.apply_to_schema(&mut column); + } + None => {} + } + projected_schema.push(column); + projected_values.push(retry_value.value); + } + + Ok(( + self.span_metadata, + Rows { + schema: projected_schema, + rows: vec![api::v1::Row { + values: projected_values, + }], + }, + )) + } +} + +/// Request-wide schema observations collected across all v1 trace chunks. +#[derive(Default)] +struct TraceRequestSchema { + columns: Vec, + column_indexes: HashMap, +} + +/// Per-column observations and the resolved target type for a trace request. +struct TraceColumnRequestSchema { + schema: ColumnSchema, + batches: BTreeMap, + target_type: Option, +} + +/// The schema and value types observed for one column in one chunk. +struct TraceBatchColumnObservation { + value_types: Vec, + schema_type: ColumnDataType, + concrete_type: ConcreteDataType, +} + +/// Maps incompatible columns to the chunk indexes that must use fallback. +type TraceSchemaExclusions = HashMap>; + +/// Table columns to add or widen before inserting the prepared chunks. +struct TraceTablePreAlter { + ensure_columns: Vec, + modify_float64_columns: Vec, + alter_existing: bool, +} + +/// Either a ready table-alter plan or the observations that cannot be unified. +enum TraceRequestSchemaPlan { + Ready(TraceTablePreAlter), + IncompatibleObservations(TraceSchemaExclusions), +} + +impl TraceTablePreAlter { + fn requires_ddl(&self) -> bool { + !self.ensure_columns.is_empty() || !self.modify_float64_columns.is_empty() + } +} + +impl TraceRequestSchema { + fn observe_batch_schema( + &mut self, + batch_index: usize, + schema: &[ColumnSchema], + batch_schema: &TraceBatchSchema, + ) { + for column in schema { + if batch_schema.has_incompatible_logical_types_for(&column.column_name) { + continue; + } + let observed_types = batch_schema.value_types(&column.column_name); + self.observe_trace_column(batch_index, column, None); + if let Some(observed_types) = observed_types { + for value_type in observed_types { + self.observe_trace_column(batch_index, column, Some(*value_type)); + } + } + } + } + + fn observe_retry_chunk( + &mut self, + batch_index: usize, + chunk: &TraceChunkRetry, + ) -> ServerResult<()> { + // Split chunks already carry row-projected schemas; untouched chunks keep + // their original batch schema while observations stay sparse. + for row in &chunk.rows { + for retry_value in &row.values { + let Some(column) = chunk.schema.get(retry_value.column_index) else { + return error::InternalSnafu { + err_msg: format!( + "trace retry column index {} is out of bounds", + retry_value.column_index + ), + } + .fail(); + }; + let value_type = retry_value + .value + .value_data + .as_ref() + .and_then(trace_value_datatype); + self.observe_trace_column(batch_index, column, value_type); + } + } + Ok(()) + } + + fn observe_trace_column( + &mut self, + batch_index: usize, + schema: &ColumnSchema, + value_type: Option, + ) { + let column_schema = self.observe_column(schema); + if let Ok(current_type) = ColumnDataType::try_from(schema.datatype) { + let observation = column_schema.batches.entry(batch_index).or_insert_with(|| { + let wrapper = + ColumnDataTypeWrapper::new(current_type, schema.datatype_extension.clone()); + TraceBatchColumnObservation { + value_types: Vec::new(), + schema_type: current_type, + concrete_type: ConcreteDataType::from(wrapper), + } + }); + push_observed_trace_type(&mut observation.value_types, current_type); + } + if let Some(value_type) = value_type { + column_schema.observe_type(batch_index, value_type); + } + } + + fn incompatible_schema_observations( + &self, + table_schema: Option<&datatypes::schema::Schema>, + ) -> TraceSchemaExclusions { + let mut exclusions = HashMap::new(); + for column in &self.columns { + let existing_type = table_schema + .and_then(|schema| schema.column_schema_by_name(&column.schema.column_name)) + .and_then(|table_col| { + ColumnDataTypeWrapper::try_from(table_col.data_type.clone()) + .ok() + .map(|wrapper| { + let datatype = wrapper.datatype(); + (datatype, ConcreteDataType::from(wrapper)) + }) + }); + let incompatible_batches = column.incompatible_schema_batches(existing_type); + if !incompatible_batches.is_empty() { + exclusions.insert(column.schema.column_name.clone(), incompatible_batches); + } + } + exclusions + } + + fn resolve_table_schema( + &mut self, + table_schema: Option<&datatypes::schema::Schema>, + ) -> TraceRequestSchemaPlan { + let mut pre_alter = TraceTablePreAlter { + ensure_columns: Vec::new(), + modify_float64_columns: Vec::new(), + alter_existing: table_schema.is_some(), + }; + + for column in &mut self.columns { + let Some(current_type) = ColumnDataType::try_from(column.schema.datatype).ok() else { + continue; + }; + let observed_types = column.observed_types(); + + let existing_type = table_schema + .and_then(|schema| schema.column_schema_by_name(&column.schema.column_name)) + .and_then(|table_col| { + ColumnDataTypeWrapper::try_from(table_col.data_type.clone()) + .ok() + .map(|wrapper| wrapper.datatype()) + }); + let fixed_type = trace_semconv_fixed_type(&column.schema.column_name); + + let needs_reconcile = observed_types + .iter() + .copied() + .any(is_trace_reconcile_candidate_type) + || existing_type + .map(is_trace_reconcile_candidate_type) + .unwrap_or(false) + || fixed_type.is_some(); + + let target_type = if needs_reconcile { + let decision = match choose_trace_reconcile_decision( + &column.schema.column_name, + &observed_types, + existing_type, + ) { + Ok(decision) => decision, + Err(_) => { + return TraceRequestSchemaPlan::IncompatibleObservations(HashMap::from([ + ( + column.schema.column_name.clone(), + column.batches.keys().copied().collect(), + ), + ])); + } + }; + decision + .map(|decision| { + if decision.requires_alter() { + pre_alter + .modify_float64_columns + .push(column.schema.column_name.clone()); + } + decision.target_type() + }) + .unwrap_or(current_type) + } else { + current_type + }; + + column.target_type = Some(target_type); + + if table_schema + .map(|schema| { + schema + .column_schema_by_name(&column.schema.column_name) + .is_none() + }) + .unwrap_or(false) + { + pre_alter.ensure_columns.push(column.target_schema()); + } + } + + if table_schema.is_none() { + pre_alter.ensure_columns = self + .columns + .iter() + .map(TraceColumnRequestSchema::target_schema) + .collect(); + } + + TraceRequestSchemaPlan::Ready(pre_alter) + } + + fn prepare_request_rewrite( + &self, + batch_index: usize, + request: &RowInsertRequest, + ) -> Result, TraceColumnRewriteError> { + let Some(rows) = request.rows.as_ref() else { + return Ok(None); + }; + let pending_rewrites = self.pending_batch_rewrites(batch_index, rows); + if pending_rewrites.is_empty() { + return Ok(None); + } + prepare_trace_column_rewrites(&rows.rows, pending_rewrites, &request.table_name).map(Some) + } + + fn apply_request_rewrite( + request: &mut RowInsertRequest, + prepared: PreparedTraceColumnRewrites, + ) { + if let Some(rows) = request.rows.as_mut() { + prepared.apply(rows); + } + } + + fn pending_batch_rewrites( + &self, + batch_index: usize, + rows: &api::v1::Rows, + ) -> Vec { + let mut pending_rewrites = Vec::new(); + for (col_idx, col_schema) in rows.schema.iter().enumerate() { + let Some(global_idx) = self.column_indexes.get(&col_schema.column_name).copied() else { + continue; + }; + let global_column = &self.columns[global_idx]; + let Some(target_type) = global_column.target_type() else { + continue; + }; + if !global_column + .batches + .get(&batch_index) + .is_some_and(|observation| { + observation + .value_types + .iter() + .any(|datatype| *datatype != target_type) + }) + { + continue; + } + pending_rewrites.push(PendingTraceColumnRewrite { + col_idx, + target_type, + column_name: global_column.schema.column_name.clone(), + }); + } + + pending_rewrites + } + + #[cfg(test)] + fn remove_batches(&mut self, batch_indexes: &HashSet) { + for column in &mut self.columns { + column + .batches + .retain(|batch_index, _| !batch_indexes.contains(batch_index)); + column.target_type = None; + } + self.columns.retain(|column| !column.batches.is_empty()); + + self.column_indexes.clear(); + for (index, column) in self.columns.iter().enumerate() { + self.column_indexes + .insert(column.schema.column_name.clone(), index); + } + } + + #[cfg(test)] + fn resolved_target_types(&self) -> Vec> { + self.columns + .iter() + .map(|column| column.target_type) + .collect() + } + + fn observe_column(&mut self, schema: &ColumnSchema) -> &mut TraceColumnRequestSchema { + if let Some(index) = self.column_indexes.get(&schema.column_name).copied() { + return &mut self.columns[index]; + } + + let index = self.columns.len(); + self.columns.push(TraceColumnRequestSchema { + schema: schema.clone(), + batches: BTreeMap::new(), + target_type: None, + }); + self.column_indexes + .insert(schema.column_name.clone(), index); + &mut self.columns[index] + } +} + +impl TraceColumnRequestSchema { + fn observe_type(&mut self, batch_index: usize, datatype: ColumnDataType) { + if let Some(observation) = self.batches.get_mut(&batch_index) { + push_observed_trace_type(&mut observation.value_types, datatype); + } + } + + fn observed_types(&self) -> Vec { + let mut observed_types = Vec::new(); + for datatype in self + .batches + .values() + .flat_map(|observation| &observation.value_types) + { + push_observed_trace_type(&mut observed_types, *datatype); + } + observed_types + } + + fn incompatible_schema_batches( + &self, + existing_type: Option<(ColumnDataType, ConcreteDataType)>, + ) -> HashSet { + let mut incompatible_batches = HashSet::new(); + if let Some((existing_datatype, existing_concrete_type)) = existing_type { + for (batch_index, observation) in &self.batches { + if trace_logical_types_incompatible( + observation.schema_type, + &observation.concrete_type, + existing_datatype, + &existing_concrete_type, + ) { + incompatible_batches.insert(*batch_index); + } + } + return incompatible_batches; + } + + let mut reconcile_candidates = Vec::<&TraceBatchColumnObservation>::new(); + for observation in self + .batches + .values() + .filter(|observation| is_trace_reconcile_candidate_type(observation.schema_type)) + { + if !reconcile_candidates.iter().any(|candidate| { + candidate.schema_type == observation.schema_type + && candidate.concrete_type == observation.concrete_type + }) { + reconcile_candidates.push(observation); + } + } + if !reconcile_candidates.is_empty() { + for (batch_index, observation) in &self.batches { + if reconcile_candidates.iter().any(|candidate| { + trace_logical_types_incompatible( + observation.schema_type, + &observation.concrete_type, + candidate.schema_type, + &candidate.concrete_type, + ) + }) { + incompatible_batches.insert(*batch_index); + } + } + return incompatible_batches; + } + + let mut schemas = Vec::<(ColumnDataType, ConcreteDataType, bool)>::new(); + for observation in self.batches.values() { + if schemas.iter().any(|(observed_datatype, observed_type, _)| { + *observed_datatype == observation.schema_type + && observed_type == &observation.concrete_type + }) { + continue; + } + + let mut conflicts = false; + for (observed_datatype, observed_type, observed_conflicts) in &mut schemas { + if trace_logical_types_incompatible( + *observed_datatype, + observed_type, + observation.schema_type, + &observation.concrete_type, + ) { + *observed_conflicts = true; + conflicts = true; + } + } + schemas.push(( + observation.schema_type, + observation.concrete_type.clone(), + conflicts, + )); + } + + for (batch_index, observation) in &self.batches { + if schemas + .iter() + .any(|(observed_datatype, observed_type, conflict)| { + *observed_datatype == observation.schema_type + && observed_type == &observation.concrete_type + && *conflict + }) + { + incompatible_batches.insert(*batch_index); + } + } + + incompatible_batches + } + + fn target_type(&self) -> Option { + self.target_type + .or_else(|| ColumnDataType::try_from(self.schema.datatype).ok()) + } + + fn target_schema(&self) -> ColumnSchema { + let mut schema = self.schema.clone(); + if let Some(target_type) = self.target_type() { + schema.datatype = target_type as i32; + } + schema + } +} + +impl Instance { + /// Ingest OTLP trace spans with chunk-level writes and span-level fallback on + /// deterministic chunk failures. + #[allow(clippy::too_many_arguments)] + pub(super) async fn ingest_trace_spans( + &self, + pipeline_handler: PipelineHandlerRef, + pipeline: &PipelineWay, + pipeline_params: &GreptimePipelineParams, + table_name: String, + groups: Vec, + conventions: &str, + ctx: QueryContextRef, + ) -> ServerResult { + let is_trace_v1_model = matches!(pipeline, PipelineWay::OtlpTraceDirectV1); + + // Only the main span table gets the identity; the derived `_services` / + // `_operations` lookup tables keep the unstamped `ctx`. + let main_ctx = { + let mut c = (*ctx).clone(); + c.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_TRACE); + c.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY); + if is_trace_v1_model { + c.set_extension(SEMANTIC_PIPELINE, TABLE_DATA_MODEL_TRACE_V1); + c.set_extension(SEMANTIC_TRACE_CONVENTIONS, conventions); + } + Arc::new(c) + }; + + let ingest_ctx = TraceChunkIngestContext { + pipeline_handler, + pipeline, + pipeline_params, + table_name: &table_name, + is_trace_v1_model, + }; + let mut ingest_state = TraceIngestState { + aux_data: TraceAuxData::default(), + outcome: TraceIngestOutcome::default(), + failure_messages: Vec::new(), + }; + + let main_result: ServerResult<()> = async { + if is_trace_v1_model { + let mut request_schema = TraceRequestSchema::default(); + let mut chunks = Vec::new(); + let mut chunk_states = Vec::new(); + // Convert each chunk once, then retain only its sparse values while + // request-wide schema planning runs. + for chunk in groups + .into_iter() + .flat_map(|group| chunk_owned(group.spans, self.trace_ingest_chunk_size)) + { + let batch_index = chunks.len(); + let span_metadata = chunk.iter().map(TraceSpanMetadata::from).collect(); + let (table_data, batch_schema) = + otlp::trace::v1::v1_to_main_table_data_with_schema(chunk)?; + let schema_state = if batch_schema.has_incompatible_logical_types() { + TraceChunkSchemaState::SpanFallbackOnly + } else { + TraceChunkSchemaState::Prepared + }; + let (schema, rows) = table_data.into_schema_and_rows(); + if schema_state == TraceChunkSchemaState::Prepared { + request_schema.observe_batch_schema(batch_index, &schema, &batch_schema); + } + chunks.push(TraceChunkRetry::try_new( + ingest_ctx.table_name.to_string(), + Rows { schema, rows }, + span_metadata, + batch_schema.into_retry_columns(), + )?); + chunk_states.push(schema_state); + } + + if !chunks.is_empty() { + self.ingest_trace_v1_prepared_chunks( + &ingest_ctx, + request_schema, + chunks, + chunk_states, + main_ctx.clone(), + &mut ingest_state, + ) + .await?; + } + } else { + for group in groups { + let chunks = chunk_owned(group.spans, self.trace_ingest_chunk_size); + for chunk in chunks { + self.ingest_trace_chunk( + &ingest_ctx, + chunk, + main_ctx.clone(), + &mut ingest_state, + ) + .await?; + } + } + } + + Ok(()) + } + .await; + + OTLP_TRACES_ROWS.inc_by(ingest_state.outcome.accepted_spans as u64); + + if !ingest_state.aux_data.is_empty() { + // Auxiliary trace tables are derived from spans whose main-table + // writes are already confirmed, so they never create new accepted + // spans and they do not affect rejected span counts. + let aux_requests = otlp::trace::to_grpc_insert_requests_for_aux_tables( + std::mem::take(&mut ingest_state.aux_data), + ingest_ctx.pipeline, + ingest_ctx.table_name, + ); + + match aux_requests { + Ok((aux_requests, _)) if !aux_requests.inserts.is_empty() => { + match self + .insert_trace_requests(aux_requests, ingest_ctx.is_trace_v1_model, ctx) + .await + { + Ok(output) => { + Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); + } + Err(err) => { + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + "aux_table_update_failed", + format!( + "Auxiliary trace tables were not fully updated ({})", + err.status_code().as_ref() + ), + ); + } + } + } + // Preserve the existing conversion-error behavior after a + // successful main write, but never mask an earlier main error. + Err(err) if main_result.is_ok() => return Err(err), + Err(err) => Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + "aux_table_update_failed", + format!( + "Auxiliary trace tables were not fully updated ({})", + err.status_code().as_ref() + ), + ), + Ok(_) => {} + } + } + + main_result?; + + ingest_state.outcome.error_message = Self::finish_trace_failure_message( + ingest_state.outcome.accepted_spans, + ingest_state.outcome.rejected_spans, + ingest_state.failure_messages, + ); + + Ok(ingest_state.outcome) + } + + /// Ingest one owned trace chunk so successful spans can be moved into the + /// accepted set without extra cloning. + async fn ingest_trace_chunk( + &self, + ingest_ctx: &TraceChunkIngestContext<'_>, + chunk: Vec, + ctx: QueryContextRef, + ingest_state: &mut TraceIngestState, + ) -> ServerResult<()> { + // Try the fast path first so healthy batches keep their original + // throughput and write amplification stays low. + let (requests, chunk_rows) = otlp::trace::to_grpc_insert_requests_from_spans( + &chunk, + ingest_ctx.pipeline, + ingest_ctx.pipeline_params, + ingest_ctx.table_name, + &ctx, + ingest_ctx.pipeline_handler.clone(), + )?; + + let output = match self + .insert_trace_requests(requests, ingest_ctx.is_trace_v1_model, ctx) + .await + { + Ok(output) => output, + Err(err) => { + // A distributed insert can partially commit before returning an + // error, so retrying or rejecting the whole chunk is unsafe. + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::Propagate.as_metric_label(), + format!( + "Propagating chunk write failure ({})", + err.status_code().as_ref() + ), + ); + return Err(err); + } + }; + + Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); + ingest_state.outcome.accepted_spans += chunk_rows; + for span in &chunk { + ingest_state.aux_data.observe_span(span); + } + + Ok(()) + } + + async fn ingest_trace_v1_prepared_chunks( + &self, + ingest_ctx: &TraceChunkIngestContext<'_>, + mut request_schema: TraceRequestSchema, + mut chunks: Vec, + mut chunk_states: Vec, + ctx: QueryContextRef, + ingest_state: &mut TraceIngestState, + ) -> ServerResult<()> { + if let Err(err) = self + .prepare_trace_v1_request_schema( + ingest_ctx.table_name, + &mut request_schema, + &mut chunks, + &mut chunk_states, + &ctx, + ) + .await + { + if matches!( + Self::classify_trace_prewrite_failure(err.status_code(), err.is_retryable()), + ChunkFailureReaction::Propagate + ) { + return Err(err); + } + + // Span-specific conversion failures were already isolated during + // prevalidation. An error escaping request-wide schema preparation + // applies to every remaining span, so repeating it per span cannot + // recover valid data. + let span_count = chunks.iter().map(|chunk| chunk.rows.len()).sum::(); + ingest_state.outcome.rejected_spans += span_count; + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::DiscardChunk.as_metric_label(), + format!( + "Discarded {} spans after pre-write request failure ({})", + span_count, + err.status_code().as_ref() + ), + ); + return Ok(()); + } + + for (batch_index, (retry, schema_state)) in chunks.into_iter().zip(chunk_states).enumerate() + { + self.ingest_trace_v1_chunk( + &request_schema, + batch_index, + retry, + schema_state, + ctx.clone(), + ingest_state, + ) + .await?; + } + + Ok(()) + } + + async fn ingest_trace_v1_chunk( + &self, + request_schema: &TraceRequestSchema, + batch_index: usize, + retry: TraceChunkRetry, + mut schema_state: TraceChunkSchemaState, + ctx: QueryContextRef, + ingest_state: &mut TraceIngestState, + ) -> ServerResult<()> { + if schema_state == TraceChunkSchemaState::SpanFallbackOnly { + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::RetryPerSpan.as_metric_label(), + "Chunk fallback triggered by incompatible binary and JSON values".to_string(), + ); + return self + .ingest_trace_v1_rows_span_by_span(retry, ctx, ingest_state) + .await; + } + + let mut request = retry.to_request()?; + let _ = Self::prepare_trace_v1_chunk_rewrite( + request_schema, + batch_index, + &mut request, + &mut schema_state, + )?; + let span_count = retry.rows.len(); + let mut requests = RowInsertRequests { + inserts: vec![request], + }; + if schema_state == TraceChunkSchemaState::ReconcilePerChunk + && let Err(err) = self.reconcile_trace_column_types(&mut requests, &ctx).await + { + match Self::classify_trace_prewrite_failure(err.status_code(), err.is_retryable()) { + ChunkFailureReaction::RetryPerSpan => { + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::RetryPerSpan.as_metric_label(), + format!("Chunk fallback triggered by {}", err.status_code().as_ref()), + ); + return self + .ingest_trace_v1_rows_span_by_span(retry, ctx, ingest_state) + .await; + } + ChunkFailureReaction::DiscardChunk => { + ingest_state.outcome.rejected_spans += span_count; + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::DiscardChunk.as_metric_label(), + format!( + "Discarded {} spans after pre-write chunk failure ({})", + span_count, + err.status_code().as_ref() + ), + ); + return Ok(()); + } + ChunkFailureReaction::Propagate => { + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::Propagate.as_metric_label(), + format!( + "Propagating pre-write chunk failure ({})", + err.status_code().as_ref() + ), + ); + return Err(err); + } + } + } + + let output = match self + .handle_trace_inserts(requests, ctx) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + { + Ok(output) => output, + Err(err) => { + // Do not retry after dispatch: writes to other peers may already + // have committed even when the aggregate request returns an error. + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::Propagate.as_metric_label(), + format!( + "Propagating chunk write failure ({})", + err.status_code().as_ref() + ), + ); + return Err(err); + } + }; + + Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); + ingest_state.outcome.accepted_spans += span_count; + retry.add_to_aux_data(&mut ingest_state.aux_data); + + Ok(()) + } + + async fn ingest_trace_v1_rows_span_by_span( + &self, + retry: TraceChunkRetry, + ctx: QueryContextRef, + ingest_state: &mut TraceIngestState, + ) -> ServerResult<()> { + for row in retry.rows { + let (span, rows) = row.into_rows(&retry.schema)?; + let mut requests = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: retry.table_name.clone(), + rows: Some(rows), + }], + }; + + if let Err(err) = self.reconcile_trace_column_types(&mut requests, &ctx).await { + if matches!( + Self::classify_trace_prewrite_failure(err.status_code(), err.is_retryable()), + ChunkFailureReaction::Propagate + ) { + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::Propagate.as_metric_label(), + format!( + "Propagating pre-write span failure for {}:{} ({})", + span.trace_id, + span.span_id, + err.status_code().as_ref() + ), + ); + return Err(err); + } + + ingest_state.outcome.rejected_spans += 1; + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + "span_rejected", + format!( + "Rejected span {}:{} ({})", + span.trace_id, + span.span_id, + err.status_code().as_ref() + ), + ); + continue; + } + + let output = match self + .handle_trace_inserts(requests, ctx.clone()) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + { + Ok(output) => output, + Err(err) => { + Self::push_trace_failure_message( + &mut ingest_state.failure_messages, + ChunkFailureReaction::Propagate.as_metric_label(), + format!( + "Propagating span write failure for {}:{} ({})", + span.trace_id, + span.span_id, + err.status_code().as_ref() + ), + ); + return Err(err); + } + }; + + Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); + ingest_state.outcome.accepted_spans += 1; + span.add_to_aux_data(&mut ingest_state.aux_data); + } + + Ok(()) + } + + /// Reconcile and insert one trace request batch. + async fn insert_trace_requests( + &self, + mut requests: RowInsertRequests, + is_trace_v1_model: bool, + ctx: QueryContextRef, + ) -> ServerResult { + if is_trace_v1_model { + self.reconcile_trace_column_types(&mut requests, &ctx) + .await?; + self.handle_trace_inserts(requests, ctx) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } else { + self.handle_log_inserts(requests, ctx) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } + } + + async fn plan_trace_v1_request_schema( + &self, + table_name: &str, + request_schema: &mut TraceRequestSchema, + ctx: &QueryContextRef, + ) -> ServerResult { + let catalog = ctx.current_catalog(); + let schema = ctx.current_schema(); + let table = self + .catalog_manager + .table(catalog, &schema, table_name, None) + .await?; + let table_schema = table.as_ref().map(|table| table.schema()); + let exclusions = request_schema.incompatible_schema_observations(table_schema.as_deref()); + if !exclusions.is_empty() { + return Ok(TraceRequestSchemaPlan::IncompatibleObservations(exclusions)); + } + + Ok(request_schema.resolve_table_schema(table_schema.as_deref())) + } + + async fn prepare_trace_v1_request_schema( + &self, + table_name: &str, + request_schema: &mut TraceRequestSchema, + chunks: &mut Vec, + chunk_states: &mut Vec, + ctx: &QueryContextRef, + ) -> ServerResult<()> { + loop { + // Fixed semconv targets are known before schema planning. Isolate invalid + // values before they can affect conflict resolution or DDL. + let exclusions = Self::prevalidate_trace_v1_fixed_columns(chunks, chunk_states)?; + if !exclusions.is_empty() { + Self::exclude_trace_v1_schema_observations( + request_schema, + chunks, + chunk_states, + &exclusions, + )?; + continue; + } + + if request_schema.columns.is_empty() { + return Ok(()); + } + + let pre_alter = match self + .plan_trace_v1_request_schema(table_name, request_schema, ctx) + .await? + { + TraceRequestSchemaPlan::Ready(pre_alter) => pre_alter, + TraceRequestSchemaPlan::IncompatibleObservations(exclusions) => { + Self::exclude_trace_v1_schema_observations( + request_schema, + chunks, + chunk_states, + &exclusions, + )?; + continue; + } + }; + if pre_alter.requires_ddl() { + let exclusions = Self::prevalidate_trace_v1_chunk_rewrites( + request_schema, + chunks, + chunk_states, + )?; + if !exclusions.is_empty() { + Self::exclude_trace_v1_schema_observations( + request_schema, + chunks, + chunk_states, + &exclusions, + )?; + continue; + } + + self.apply_trace_v1_pre_alter(ctx, table_name, pre_alter) + .await?; + let final_plan = self + .plan_trace_v1_request_schema(table_name, request_schema, ctx) + .await?; + let schema_converged = matches!( + final_plan, + TraceRequestSchemaPlan::Ready(final_plan) if !final_plan.requires_ddl() + ); + if !schema_converged { + for state in chunk_states.iter_mut() { + state.mark_for_reconcile(); + } + return Ok(()); + } + + return Ok(()); + } + + return Ok(()); + } + } + + fn prevalidate_trace_v1_fixed_columns( + chunks: &[TraceChunkRetry], + chunk_states: &[TraceChunkSchemaState], + ) -> ServerResult { + let mut exclusions = HashMap::>::new(); + for (batch_index, (chunk, state)) in chunks.iter().zip(chunk_states).enumerate() { + if *state != TraceChunkSchemaState::Prepared { + continue; + } + + 'rows: for row in &chunk.rows { + for retry_value in &row.values { + let Some(column) = chunk.schema.get(retry_value.column_index) else { + return error::InternalSnafu { + err_msg: format!( + "trace retry column index {} is out of bounds", + retry_value.column_index + ), + } + .fail(); + }; + let Some(target_type) = trace_semconv_fixed_type(&column.column_name) else { + continue; + }; + let Some(request_type) = retry_value + .value + .value_data + .as_ref() + .and_then(trace_value_datatype) + else { + continue; + }; + if request_type != target_type + && coerce_value_data( + &retry_value.value.value_data, + target_type, + request_type, + ) + .is_err() + { + exclusions + .entry(column.column_name.clone()) + .or_default() + .insert(batch_index); + break 'rows; + } + } + } + } + Ok(exclusions) + } + + fn prevalidate_trace_v1_chunk_rewrites( + request_schema: &TraceRequestSchema, + chunks: &[TraceChunkRetry], + chunk_states: &mut [TraceChunkSchemaState], + ) -> ServerResult { + let mut exclusions = HashMap::>::new(); + for (batch_index, (chunk, state)) in chunks.iter().zip(chunk_states).enumerate() { + if *state != TraceChunkSchemaState::Prepared { + continue; + } + let mut request = chunk.to_request()?; + if let Some(column_name) = Self::prepare_trace_v1_chunk_rewrite( + request_schema, + batch_index, + &mut request, + state, + )? { + exclusions + .entry(column_name) + .or_default() + .insert(batch_index); + } + } + Ok(exclusions) + } + + fn prepare_trace_v1_chunk_rewrite( + request_schema: &TraceRequestSchema, + batch_index: usize, + request: &mut RowInsertRequest, + state: &mut TraceChunkSchemaState, + ) -> ServerResult> { + if *state != TraceChunkSchemaState::Prepared { + return Ok(None); + } + + match request_schema.prepare_request_rewrite(batch_index, request) { + Ok(Some(rewrite)) => { + TraceRequestSchema::apply_request_rewrite(request, rewrite); + Ok(None) + } + Ok(None) => Ok(None), + Err(failure) => { + if !matches!( + Self::classify_trace_prewrite_failure( + failure.error.status_code(), + failure.error.is_retryable(), + ), + ChunkFailureReaction::RetryPerSpan + ) { + return Err(failure.error); + } + state.mark_for_reconcile(); + Ok(Some(failure.column_name)) + } + } + } + + fn exclude_trace_v1_schema_observations( + request_schema: &mut TraceRequestSchema, + chunks: &mut Vec, + chunk_states: &mut Vec, + exclusions: &TraceSchemaExclusions, + ) -> ServerResult<()> { + if chunks.len() != chunk_states.len() { + return error::InternalSnafu { + err_msg: format!( + "trace chunk count {} does not match state count {}", + chunks.len(), + chunk_states.len() + ), + } + .fail(); + } + + let batch_indexes = exclusions + .values() + .flatten() + .copied() + .collect::>(); + let previous_chunks = std::mem::take(chunks); + let previous_states = std::mem::take(chunk_states); + for (batch_index, (chunk, mut state)) in + previous_chunks.into_iter().zip(previous_states).enumerate() + { + if !batch_indexes.contains(&batch_index) { + chunks.push(chunk); + chunk_states.push(state); + } else if chunk.rows.len() <= 1 { + state.mark_for_reconcile(); + chunks.push(chunk); + chunk_states.push(state); + } else { + let split_chunks = chunk.into_single_span_chunks()?; + chunk_states.extend(std::iter::repeat_n( + TraceChunkSchemaState::Prepared, + split_chunks.len(), + )); + chunks.extend(split_chunks); + } + } + + *request_schema = TraceRequestSchema::default(); + for (batch_index, (chunk, state)) in chunks.iter().zip(chunk_states.iter()).enumerate() { + if *state == TraceChunkSchemaState::Prepared { + request_schema.observe_retry_chunk(batch_index, chunk)?; + } + } + + Ok(()) + } + + async fn apply_trace_v1_pre_alter( + &self, + ctx: &QueryContextRef, + table_name: &str, + pre_alter: TraceTablePreAlter, + ) -> ServerResult<()> { + let TraceTablePreAlter { + ensure_columns, + modify_float64_columns, + alter_existing, + } = pre_alter; + + if !ensure_columns.is_empty() { + self.inserter + .ensure_trace_table_on_demand( + table_name, + ensure_columns, + alter_existing, + ctx, + &self.statement_executor, + ) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)?; + } + + if !modify_float64_columns.is_empty() { + self.alter_trace_table_columns_to_float64(ctx, table_name, &modify_float64_columns) + .await?; + } + + Ok(()) + } + + fn classify_trace_prewrite_failure( + status: StatusCode, + retryable: bool, + ) -> ChunkFailureReaction { + if retryable { + return ChunkFailureReaction::Propagate; + } + + match status { + StatusCode::InvalidArguments + | StatusCode::InvalidSyntax + | StatusCode::Unsupported + | StatusCode::TableNotFound + | StatusCode::TableColumnNotFound => ChunkFailureReaction::RetryPerSpan, + StatusCode::DatabaseNotFound => ChunkFailureReaction::DiscardChunk, + _ => ChunkFailureReaction::Propagate, + } + } + + fn add_trace_write_cost(outcome: &mut TraceIngestOutcome, cost: usize) { + outcome.write_cost += cost; + } + + fn push_trace_failure_message(messages: &mut Vec, label: &str, message: String) { + OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).inc(); + + if messages.len() < TRACE_FAILURE_MESSAGE_LIMIT { + messages.push(message); + } else if messages.len() == TRACE_FAILURE_MESSAGE_LIMIT { + tracing::debug!( + label, + limit = TRACE_FAILURE_MESSAGE_LIMIT, + "Trace ingest failure message limit reached; suppressing additional failure details" + ); + } + } + + fn finish_trace_failure_message( + accepted_spans: usize, + rejected_spans: usize, + messages: Vec, + ) -> Option { + if rejected_spans == 0 && messages.is_empty() { + return None; + } + + let mut summary = format!( + "Accepted {} spans, rejected {} spans", + accepted_spans, rejected_spans + ); + + if !messages.is_empty() { + summary.push_str(": "); + summary.push_str(&messages.join("; ")); + } + + Some(summary) + } + + /// Widen existing trace table columns to Float64 before request rewrite. + async fn alter_trace_table_columns_to_float64( + &self, + ctx: &QueryContextRef, + table_name: &str, + column_names: &[String], + ) -> ServerResult<()> { + let catalog_name = ctx.current_catalog().to_string(); + let schema_name = ctx.current_schema(); + let alter_expr = AlterTableExpr { + catalog_name: catalog_name.clone(), + schema_name: schema_name.clone(), + table_name: table_name.to_string(), + kind: Some(Kind::ModifyColumnTypes(ModifyColumnTypes { + modify_column_types: column_names + .iter() + .map(|column_name| ModifyColumnType { + column_name: column_name.clone(), + target_type: ColumnDataType::Float64 as i32, + target_type_extension: None, + }) + .collect(), + })), + }; + + if let Err(err) = self + .statement_executor + .alter_table_inner(alter_expr, ctx.clone()) + .await + { + let table = self + .catalog_manager + .table(&catalog_name, &schema_name, table_name, None) + .await + .map_err(servers::error::Error::from)?; + let alter_already_applied = table + .map(|table| { + let table_schema = table.schema(); + column_names.iter().all(|column_name| { + table_schema + .column_schema_by_name(column_name) + .and_then(|table_col| { + ColumnDataTypeWrapper::try_from(table_col.data_type.clone()) + .ok() + .map(|wrapper| wrapper.datatype()) + }) + == Some(ColumnDataType::Float64) + }) + }) + .unwrap_or(false); + + if alter_already_applied { + return Ok(()); + } + + warn!( + table_name, + columns = ?column_names, + error = %err, + "failed to widen trace columns before insert" + ); + + return Err(wrap_trace_alter_failure(err)); + } + + Ok(()) + } + + /// Coerce request column types and values to match the existing table schema + /// for compatible type pairs. Existing table schema wins when present; + /// otherwise the full request batch decides a stable target type. + async fn reconcile_trace_column_types( + &self, + requests: &mut RowInsertRequests, + ctx: &QueryContextRef, + ) -> ServerResult<()> { + let catalog = ctx.current_catalog(); + let schema = ctx.current_schema(); + + for req in &mut requests.inserts { + let table = self + .catalog_manager + .table(catalog, &schema, &req.table_name, None) + .await?; + + let Some(rows) = req.rows.as_mut() else { + continue; + }; + + let table_schema = table.map(|table| table.schema()); + let mut pending_rewrites = Vec::new(); + let mut pending_alter_columns = Vec::new(); + + for (col_idx, col_schema) in rows.schema.iter().enumerate() { + let Some(current_type) = ColumnDataType::try_from(col_schema.datatype).ok() else { + continue; + }; + + let mut observed_types = Vec::new(); + push_observed_trace_type(&mut observed_types, current_type); + + // Scan the full request first so the final type decision is not affected + // by row order inside the batch. + for row in &rows.rows { + let Some(value) = row + .values + .get(col_idx) + .and_then(|value| value.value_data.as_ref()) + else { + continue; + }; + + let Some(value_type) = trace_value_datatype(value) else { + continue; + }; + push_observed_trace_type(&mut observed_types, value_type); + } + + let existing_schema_type = table_schema + .as_ref() + .and_then(|schema| schema.column_schema_by_name(&col_schema.column_name)) + .and_then(|table_col| { + ColumnDataTypeWrapper::try_from(table_col.data_type.clone()) + .ok() + .map(|wrapper| { + let datatype = wrapper.datatype(); + (datatype, ConcreteDataType::from(wrapper)) + }) + }); + let request_concrete_type = ConcreteDataType::from(ColumnDataTypeWrapper::new( + current_type, + col_schema.datatype_extension.clone(), + )); + if let Some((existing_datatype, existing_concrete_type)) = + existing_schema_type.as_ref() + && trace_logical_types_incompatible( + current_type, + &request_concrete_type, + *existing_datatype, + existing_concrete_type, + ) + { + return error::InvalidParameterSnafu { + reason: format!( + "incompatible logical types for trace column '{}' in table '{}'", + col_schema.column_name, req.table_name + ), + } + .fail(); + } + let existing_type = existing_schema_type.map(|(datatype, _)| datatype); + let fixed_type = trace_semconv_fixed_type(&col_schema.column_name); + + if !observed_types + .iter() + .copied() + .any(is_trace_reconcile_candidate_type) + && existing_type + .map(|datatype| !is_trace_reconcile_candidate_type(datatype)) + .unwrap_or(true) + && fixed_type.is_none() + { + continue; + } + + // Decide the final type once per column, then rewrite all affected cells + // together in one row pass below. + let Some(decision) = choose_trace_reconcile_decision( + &col_schema.column_name, + &observed_types, + existing_type, + ) + .map_err(|_| { + enrich_trace_reconcile_error( + &req.table_name, + &col_schema.column_name, + &observed_types, + existing_type, + fixed_type, + ) + })? + else { + continue; + }; + let target_type = decision.target_type(); + + if !decision.requires_alter() + && observed_types + .iter() + .all(|observed| *observed == target_type) + && col_schema.datatype == target_type as i32 + { + continue; + } + + if decision.requires_alter() + && !pending_alter_columns.contains(&col_schema.column_name) + { + pending_alter_columns.push(col_schema.column_name.clone()); + } + + pending_rewrites.push(PendingTraceColumnRewrite { + col_idx, + target_type, + column_name: col_schema.column_name.clone(), + }); + } + + if pending_rewrites.is_empty() { + continue; + } + + let prepared_rewrites = + prepare_trace_column_rewrites(&rows.rows, pending_rewrites, &req.table_name) + .map_err(|failure| failure.error)?; + + if !pending_alter_columns.is_empty() { + self.alter_trace_table_columns_to_float64( + ctx, + &req.table_name, + &pending_alter_columns, + ) + .await?; + } + + prepared_rewrites.apply(rows); + } + + Ok(()) + } +} + +fn trace_logical_types_incompatible( + left_datatype: ColumnDataType, + left_concrete_type: &ConcreteDataType, + right_datatype: ColumnDataType, + right_concrete_type: &ConcreteDataType, +) -> bool { + left_concrete_type != right_concrete_type + && (left_datatype == right_datatype + || !is_trace_reconcile_candidate_type(left_datatype) + || !is_trace_reconcile_candidate_type(right_datatype)) +} + +fn chunk_owned(items: Vec, chunk_size: usize) -> Vec> { + if items.is_empty() { + return Vec::new(); + } + + if chunk_size == 0 { + return vec![items]; + } + + let mut chunks = Vec::with_capacity(items.len().div_ceil(chunk_size)); + let mut iter = items.into_iter(); + while iter.len() > 0 { + chunks.push(iter.by_ref().take(chunk_size).collect()); + } + chunks +} + +/// Preserve the original alter failure status so chunk retry behavior stays correct. +fn wrap_trace_alter_failure(err: E) -> servers::error::Error +where + E: ErrorExt + Send + Sync + 'static, +{ + error::ExecuteGrpcQuerySnafu.into_error(BoxedError::new(err)) +} + +/// Derives `trace.conventions` from the request's resource/scope `schema_url`s. +/// A single distinct non-empty value is concrete; multiple distinct values are +/// `mixed`; none is `unknown`. `schema_url` is row-level in OTLP, so the +/// table-level value is best-effort per the RFC conflict rule. +pub(super) fn trace_conventions(request: &ExportTraceServiceRequest) -> String { + let mut seen: Option<&str> = None; + let mut mixed = false; + + for resource_spans in &request.resource_spans { + let urls = std::iter::once(resource_spans.schema_url.as_str()).chain( + resource_spans + .scope_spans + .iter() + .map(|s| s.schema_url.as_str()), + ); + for url in urls { + if url.is_empty() { + continue; + } + match seen { + None => seen = Some(url), + Some(prev) if prev == url => {} + Some(_) => { + mixed = true; + break; + } + } + } + if mixed { + break; + } + } + + if mixed { + SEMANTIC_VALUE_MIXED.to_string() + } else { + seen.map(str::to_string) + .unwrap_or_else(|| SEMANTIC_VALUE_UNKNOWN.to_string()) + } +} + +#[cfg(test)] +mod tests; diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs new file mode 100644 index 00000000000..949ab1535ac --- /dev/null +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -0,0 +1,1212 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::{HashMap, HashSet}; + +use api::v1::column_data_type_extension::TypeExt; +use api::v1::value::ValueData; +use api::v1::{ + ColumnDataType, ColumnDataTypeExtension, ColumnSchema, JsonTypeExtension, Row, + RowInsertRequest, Rows, SemanticType, Value, +}; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ + ColumnSchema as DatatypesColumnSchema, SchemaBuilder as DatatypesSchemaBuilder, +}; +use opentelemetry_proto::tonic::common::v1::any_value::Value as OtlpValue; +use opentelemetry_proto::tonic::common::v1::{AnyValue, ArrayValue, KeyValue}; +use servers::otlp::trace::SERVICE_NAME_COLUMN; +use servers::otlp::trace::attributes::Attributes; +use servers::otlp::trace::span::{SpanEvents, SpanLinks, TraceSpan}; +use servers::otlp::trace::v1::{TraceBinaryType, TraceRetryColumn}; +use servers::query_handler::TraceIngestOutcome; + +use super::{ + ChunkFailureReaction, Instance, TraceChunkRetry, TraceChunkSchemaState, TraceRequestSchema, + TraceRequestSchemaPlan, TraceSpanMetadata, TraceTablePreAlter, chunk_owned, + wrap_trace_alter_failure, +}; +use crate::metrics::OTLP_TRACES_FAILURE_COUNT; + +#[test] +fn test_chunk_owned() { + let chunks = chunk_owned(vec![1, 2, 3], 2); + assert_eq!(chunks.iter().map(Vec::len).collect::>(), vec![2, 1]); + + let chunks = chunk_owned(vec![1, 2, 3], 0); + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].len(), 3); + + assert!(chunk_owned::(Vec::new(), 0).is_empty()); +} + +#[test] +fn test_classify_trace_prewrite_failure() { + for status in [ + StatusCode::InvalidArguments, + StatusCode::InvalidSyntax, + StatusCode::Unsupported, + StatusCode::TableColumnNotFound, + StatusCode::TableNotFound, + ] { + assert_eq!( + Instance::classify_trace_prewrite_failure(status, false), + ChunkFailureReaction::RetryPerSpan + ); + } + + assert_eq!( + Instance::classify_trace_prewrite_failure(StatusCode::DatabaseNotFound, false), + ChunkFailureReaction::DiscardChunk + ); + + for status in [ + StatusCode::DeadlineExceeded, + StatusCode::Cancelled, + StatusCode::StorageUnavailable, + StatusCode::Internal, + StatusCode::RegionNotReady, + StatusCode::TableUnavailable, + StatusCode::RegionBusy, + StatusCode::RuntimeResourcesExhausted, + StatusCode::Unknown, + ] { + assert_eq!( + Instance::classify_trace_prewrite_failure(status, false), + ChunkFailureReaction::Propagate + ); + } + + // Retry hints take precedence over status-based partial success handling. + assert_eq!( + Instance::classify_trace_prewrite_failure(StatusCode::InvalidArguments, true), + ChunkFailureReaction::Propagate + ); + assert_eq!( + Instance::classify_trace_prewrite_failure(StatusCode::DatabaseNotFound, true), + ChunkFailureReaction::Propagate + ); +} + +#[test] +fn test_add_trace_write_cost() { + let mut outcome = TraceIngestOutcome::default(); + Instance::add_trace_write_cost(&mut outcome, 3); + Instance::add_trace_write_cost(&mut outcome, 5); + assert_eq!(outcome.write_cost, 8); +} + +#[test] +fn test_finish_trace_failure_message() { + let message = Instance::finish_trace_failure_message( + 3, + 2, + vec!["Rejected span trace:span (InvalidArguments)".to_string()], + ) + .unwrap(); + assert!(message.contains("Accepted 3 spans, rejected 2 spans")); + assert!(message.contains("Rejected span trace:span")); + + assert_eq!(Instance::finish_trace_failure_message(2, 0, vec![]), None); +} + +#[test] +fn test_finish_trace_failure_message_without_detail_messages() { + assert_eq!( + Instance::finish_trace_failure_message(0, 2, vec![]), + Some("Accepted 0 spans, rejected 2 spans".to_string()) + ); +} + +#[test] +fn test_push_trace_failure_message_increments_labeled_counter() { + let label = "retry_per_span_counter_test"; + let initial = OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(); + let mut messages = Vec::new(); + + Instance::push_trace_failure_message( + &mut messages, + label, + "Chunk fallback triggered by InvalidArguments".to_string(), + ); + + assert_eq!(messages.len(), 1); + assert_eq!( + OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(), + initial + 1 + ); +} + +#[test] +fn test_push_trace_failure_message_caps_recorded_messages() { + let label = "retry_per_span_limit_test"; + let mut messages = Vec::new(); + + for idx in 0..=4 { + Instance::push_trace_failure_message(&mut messages, label, format!("failure-{idx}")); + } + + assert_eq!(messages.len(), 4); + assert_eq!( + messages, + vec![ + "failure-0".to_string(), + "failure-1".to_string(), + "failure-2".to_string(), + "failure-3".to_string() + ] + ); +} + +#[test] +fn test_wrap_trace_alter_failure_preserves_status_code() { + let err = wrap_trace_alter_failure( + servers::error::TableNotFoundSnafu { + catalog: "greptime".to_string(), + schema: "public".to_string(), + table: "trace_type_missing".to_string(), + } + .build(), + ); + + assert_eq!(err.status_code(), StatusCode::TableNotFound); +} + +fn field_schema(name: &str, datatype: ColumnDataType) -> ColumnSchema { + ColumnSchema { + column_name: name.to_string(), + datatype: datatype as i32, + semantic_type: SemanticType::Field as i32, + ..Default::default() + } +} + +fn tag_schema(name: &str, datatype: ColumnDataType) -> ColumnSchema { + ColumnSchema { + semantic_type: SemanticType::Tag as i32, + ..field_schema(name, datatype) + } +} + +fn json_field_schema(name: &str) -> ColumnSchema { + ColumnSchema { + datatype_extension: Some(ColumnDataTypeExtension { + type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), + }), + ..field_schema(name, ColumnDataType::Binary) + } +} + +fn row(values: Vec>) -> Row { + Row { + values: values + .into_iter() + .map(|value_data| Value { value_data }) + .collect(), + } +} + +fn row_insert_request( + table_name: &str, + schema: Vec, + rows: Vec, +) -> RowInsertRequest { + RowInsertRequest { + table_name: table_name.to_string(), + rows: Some(Rows { schema, rows }), + } +} + +fn trace_attribute(key: &str, value: OtlpValue) -> KeyValue { + KeyValue { + key: key.to_string(), + value: Some(AnyValue { value: Some(value) }), + } +} + +fn trace_span(span_id: &str, span_attributes: Vec) -> TraceSpan { + TraceSpan { + service_name: Some("service".to_string()), + trace_id: "trace-id".to_string(), + span_id: span_id.to_string(), + parent_span_id: None, + resource_attributes: Attributes::from(vec![]), + scope_name: "scope".to_string(), + scope_version: "v1".to_string(), + scope_attributes: Attributes::from(vec![]), + trace_state: String::new(), + span_name: "operation".to_string(), + span_kind: "SPAN_KIND_SERVER".to_string(), + span_status_code: "STATUS_CODE_UNSET".to_string(), + span_status_message: String::new(), + span_attributes: Attributes::from(span_attributes), + span_events: SpanEvents::from(vec![]), + span_links: SpanLinks::from(vec![]), + start_in_nanosecond: 1, + end_in_nanosecond: 2, + } +} + +fn ready_plan(plan: TraceRequestSchemaPlan) -> TraceTablePreAlter { + let TraceRequestSchemaPlan::Ready(plan) = plan else { + panic!("expected a ready trace schema plan"); + }; + plan +} + +fn trace_chunk_retry(request: &RowInsertRequest, span_ids: &[&str]) -> TraceChunkRetry { + TraceChunkRetry::try_new( + request.table_name.clone(), + request.rows.clone().unwrap(), + span_ids + .iter() + .map(|span_id| TraceSpanMetadata { + trace_id: "trace-id".to_string(), + span_id: (*span_id).to_string(), + operation: None, + }) + .collect(), + HashMap::new(), + ) + .unwrap() +} + +fn prepare_and_apply( + request_schema: &TraceRequestSchema, + batch_index: usize, + request: &mut RowInsertRequest, +) { + if let Some(prepared) = request_schema + .prepare_request_rewrite(batch_index, request) + .unwrap() + { + TraceRequestSchema::apply_request_rewrite(request, prepared); + } +} + +#[test] +fn test_trace_chunk_retry_drops_absent_sparse_columns() { + let schema = vec![ + field_schema("parent_span_id", ColumnDataType::String), + field_schema("service_name", ColumnDataType::String), + field_schema("span_attributes.absent", ColumnDataType::Int64), + field_schema("scope_attributes.present", ColumnDataType::Boolean), + ]; + let requests = row_insert_request( + "traces", + schema, + vec![ + row(vec![None, None, None, Some(ValueData::BoolValue(true))]), + row(vec![ + None, + Some(ValueData::StringValue("service".to_string())), + Some(ValueData::I64Value(1)), + None, + ]), + ], + ); + let retry = TraceChunkRetry::try_new( + requests.table_name, + requests.rows.unwrap(), + ["span-1", "span-2"] + .into_iter() + .map(|span_id| TraceSpanMetadata { + trace_id: "trace-id".to_string(), + span_id: span_id.to_string(), + operation: None, + }) + .collect(), + HashMap::from([ + ( + "service_name".to_string(), + TraceRetryColumn { + present_rows: vec![1], + binary_types: Vec::new(), + }, + ), + ( + "span_attributes.absent".to_string(), + TraceRetryColumn { + present_rows: vec![1], + binary_types: Vec::new(), + }, + ), + ( + "scope_attributes.present".to_string(), + TraceRetryColumn { + present_rows: vec![0], + binary_types: Vec::new(), + }, + ), + ]), + ) + .unwrap(); + assert_eq!(retry.rows[0].values.len(), 2); + + let mut rows = retry.rows.into_iter(); + let (_, projected) = rows.next().unwrap().into_rows(&retry.schema).unwrap(); + + assert_eq!( + projected + .schema + .iter() + .map(|column| column.column_name.as_str()) + .collect::>(), + ["parent_span_id", "scope_attributes.present"] + ); + assert_eq!(projected.rows[0].values.len(), 2); + assert_eq!(projected.rows[0].values[0].value_data, None); + assert_eq!( + projected.rows[0].values[1].value_data, + Some(ValueData::BoolValue(true)) + ); +} + +#[test] +fn test_trace_chunk_retry_restores_row_logical_types() { + let column_name = "span_attributes.payload"; + let requests = row_insert_request( + "traces", + vec![field_schema(column_name, ColumnDataType::Binary)], + vec![ + row(vec![Some(ValueData::BinaryValue(vec![1, 2, 3]))]), + row(vec![Some(ValueData::BinaryValue(vec![4, 5, 6]))]), + row(vec![Some(ValueData::StringValue("value".to_string()))]), + ], + ); + let metadata = ["span-1", "span-2", "span-3"] + .into_iter() + .map(|span_id| TraceSpanMetadata { + trace_id: "trace-id".to_string(), + span_id: span_id.to_string(), + operation: None, + }) + .collect(); + let retry = TraceChunkRetry::try_new( + requests.table_name, + requests.rows.unwrap(), + metadata, + HashMap::from([( + column_name.to_string(), + TraceRetryColumn { + present_rows: vec![0, 1, 2], + binary_types: vec![(0, TraceBinaryType::Binary), (1, TraceBinaryType::Json)], + }, + )]), + ) + .unwrap(); + + let projected = retry + .rows + .into_iter() + .map(|row| row.into_rows(&retry.schema).unwrap().1) + .collect::>(); + + assert!(projected[0].schema[0].datatype_extension.is_none()); + assert!(matches!( + projected[1].schema[0] + .datatype_extension + .as_ref() + .and_then(|extension| extension.type_ext.as_ref()), + Some(TypeExt::JsonType(_)) + )); + assert_eq!( + projected[2].schema[0].datatype, + ColumnDataType::String as i32 + ); + assert!(projected[2].schema[0].datatype_extension.is_none()); +} + +#[test] +fn test_trace_chunk_retry_reconstructs_converter_request() { + let spans = vec![ + trace_span( + "span-1", + vec![ + trace_attribute("first", OtlpValue::IntValue(1)), + trace_attribute("payload", OtlpValue::BytesValue(vec![1, 2, 3])), + ], + ), + trace_span( + "span-2", + vec![ + trace_attribute("second", OtlpValue::DoubleValue(2.5)), + trace_attribute( + "payload", + OtlpValue::ArrayValue(ArrayValue { + values: vec![AnyValue { + value: Some(OtlpValue::StringValue("value".to_string())), + }], + }), + ), + ], + ), + ]; + let metadata = spans.iter().map(TraceSpanMetadata::from).collect(); + let (expected_table_data, _) = + servers::otlp::trace::v1::v1_to_main_table_data_with_schema(spans.clone()).unwrap(); + let (expected_schema, mut expected_rows) = expected_table_data.into_schema_and_rows(); + for row in &mut expected_rows { + row.values.resize(expected_schema.len(), Value::default()); + } + let expected = row_insert_request("traces", expected_schema, expected_rows); + let (table_data, batch_schema) = + servers::otlp::trace::v1::v1_to_main_table_data_with_schema(spans).unwrap(); + let (schema, rows) = table_data.into_schema_and_rows(); + let retry = TraceChunkRetry::try_new( + "traces".to_string(), + Rows { schema, rows }, + metadata, + batch_schema.into_retry_columns(), + ) + .unwrap(); + + assert_eq!(retry.to_request().unwrap(), expected); +} + +#[test] +fn test_trace_request_schema_detects_datatype_extension_conflicts() { + let column_name = "span_attributes.payload"; + let binary_schema = field_schema(column_name, ColumnDataType::Binary); + let json_schema = json_field_schema(column_name); + let legacy_json_schema = ColumnSchema { + datatype: ColumnDataType::Json as i32, + ..json_schema.clone() + }; + + let mut request_schema = TraceRequestSchema::default(); + request_schema.observe_trace_column(0, &binary_schema, None); + request_schema.observe_trace_column(1, &json_schema, None); + assert_eq!( + request_schema.incompatible_schema_observations(None), + HashMap::from([(column_name.to_string(), HashSet::from([0, 1]))]) + ); + + let mut equivalent_json_request_schema = TraceRequestSchema::default(); + equivalent_json_request_schema.observe_trace_column(0, &json_schema, None); + equivalent_json_request_schema.observe_trace_column(1, &legacy_json_schema, None); + assert!( + equivalent_json_request_schema + .incompatible_schema_observations(None) + .is_empty() + ); + + let mut incompatible_json_request_schema = TraceRequestSchema::default(); + incompatible_json_request_schema.observe_trace_column(0, &binary_schema, None); + incompatible_json_request_schema.observe_trace_column(1, &legacy_json_schema, None); + assert_eq!( + incompatible_json_request_schema.incompatible_schema_observations(None), + HashMap::from([(column_name.to_string(), HashSet::from([0, 1]))]) + ); + + let existing_json_schema = + DatatypesSchemaBuilder::try_from_columns(vec![DatatypesColumnSchema::new( + column_name, + ConcreteDataType::json_datatype(), + true, + )]) + .unwrap() + .build() + .unwrap(); + let mut binary_request_schema = TraceRequestSchema::default(); + binary_request_schema.observe_trace_column(0, &binary_schema, None); + assert_eq!( + binary_request_schema.incompatible_schema_observations(Some(&existing_json_schema)), + HashMap::from([(column_name.to_string(), HashSet::from([0]))]) + ); + + let mut json_request_schema = TraceRequestSchema::default(); + json_request_schema.observe_trace_column(0, &json_schema, None); + assert!( + json_request_schema + .incompatible_schema_observations(Some(&existing_json_schema)) + .is_empty() + ); +} + +#[test] +fn test_trace_request_schema_isolates_non_candidate_batch() { + let column_name = "span_attributes.payload"; + let string_schema = field_schema(column_name, ColumnDataType::String); + let binary_schema = field_schema(column_name, ColumnDataType::Binary); + let existing_string_schema = + DatatypesSchemaBuilder::try_from_columns(vec![DatatypesColumnSchema::new( + column_name, + ConcreteDataType::string_datatype(), + true, + )]) + .unwrap() + .build() + .unwrap(); + + let mut request_schema = TraceRequestSchema::default(); + request_schema.observe_trace_column(0, &string_schema, Some(ColumnDataType::String)); + request_schema.observe_trace_column(1, &binary_schema, Some(ColumnDataType::Binary)); + + assert_eq!( + request_schema.incompatible_schema_observations(Some(&existing_string_schema)), + HashMap::from([(column_name.to_string(), HashSet::from([1]))]) + ); +} + +#[test] +fn test_trace_request_schema_isolates_non_candidate_new_column_batch() { + let column_name = "span_attributes.payload"; + for datatypes in [ + [ + ColumnDataType::String, + ColumnDataType::Int64, + ColumnDataType::Binary, + ], + [ + ColumnDataType::Binary, + ColumnDataType::Int64, + ColumnDataType::String, + ], + ] { + let binary_batch = datatypes + .iter() + .position(|datatype| *datatype == ColumnDataType::Binary) + .unwrap(); + let mut request_schema = TraceRequestSchema::default(); + for (batch_index, datatype) in datatypes.into_iter().enumerate() { + request_schema.observe_trace_column( + batch_index, + &field_schema(column_name, datatype), + Some(datatype), + ); + } + + assert_eq!( + request_schema.incompatible_schema_observations(None), + HashMap::from([(column_name.to_string(), HashSet::from([binary_batch]),)]) + ); + + request_schema.remove_batches(&HashSet::from([binary_batch])); + let plan = ready_plan(request_schema.resolve_table_schema(None)); + assert_eq!(plan.ensure_columns.len(), 1); + assert_eq!( + plan.ensure_columns[0].datatype, + ColumnDataType::Int64 as i32 + ); + } +} + +#[test] +fn test_trace_request_schema_isolates_unsupported_type_batches() { + let mixed_column = "span_attributes.mixed"; + let unrelated_column = "span_attributes.unrelated"; + let int_schema = field_schema(mixed_column, ColumnDataType::Int64); + let bool_schema = field_schema(mixed_column, ColumnDataType::Boolean); + let string_schema = field_schema(unrelated_column, ColumnDataType::String); + let mut request_schema = TraceRequestSchema::default(); + request_schema.observe_trace_column(0, &int_schema, Some(ColumnDataType::Int64)); + request_schema.observe_trace_column(0, &string_schema, Some(ColumnDataType::String)); + request_schema.observe_trace_column(1, &bool_schema, Some(ColumnDataType::Boolean)); + request_schema.observe_trace_column(1, &string_schema, Some(ColumnDataType::String)); + + let TraceRequestSchemaPlan::IncompatibleObservations(exclusions) = + request_schema.resolve_table_schema(None) + else { + panic!("expected incompatible trace schema batches"); + }; + assert_eq!( + exclusions, + HashMap::from([(mixed_column.to_string(), HashSet::from([0, 1]))]) + ); + + let batch_indexes = exclusions + .values() + .flatten() + .copied() + .collect::>(); + request_schema.remove_batches(&batch_indexes); + let plan = ready_plan(request_schema.resolve_table_schema(None)); + assert!(!plan.requires_ddl()); +} + +#[test] +fn test_trace_request_schema_drops_failed_batch_observations() { + let table_name = "trace_failed_batch"; + let column_name = "span_attributes.value"; + let failed_only_column = "span_attributes.server.port"; + let safe_column = "span_attributes.safe"; + let string_schema = field_schema(column_name, ColumnDataType::String); + let float_schema = field_schema(column_name, ColumnDataType::Float64); + let failed_only_schema = field_schema(failed_only_column, ColumnDataType::String); + let service_schema = tag_schema(SERVICE_NAME_COLUMN, ColumnDataType::String); + let safe_schema = field_schema(safe_column, ColumnDataType::Boolean); + + let mut good_request = row_insert_request( + table_name, + vec![string_schema.clone()], + vec![row(vec![Some(ValueData::StringValue("1.5".to_string()))])], + ); + let mut bad_request = row_insert_request( + table_name, + vec![ + string_schema.clone(), + failed_only_schema.clone(), + service_schema.clone(), + safe_schema.clone(), + ], + vec![row(vec![ + Some(ValueData::StringValue("not-a-number".to_string())), + Some(ValueData::StringValue("invalid-port".to_string())), + Some(ValueData::StringValue("frontend".to_string())), + Some(ValueData::BoolValue(true)), + ])], + ); + let original_bad_request = bad_request.clone(); + + let mut request_schema = TraceRequestSchema::default(); + request_schema.observe_trace_column(0, &string_schema, Some(ColumnDataType::String)); + request_schema.observe_trace_column(0, &float_schema, Some(ColumnDataType::Float64)); + request_schema.observe_trace_column(1, &string_schema, Some(ColumnDataType::String)); + request_schema.observe_trace_column(1, &failed_only_schema, Some(ColumnDataType::String)); + request_schema.observe_trace_column(1, &service_schema, Some(ColumnDataType::String)); + request_schema.observe_trace_column(1, &safe_schema, Some(ColumnDataType::Boolean)); + ready_plan(request_schema.resolve_table_schema(None)); + + let mut chunk_states = vec![ + TraceChunkSchemaState::Prepared, + TraceChunkSchemaState::Prepared, + ]; + assert_eq!( + Instance::prepare_trace_v1_chunk_rewrite( + &request_schema, + 0, + &mut good_request, + &mut chunk_states[0], + ) + .unwrap(), + None + ); + let failed_column = Instance::prepare_trace_v1_chunk_rewrite( + &request_schema, + 1, + &mut bad_request, + &mut chunk_states[1], + ) + .unwrap(); + let exclusions = HashMap::from([(failed_column.unwrap(), HashSet::from([1]))]); + assert_eq!( + exclusions, + HashMap::from([(column_name.to_string(), HashSet::from([1]))]) + ); + assert_eq!(chunk_states[0], TraceChunkSchemaState::Prepared); + assert_eq!(chunk_states[1], TraceChunkSchemaState::ReconcilePerChunk); + + let mut chunks = vec![ + trace_chunk_retry(&good_request, &["good-span"]), + trace_chunk_retry(&bad_request, &["bad-span"]), + ]; + Instance::exclude_trace_v1_schema_observations( + &mut request_schema, + &mut chunks, + &mut chunk_states, + &exclusions, + ) + .unwrap(); + assert!( + !request_schema + .column_indexes + .contains_key(failed_only_column) + ); + assert!( + !request_schema + .column_indexes + .contains_key(SERVICE_NAME_COLUMN) + ); + assert!(!request_schema.column_indexes.contains_key(safe_column)); + let plan = ready_plan(request_schema.resolve_table_schema(None)); + assert_eq!(plan.ensure_columns.len(), 1); + assert_eq!(plan.ensure_columns[0].column_name, column_name); + assert_eq!( + Instance::prepare_trace_v1_chunk_rewrite( + &request_schema, + 1, + &mut bad_request, + &mut chunk_states[1], + ) + .unwrap(), + None + ); + + let good_rows = good_request.rows.as_ref().unwrap(); + assert_eq!(good_rows.schema[0].datatype, ColumnDataType::Float64 as i32); + assert_eq!( + good_rows.rows[0].values[0].value_data, + Some(ValueData::F64Value(1.5)) + ); + assert_eq!(bad_request, original_bad_request); +} + +#[test] +fn test_trace_request_schema_splits_failed_chunk_before_exclusion() { + let table_name = "trace_failed_chunk"; + let value_column = "span_attributes.value"; + let port_column = "span_attributes.server.port"; + let boolean_schema = field_schema(value_column, ColumnDataType::Boolean); + let int_schema = field_schema(value_column, ColumnDataType::Int64); + let string_schema = field_schema(value_column, ColumnDataType::String); + let port_string_schema = field_schema(port_column, ColumnDataType::String); + let port_int_schema = field_schema(port_column, ColumnDataType::Int64); + let requests = [ + row_insert_request( + table_name, + vec![boolean_schema, port_string_schema], + vec![ + row(vec![ + Some(ValueData::BoolValue(true)), + Some(ValueData::StringValue("1".to_string())), + ]), + row(vec![ + Some(ValueData::BoolValue(false)), + Some(ValueData::StringValue("bad".to_string())), + ]), + ], + ), + row_insert_request( + table_name, + vec![int_schema, port_int_schema.clone()], + vec![row(vec![ + Some(ValueData::I64Value(2)), + Some(ValueData::I64Value(2)), + ])], + ), + row_insert_request( + table_name, + vec![string_schema, port_int_schema], + vec![row(vec![ + Some(ValueData::StringValue("3".to_string())), + Some(ValueData::I64Value(3)), + ])], + ), + ]; + let mut chunks = vec![ + trace_chunk_retry(&requests[0], &["good-span", "bad-span"]), + trace_chunk_retry(&requests[1], &["int-span"]), + trace_chunk_retry(&requests[2], &["string-span"]), + ]; + let mut chunk_states = vec![TraceChunkSchemaState::Prepared; chunks.len()]; + let mut request_schema = TraceRequestSchema::default(); + for (batch_index, chunk) in chunks.iter().enumerate() { + request_schema + .observe_retry_chunk(batch_index, chunk) + .unwrap(); + } + + ready_plan(request_schema.resolve_table_schema(None)); + let value_index = request_schema.column_indexes[value_column]; + assert_eq!( + request_schema.columns[value_index].target_type, + Some(ColumnDataType::String) + ); + + let exclusions = + Instance::prevalidate_trace_v1_chunk_rewrites(&request_schema, &chunks, &mut chunk_states) + .unwrap(); + assert_eq!( + exclusions, + HashMap::from([(port_column.to_string(), HashSet::from([0]))]) + ); + Instance::exclude_trace_v1_schema_observations( + &mut request_schema, + &mut chunks, + &mut chunk_states, + &exclusions, + ) + .unwrap(); + assert_eq!(chunks.len(), 4); + assert!(chunks.iter().all(|chunk| chunk.rows.len() == 1)); + + ready_plan(request_schema.resolve_table_schema(None)); + let exclusions = + Instance::prevalidate_trace_v1_chunk_rewrites(&request_schema, &chunks, &mut chunk_states) + .unwrap(); + assert_eq!( + exclusions, + HashMap::from([(port_column.to_string(), HashSet::from([1]))]) + ); + Instance::exclude_trace_v1_schema_observations( + &mut request_schema, + &mut chunks, + &mut chunk_states, + &exclusions, + ) + .unwrap(); + + ready_plan(request_schema.resolve_table_schema(None)); + let value_index = request_schema.column_indexes[value_column]; + assert_eq!( + request_schema.columns[value_index].target_type, + Some(ColumnDataType::String) + ); + assert_eq!(chunk_states[0], TraceChunkSchemaState::Prepared); + assert_eq!(chunk_states[1], TraceChunkSchemaState::ReconcilePerChunk); +} + +#[test] +fn test_trace_request_schema_prevalidates_fixed_columns_before_conflicts() { + let table_name = "trace_fixed_column_conflict"; + let conflict_column = "span_attributes.conflict"; + let value_column = "span_attributes.value"; + let port_column = "span_attributes.server.port"; + let requests = [ + row_insert_request( + table_name, + vec![ + field_schema(conflict_column, ColumnDataType::String), + field_schema(value_column, ColumnDataType::Boolean), + ], + vec![row(vec![ + Some(ValueData::StringValue("text".to_string())), + Some(ValueData::BoolValue(true)), + ])], + ), + row_insert_request( + table_name, + vec![ + field_schema(conflict_column, ColumnDataType::Binary), + field_schema(port_column, ColumnDataType::String), + ], + vec![row(vec![ + Some(ValueData::BinaryValue(vec![1])), + Some(ValueData::StringValue("bad".to_string())), + ])], + ), + row_insert_request( + table_name, + vec![field_schema(value_column, ColumnDataType::Int64)], + vec![row(vec![Some(ValueData::I64Value(2))])], + ), + row_insert_request( + table_name, + vec![field_schema(value_column, ColumnDataType::String)], + vec![row(vec![Some(ValueData::StringValue("3".to_string()))])], + ), + ]; + let mut chunks = requests + .iter() + .map(|request| trace_chunk_retry(request, &["span"])) + .collect::>(); + let mut chunk_states = vec![TraceChunkSchemaState::Prepared; chunks.len()]; + let mut request_schema = TraceRequestSchema::default(); + for (batch_index, chunk) in chunks.iter().enumerate() { + request_schema + .observe_retry_chunk(batch_index, chunk) + .unwrap(); + } + + assert_eq!( + request_schema.incompatible_schema_observations(None), + HashMap::from([(conflict_column.to_string(), HashSet::from([1]),)]) + ); + let exclusions = Instance::prevalidate_trace_v1_fixed_columns(&chunks, &chunk_states).unwrap(); + assert_eq!( + exclusions, + HashMap::from([(port_column.to_string(), HashSet::from([1]))]) + ); + Instance::exclude_trace_v1_schema_observations( + &mut request_schema, + &mut chunks, + &mut chunk_states, + &exclusions, + ) + .unwrap(); + + assert!( + request_schema + .incompatible_schema_observations(None) + .is_empty() + ); + ready_plan(request_schema.resolve_table_schema(None)); + let value_index = request_schema.column_indexes[value_column]; + assert_eq!( + request_schema.columns[value_index].target_type, + Some(ColumnDataType::String) + ); + assert_eq!(chunk_states[0], TraceChunkSchemaState::Prepared); + assert_eq!(chunk_states[1], TraceChunkSchemaState::ReconcilePerChunk); +} + +#[test] +fn test_trace_request_schema_uses_global_columns_and_types() { + let table_name = "trace_global_schema"; + let attr_num = "span_attributes.attr_num"; + let attr_later = "span_attributes.attr_later"; + let mut first_chunk = row_insert_request( + table_name, + vec![field_schema(attr_num, ColumnDataType::Int64)], + vec![row(vec![Some(ValueData::I64Value(1))])], + ); + let mut second_chunk = row_insert_request( + table_name, + vec![ + field_schema(attr_num, ColumnDataType::Float64), + field_schema(attr_later, ColumnDataType::Boolean), + ], + vec![row(vec![ + Some(ValueData::F64Value(2.5)), + Some(ValueData::BoolValue(true)), + ])], + ); + + let mut request_schema = TraceRequestSchema::default(); + let attr_num_i64 = field_schema(attr_num, ColumnDataType::Int64); + let attr_num_f64 = field_schema(attr_num, ColumnDataType::Float64); + let attr_later_bool = field_schema(attr_later, ColumnDataType::Boolean); + request_schema.observe_trace_column(0, &attr_num_i64, Some(ColumnDataType::Int64)); + request_schema.observe_trace_column(1, &attr_num_f64, Some(ColumnDataType::Float64)); + request_schema.observe_trace_column(1, &attr_later_bool, Some(ColumnDataType::Boolean)); + + let existing_schema = + DatatypesSchemaBuilder::try_from_columns(vec![DatatypesColumnSchema::new( + attr_num, + ConcreteDataType::int64_datatype(), + true, + )]) + .unwrap() + .build() + .unwrap(); + let pre_alter = ready_plan(request_schema.resolve_table_schema(Some(&existing_schema))); + assert_eq!(pre_alter.modify_float64_columns, vec![attr_num.to_string()]); + assert_eq!(pre_alter.ensure_columns.len(), 1); + assert_eq!(pre_alter.ensure_columns[0].column_name, attr_later); + + prepare_and_apply(&request_schema, 0, &mut first_chunk); + prepare_and_apply(&request_schema, 1, &mut second_chunk); + + let rows = first_chunk.rows.as_ref().unwrap(); + assert_eq!(rows.schema.len(), 1); + assert_eq!(rows.schema[0].column_name, attr_num); + assert_eq!(rows.schema[0].datatype, ColumnDataType::Float64 as i32); + assert_eq!( + rows.rows[0].values[0].value_data, + Some(ValueData::F64Value(1.0)) + ); + + let rows = second_chunk.rows.as_ref().unwrap(); + assert_eq!(rows.schema.len(), 2); + assert_eq!(rows.schema[0].column_name, attr_num); + assert_eq!(rows.schema[0].datatype, ColumnDataType::Float64 as i32); + assert_eq!(rows.schema[1].column_name, attr_later); + assert_eq!(rows.schema[1].datatype, ColumnDataType::Boolean as i32); + assert_eq!( + rows.rows[0].values[0].value_data, + Some(ValueData::F64Value(2.5)) + ); + assert_eq!( + rows.rows[0].values[1].value_data, + Some(ValueData::BoolValue(true)) + ); +} + +#[test] +fn test_trace_request_schema_resolves_three_compatible_types() { + let column_name = "span_attributes.value"; + let mut request_schema = TraceRequestSchema::default(); + for (batch_index, datatype) in [ + ColumnDataType::String, + ColumnDataType::Int64, + ColumnDataType::Float64, + ] + .into_iter() + .enumerate() + { + let schema = field_schema(column_name, datatype); + request_schema.observe_trace_column(batch_index, &schema, Some(datatype)); + } + + let plan = ready_plan(request_schema.resolve_table_schema(None)); + assert_eq!(plan.ensure_columns.len(), 1); + assert_eq!( + plan.ensure_columns[0].datatype, + ColumnDataType::Float64 as i32 + ); +} + +#[test] +fn test_trace_request_schema_creates_table_without_padding_batches() { + let table_name = "trace_create_schema"; + let attr_num = "span_attributes.attr_num"; + let attr_later = "span_attributes.attr_later"; + let mut first_chunk = row_insert_request( + table_name, + vec![field_schema(attr_num, ColumnDataType::Int64)], + vec![row(vec![Some(ValueData::I64Value(1))])], + ); + let mut second_chunk = row_insert_request( + table_name, + vec![ + field_schema(attr_num, ColumnDataType::Float64), + field_schema(attr_later, ColumnDataType::Boolean), + ], + vec![row(vec![ + Some(ValueData::F64Value(2.5)), + Some(ValueData::BoolValue(true)), + ])], + ); + + let mut request_schema = TraceRequestSchema::default(); + let attr_num_i64 = field_schema(attr_num, ColumnDataType::Int64); + let attr_num_f64 = field_schema(attr_num, ColumnDataType::Float64); + let attr_later_bool = field_schema(attr_later, ColumnDataType::Boolean); + request_schema.observe_trace_column(0, &attr_num_i64, Some(ColumnDataType::Int64)); + request_schema.observe_trace_column(1, &attr_num_f64, Some(ColumnDataType::Float64)); + request_schema.observe_trace_column(1, &attr_later_bool, Some(ColumnDataType::Boolean)); + + let pre_alter = ready_plan(request_schema.resolve_table_schema(None)); + assert!(pre_alter.modify_float64_columns.is_empty()); + assert_eq!(pre_alter.ensure_columns.len(), 2); + assert_eq!(pre_alter.ensure_columns[0].column_name, attr_num); + assert_eq!( + pre_alter.ensure_columns[0].datatype, + ColumnDataType::Float64 as i32 + ); + assert_eq!(pre_alter.ensure_columns[1].column_name, attr_later); + + prepare_and_apply(&request_schema, 0, &mut first_chunk); + prepare_and_apply(&request_schema, 1, &mut second_chunk); + + let rows = first_chunk.rows.as_ref().unwrap(); + assert_eq!(rows.schema.len(), 1); + assert_eq!(rows.schema[0].column_name, attr_num); + assert_eq!(rows.schema[0].datatype, ColumnDataType::Float64 as i32); + assert_eq!( + rows.rows[0].values[0].value_data, + Some(ValueData::F64Value(1.0)) + ); + assert_eq!(rows.rows[0].values.len(), 1); + + let rows = second_chunk.rows.as_ref().unwrap(); + assert_eq!(rows.schema.len(), 2); + assert_eq!(rows.schema[0].column_name, attr_num); + assert_eq!(rows.schema[1].column_name, attr_later); +} + +#[test] +fn test_trace_request_schema_re_resolves_after_concurrent_create() { + let table_name = "trace_concurrent_create"; + let column_name = "span_attributes.value"; + let string_schema = field_schema(column_name, ColumnDataType::String); + let mut requests = row_insert_request( + table_name, + vec![string_schema.clone()], + vec![row(vec![Some(ValueData::StringValue("42".to_string()))])], + ); + + let mut request_schema = TraceRequestSchema::default(); + request_schema.observe_trace_column(0, &string_schema, Some(ColumnDataType::String)); + let initial_plan = ready_plan(request_schema.resolve_table_schema(None)); + assert!(!initial_plan.ensure_columns.is_empty()); + assert!(!initial_plan.alter_existing); + let initial_targets = request_schema.resolved_target_types(); + + let concurrent_schema = + DatatypesSchemaBuilder::try_from_columns(vec![DatatypesColumnSchema::new( + column_name, + ConcreteDataType::int64_datatype(), + true, + )]) + .unwrap() + .build() + .unwrap(); + let final_plan = ready_plan(request_schema.resolve_table_schema(Some(&concurrent_schema))); + assert!(!final_plan.requires_ddl()); + assert!(final_plan.alter_existing); + assert_ne!(initial_targets, request_schema.resolved_target_types()); + + prepare_and_apply(&request_schema, 0, &mut requests); + let rows = requests.rows.as_ref().unwrap(); + assert_eq!(rows.schema[0].datatype, ColumnDataType::Int64 as i32); + assert_eq!( + rows.rows[0].values[0].value_data, + Some(ValueData::I64Value(42)) + ); +} + +#[test] +fn test_trace_request_schema_keeps_request_intact_on_coercion_failure() { + let table_name = "trace_atomic_rewrite"; + let column_name = "span_attributes.value"; + let string_schema = field_schema(column_name, ColumnDataType::String); + let float_schema = field_schema(column_name, ColumnDataType::Float64); + let requests = row_insert_request( + table_name, + vec![string_schema.clone()], + vec![row(vec![Some(ValueData::StringValue( + "not-a-number".to_string(), + ))])], + ); + let original = requests.clone(); + + let mut request_schema = TraceRequestSchema::default(); + request_schema.observe_trace_column(0, &string_schema, Some(ColumnDataType::String)); + request_schema.observe_trace_column(0, &float_schema, Some(ColumnDataType::Float64)); + ready_plan(request_schema.resolve_table_schema(None)); + + assert!( + request_schema + .prepare_request_rewrite(0, &requests) + .is_err() + ); + assert_eq!(requests, original); +} + +use opentelemetry_proto::tonic::trace::v1::{ResourceSpans, ScopeSpans}; + +use super::{ExportTraceServiceRequest, trace_conventions}; + +fn resource_spans(resource_url: &str, scope_urls: &[&str]) -> ResourceSpans { + ResourceSpans { + schema_url: resource_url.to_string(), + scope_spans: scope_urls + .iter() + .map(|u| ScopeSpans { + schema_url: u.to_string(), + ..Default::default() + }) + .collect(), + ..Default::default() + } +} + +#[test] +fn test_trace_conventions() { + let unknown = ExportTraceServiceRequest::default(); + assert_eq!(trace_conventions(&unknown), "unknown"); + + let url = "https://opentelemetry.io/schemas/1.27.0"; + let single = ExportTraceServiceRequest { + resource_spans: vec![resource_spans("", &[url, url])], + }; + assert_eq!(trace_conventions(&single), url); + + let resource_level = ExportTraceServiceRequest { + resource_spans: vec![resource_spans(url, &[""])], + }; + assert_eq!(trace_conventions(&resource_level), url); + + let conflicting = ExportTraceServiceRequest { + resource_spans: vec![resource_spans( + "", + &[url, "https://opentelemetry.io/schemas/1.30.0"], + )], + }; + assert_eq!(trace_conventions(&conflicting), "mixed"); +} diff --git a/src/frontend/src/instance/otlp/trace_types.rs b/src/frontend/src/instance/otlp/trace_types.rs index e46e16fbd5c..0e2df25053e 100644 --- a/src/frontend/src/instance/otlp/trace_types.rs +++ b/src/frontend/src/instance/otlp/trace_types.rs @@ -12,7 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::{ColumnDataType, Row}; +use api::v1::value::ValueData; +use api::v1::{ColumnDataType, Row, Rows}; use servers::error::{self, Result as ServerResult}; use servers::otlp::coerce::{ coerce_value_data, is_supported_trace_coercion, resolve_new_trace_column_type, @@ -42,12 +43,54 @@ impl TraceReconcileDecision { } } +/// Describes a column rewrite before its row values have been validated. +#[derive(Debug)] pub(super) struct PendingTraceColumnRewrite { pub(super) col_idx: usize, pub(super) target_type: ColumnDataType, pub(super) column_name: String, } +/// Holds the schema and value rewrites prepared for atomic application. +#[derive(Debug)] +pub(super) struct PreparedTraceColumnRewrites { + columns: Vec, + values: Vec, +} + +/// Reports the column whose trace value could not be rewritten. +#[derive(Debug)] +pub(super) struct TraceColumnRewriteError { + pub(super) error: servers::error::Error, + pub(super) column_name: String, +} + +/// Updates one request column to its reconciled datatype. +#[derive(Debug)] +struct PreparedTraceColumnRewrite { + col_idx: usize, + target_type: ColumnDataType, +} + +/// Replaces one trace value with its precomputed coerced value. +#[derive(Debug)] +struct PreparedTraceValueRewrite { + row_idx: usize, + col_idx: usize, + value_data: Option, +} + +impl PreparedTraceColumnRewrites { + pub(super) fn apply(self, rows: &mut Rows) { + for column in self.columns { + rows.schema[column.col_idx].datatype = column.target_type as i32; + } + for value in self.values { + rows.rows[value.row_idx].values[value.col_idx].value_data = value.value_data; + } + } +} + /// Picks the reconciliation action for one trace column. /// /// Existing table schema is authoritative unless the only incompatible case is @@ -129,14 +172,22 @@ fn choose_fixed_trace_reconcile_decision( .fail() } -/// Validate all pending trace column rewrites before any schema mutation happens. -pub(super) fn validate_trace_column_rewrites( +/// Prepares an atomic rewrite plan without mutating the input rows. +/// +/// For each pending column rewrite, this precomputes every required value +/// coercion. Missing, null, and already-correct values are skipped. If any +/// coercion fails, it returns the failing column and leaves all rows unchanged. +/// +/// Target types must already have been selected by `TraceRequestSchema`. +/// Call [`PreparedTraceColumnRewrites::apply`] to update the schema and values. +pub(super) fn prepare_trace_column_rewrites( rows: &[Row], - pending_rewrites: &[PendingTraceColumnRewrite], + pending_rewrites: Vec, table_name: &str, -) -> ServerResult<()> { - for row in rows { - for pending_rewrite in pending_rewrites { +) -> Result { + let mut values = Vec::new(); + for (row_idx, row) in rows.iter().enumerate() { + for pending_rewrite in &pending_rewrites { let Some(value) = row.values.get(pending_rewrite.col_idx) else { continue; }; @@ -148,23 +199,39 @@ pub(super) fn validate_trace_column_rewrites( continue; } - coerce_value_data(&value.value_data, pending_rewrite.target_type, request_type) - .map_err(|_| { - error::InvalidParameterSnafu { - reason: format!( - "failed to coerce trace column '{}' in table '{}' from {:?} to {:?}", - pending_rewrite.column_name, - table_name, - request_type, - pending_rewrite.target_type - ), - } - .build() - })?; + let value_data = + coerce_value_data(&value.value_data, pending_rewrite.target_type, request_type) + .map_err(|_| { + TraceColumnRewriteError { + error: error::InvalidParameterSnafu { + reason: format!( + "failed to coerce trace column '{}' in table '{}' from {:?} to {:?}", + pending_rewrite.column_name, + table_name, + request_type, + pending_rewrite.target_type + ), + } + .build(), + column_name: pending_rewrite.column_name.clone(), + } + })?; + values.push(PreparedTraceValueRewrite { + row_idx, + col_idx: pending_rewrite.col_idx, + value_data, + }); } } - Ok(()) + let columns = pending_rewrites + .into_iter() + .map(|rewrite| PreparedTraceColumnRewrite { + col_idx: rewrite.col_idx, + target_type: rewrite.target_type, + }) + .collect(); + Ok(PreparedTraceColumnRewrites { columns, values }) } pub(super) fn enrich_trace_reconcile_error( @@ -228,14 +295,14 @@ pub(super) fn push_observed_trace_type( #[cfg(test)] mod tests { use api::v1::value::ValueData; - use api::v1::{ColumnDataType, Row, Value}; + use api::v1::{ColumnDataType, ColumnSchema, Row, Rows, Value}; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use super::{ PendingTraceColumnRewrite, TraceReconcileDecision, choose_trace_reconcile_decision, - enrich_trace_reconcile_error, is_trace_reconcile_candidate_type, push_observed_trace_type, - validate_trace_column_rewrites, + enrich_trace_reconcile_error, is_trace_reconcile_candidate_type, + prepare_trace_column_rewrites, push_observed_trace_type, }; #[test] @@ -375,7 +442,7 @@ mod tests { } #[test] - fn test_validate_trace_column_rewrites_rejects_invalid_string_parse() { + fn test_prepare_trace_column_rewrites_rejects_invalid_string_parse() { let rows = vec![Row { values: vec![Value { value_data: Some(ValueData::StringValue("not_a_number".to_string())), @@ -387,29 +454,49 @@ mod tests { column_name: "span_attributes.attr_int".to_string(), }]; - let err = validate_trace_column_rewrites(&rows, &pending_rewrites, "trace_type_atomicity") + let err = prepare_trace_column_rewrites(&rows, pending_rewrites, "trace_type_atomicity") .unwrap_err(); - assert_eq!(err.status_code(), StatusCode::InvalidArguments); + assert_eq!(err.error.status_code(), StatusCode::InvalidArguments); + assert_eq!(err.column_name, "span_attributes.attr_int"); } #[test] - fn test_validate_trace_column_rewrites_whitelisted_values_validate_against_fixed_type() { - let rows = vec![Row { - values: vec![Value { - value_data: Some(ValueData::StringValue("503".to_string())), + fn test_prepare_trace_column_rewrites_applies_prepared_values() { + let mut rows = Rows { + schema: vec![ColumnSchema { + datatype: ColumnDataType::String as i32, + ..Default::default() }], - }]; + rows: vec![Row { + values: vec![Value { + value_data: Some(ValueData::StringValue("503".to_string())), + }], + }], + }; let pending_rewrites = vec![PendingTraceColumnRewrite { col_idx: 0, target_type: ColumnDataType::Int64, column_name: "span_attributes.http.response.status_code".to_string(), }]; - validate_trace_column_rewrites(&rows, &pending_rewrites, "trace_type_atomicity").unwrap(); + let prepared = + prepare_trace_column_rewrites(&rows.rows, pending_rewrites, "trace_type_atomicity") + .unwrap(); + assert_eq!( + rows.rows[0].values[0].value_data, + Some(ValueData::StringValue("503".to_string())) + ); + + prepared.apply(&mut rows); + assert_eq!(rows.schema[0].datatype, ColumnDataType::Int64 as i32); + assert_eq!( + rows.rows[0].values[0].value_data, + Some(ValueData::I64Value(503)) + ); } #[test] - fn test_validate_trace_column_rewrites_whitelisted_boolean_rejects_invalid_string_parse() { + fn test_prepare_trace_column_rewrites_boolean_rejects_invalid_string_parse() { let rows = vec![Row { values: vec![Value { value_data: Some(ValueData::StringValue("not_a_bool".to_string())), @@ -421,9 +508,13 @@ mod tests { column_name: "span_attributes.messaging.destination.temporary".to_string(), }]; - let err = validate_trace_column_rewrites(&rows, &pending_rewrites, "trace_type_atomicity") + let err = prepare_trace_column_rewrites(&rows, pending_rewrites, "trace_type_atomicity") .unwrap_err(); - assert_eq!(err.status_code(), StatusCode::InvalidArguments); + assert_eq!(err.error.status_code(), StatusCode::InvalidArguments); + assert_eq!( + err.column_name, + "span_attributes.messaging.destination.temporary" + ); } #[test] diff --git a/src/log-query/src/log_query.rs b/src/log-query/src/log_query.rs index 2955b7bdedf..78e622200d4 100644 --- a/src/log-query/src/log_query.rs +++ b/src/log-query/src/log_query.rs @@ -214,9 +214,9 @@ impl TimeFilter { (Some(start), Some(end), _) => { // Both 'start' and 'end' are provided let (start, _) = Self::parse_datetime(start)?; - let (end, _) = Self::parse_datetime(end)?; + let (end, inferred_end) = Self::parse_datetime(end)?; start_dt = Some(start); - end_dt = Some(end); + end_dt = Some(inferred_end.unwrap_or(end)); } (Some(start), None, Some(span)) => { let (start, _) = Self::parse_datetime(start)?; @@ -276,41 +276,39 @@ impl TimeFilter { if let Ok(dt) = DateTime::parse_from_rfc3339(s) { Ok((dt.with_timezone(&Utc), None)) } else { - let formats = ["%Y-%m-%d", "%Y-%m", "%Y"]; - for format in &formats { - if let Ok(naive_date) = NaiveDate::parse_from_str(s, format) { - let start = Utc.from_utc_datetime( - &naive_date.and_time(NaiveTime::from_hms_opt(0, 0, 0).unwrap()), - ); - let end = match *format { - "%Y-%m-%d" => start + Duration::days(1), - "%Y-%m" => { - let next_month = if naive_date.month() == 12 { - NaiveDate::from_ymd_opt(naive_date.year() + 1, 1, 1).unwrap() - } else { - NaiveDate::from_ymd_opt( - naive_date.year(), - naive_date.month() + 1, - 1, - ) - .unwrap() - }; - Utc.from_utc_datetime(&next_month.and_hms_opt(0, 0, 0).unwrap()) - } - "%Y" => { - let next_year = - NaiveDate::from_ymd_opt(naive_date.year() + 1, 1, 1).unwrap(); - Utc.from_utc_datetime(&next_year.and_hms_opt(0, 0, 0).unwrap()) - } - _ => unreachable!(), - }; - return Ok((start, Some(end))); + let (naive_date, format) = if let Ok(date) = NaiveDate::parse_from_str(s, "%Y-%m-%d") { + (date, "%Y-%m-%d") + } else if let Ok(date) = NaiveDate::parse_from_str(&format!("{s}-01"), "%Y-%m-%d") { + (date, "%Y-%m") + } else if let Ok(date) = NaiveDate::parse_from_str(&format!("{s}-01-01"), "%Y-%m-%d") { + (date, "%Y") + } else { + return Err(InvalidDateFormatSnafu { + input: s.to_string(), } - } - Err(InvalidDateFormatSnafu { - input: s.to_string(), - } - .build()) + .build()); + }; + + let start = Utc + .from_utc_datetime(&naive_date.and_time(NaiveTime::from_hms_opt(0, 0, 0).unwrap())); + let end = match format { + "%Y-%m-%d" => start + Duration::days(1), + "%Y-%m" => { + let next_month = if naive_date.month() == 12 { + NaiveDate::from_ymd_opt(naive_date.year() + 1, 1, 1).unwrap() + } else { + NaiveDate::from_ymd_opt(naive_date.year(), naive_date.month() + 1, 1) + .unwrap() + }; + Utc.from_utc_datetime(&next_month.and_hms_opt(0, 0, 0).unwrap()) + } + "%Y" => { + let next_year = NaiveDate::from_ymd_opt(naive_date.year() + 1, 1, 1).unwrap(); + Utc.from_utc_datetime(&next_year.and_hms_opt(0, 0, 0).unwrap()) + } + _ => unreachable!(), + }; + Ok((start, Some(end))) } } @@ -545,4 +543,44 @@ mod tests { let result = tf.canonicalize(); assert!(matches!(result, Err(Error::EndBeforeStart { .. }))); } + + #[test] + fn test_canonicalize_date_end() { + let mut date_end = TimeFilter { + start: Some("2023-10-01".to_string()), + end: Some("2023-10-02".to_string()), + span: None, + }; + date_end.canonicalize().unwrap(); + assert_eq!(date_end.end.as_deref(), Some("2023-10-03T00:00:00+00:00")); + + let mut timestamp_end = TimeFilter { + start: Some("2023-10-01".to_string()), + end: Some("2023-10-02T12:34:56Z".to_string()), + span: None, + }; + timestamp_end.canonicalize().unwrap(); + assert_eq!( + timestamp_end.end.as_deref(), + Some("2023-10-02T12:34:56+00:00") + ); + + let mut month = TimeFilter { + start: Some("2023-10".to_string()), + end: Some("2023-10".to_string()), + span: None, + }; + month.canonicalize().unwrap(); + assert_eq!(month.start.as_deref(), Some("2023-10-01T00:00:00+00:00")); + assert_eq!(month.end.as_deref(), Some("2023-11-01T00:00:00+00:00")); + + let mut year = TimeFilter { + start: Some("2023".to_string()), + end: Some("2023".to_string()), + span: None, + }; + year.canonicalize().unwrap(); + assert_eq!(year.start.as_deref(), Some("2023-01-01T00:00:00+00:00")); + assert_eq!(year.end.as_deref(), Some("2024-01-01T00:00:00+00:00")); + } } diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index 01484a7d902..591a3c88ea5 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -511,6 +511,10 @@ impl LogStore for KafkaLogStore { Ok(()) } + async fn obsolete_all(&self, _provider: &Provider, _region_id: RegionId) -> Result<()> { + Ok(()) + } + /// Returns the highest entry id of the specified topic in remote WAL. fn latest_entry_id(&self, provider: &Provider) -> Result { let provider = provider diff --git a/src/log-store/src/kafka/util/record.rs b/src/log-store/src/kafka/util/record.rs index 720f989139d..aefb098e0db 100644 --- a/src/log-store/src/kafka/util/record.rs +++ b/src/log-store/src/kafka/util/record.rs @@ -227,13 +227,17 @@ pub fn remaining_entries( /// - Emits a [RecordType::Full] type record immediately. /// /// For type of [Entry::MultiplePart] Entry: -/// - Emits a complete or incomplete [Entry] while the next same [RegionId] record arrives. +/// - Emits a complete [Entry] immediately when a [RecordType::Last] record arrives with +/// buffered parts. +/// - Emits an incomplete [Entry] when a new [RecordType::First] record replaces +/// buffered incomplete records for the same [RegionId]. /// -/// **Incomplete Entry:** -/// If the records arrive in the following order, it emits **the incomplete [Entry]** when the next record arrives. -/// - **[RecordType::First], [RecordType::Middle]**, [RecordType::First] -/// - **[RecordType::Middle]**, [RecordType::First] -/// - **[RecordType::Last]** +/// **Discarded records:** +/// - A standalone [RecordType::First] record is discarded when another [RecordType::First] +/// record with the same payload for the same [RegionId] arrives. +/// - A [RecordType::Last] record without buffered parts is discarded. +/// +/// A trailing incomplete entry is emitted by [`remaining_entries`]. pub(crate) fn maybe_emit_entry( provider: &Arc, record: Record, @@ -244,7 +248,16 @@ pub(crate) fn maybe_emit_entry( RecordType::Full => entry = Some(convert_to_naive_entry(provider.clone(), record)), RecordType::First => { let region_id = record.meta.ns.region_id.into(); + let duplicate = buffered_records.get(®ion_id).is_some_and(|records| { + records.len() == 1 + && records[0].meta.tp == RecordType::First + && records[0].data == record.data + }); if let Some(records) = buffered_records.insert(region_id, vec![record]) { + if duplicate { + // A duplicate standalone First cannot form a complete entry. + return Ok(None); + } // Incomplete entry entry = Some(convert_to_multiple_entry( provider.clone(), @@ -261,6 +274,11 @@ pub(crate) fn maybe_emit_entry( if !records.is_empty() { // Safety: the records are guaranteed not empty if the key exists. let last_record = records.last().unwrap(); + if matches!(last_record.meta.tp, RecordType::Middle(last_seq) if last_seq == seq) + && last_record.data == record.data + { + return Ok(None); + } let legal = match last_record.meta.tp { // Legal if this record follows a First record. RecordType::First => seq == 1, @@ -290,14 +308,10 @@ pub(crate) fn maybe_emit_entry( provider.clone(), region_id, records, - )) + )); } else { - // Incomplete entry - entry = Some(convert_to_multiple_entry( - provider.clone(), - region_id, - vec![record], - )) + // Intentionally discard a Last record without buffered parts. It is either a + // duplicate Last record or the tail of an incomplete multipart entry. } } } @@ -359,7 +373,16 @@ mod tests { .unwrap() .is_none() ); - let record = new_test_record(RecordType::First, 2, region_id.as_u64(), vec![2; 100]); + let record = new_test_record(RecordType::First, 2, region_id.as_u64(), vec![1; 100]); + // An identical First is a duplicate and does not emit an incomplete entry. + assert!( + maybe_emit_entry(&provider, record, &mut buffer) + .unwrap() + .is_none() + ); + assert_eq!(buffer[®ion_id][0].meta.entry_id, 2); + + let record = new_test_record(RecordType::First, 3, region_id.as_u64(), vec![2; 100]); let incomplete_entry = maybe_emit_entry(&provider, record, &mut buffer) .unwrap() .unwrap(); @@ -369,29 +392,21 @@ mod tests { Entry::MultiplePart(MultiplePartEntry { provider: Provider::Kafka(provider.clone()), region_id, - entry_id: 1, + entry_id: 2, headers: vec![MultiplePartHeader::First], parts: vec![vec![1; 100]], }) ); - // `Last` overwrite `None` + // `Last` after `None` let mut buffer = HashMap::new(); let record = new_test_record(RecordType::Last, 1, region_id.as_u64(), vec![1; 100]); - let incomplete_entry = maybe_emit_entry(&provider, record, &mut buffer) - .unwrap() - .unwrap(); - - assert_eq!( - incomplete_entry, - Entry::MultiplePart(MultiplePartEntry { - provider: Provider::Kafka(provider.clone()), - region_id, - entry_id: 1, - headers: vec![MultiplePartHeader::Last], - parts: vec![vec![1; 100]], - }) + assert!( + maybe_emit_entry(&provider, record, &mut buffer) + .unwrap() + .is_none() ); + assert!(buffer.is_empty()); // `First` overwrite `Middle(0)` let mut buffer = HashMap::new(); @@ -451,6 +466,99 @@ mod tests { assert_matches!(err, error::Error::IllegalSequence { .. }); } + #[test] + fn test_maybe_emit_entry_deduplicates_multipart_records() { + let provider = Arc::new(KafkaProvider::new("my_topic".to_string())); + let region_id = RegionId::new(1, 1); + let mut buffer = HashMap::new(); + + for record in [ + new_test_record(RecordType::First, 1, region_id.as_u64(), vec![1; 100]), + new_test_record(RecordType::Middle(1), 1, region_id.as_u64(), vec![2; 100]), + new_test_record(RecordType::Middle(1), 1, region_id.as_u64(), vec![2; 100]), + ] { + assert!( + maybe_emit_entry(&provider, record, &mut buffer) + .unwrap() + .is_none() + ); + } + + let last = new_test_record(RecordType::Last, 1, region_id.as_u64(), vec![3; 100]); + let entry = maybe_emit_entry(&provider, last, &mut buffer) + .unwrap() + .unwrap(); + + assert_eq!( + entry, + Entry::MultiplePart(MultiplePartEntry { + provider: Provider::Kafka(provider.clone()), + region_id, + entry_id: 1, + headers: vec![ + MultiplePartHeader::First, + MultiplePartHeader::Middle(1), + MultiplePartHeader::Last, + ], + parts: vec![vec![1; 100], vec![2; 100], vec![3; 100]], + }) + ); + let duplicate = new_test_record(RecordType::Last, 1, region_id.as_u64(), vec![3; 100]); + assert!( + maybe_emit_entry(&provider, duplicate, &mut buffer) + .unwrap() + .is_none() + ); + } + + #[test] + fn test_maybe_emit_entry_rejects_conflicting_duplicate_middle_record() { + let provider = Arc::new(KafkaProvider::new("my_topic".to_string())); + let region_id = RegionId::new(1, 1); + let mut buffer = HashMap::new(); + + for record in [ + new_test_record(RecordType::First, 1, region_id.as_u64(), vec![1; 100]), + new_test_record(RecordType::Middle(1), 1, region_id.as_u64(), vec![2; 100]), + ] { + assert!( + maybe_emit_entry(&provider, record, &mut buffer) + .unwrap() + .is_none() + ); + } + + let duplicate = new_test_record(RecordType::Middle(1), 1, region_id.as_u64(), vec![3; 100]); + let err = maybe_emit_entry(&provider, duplicate, &mut buffer).unwrap_err(); + + assert_matches!(err, error::Error::IllegalSequence { .. }); + } + + #[test] + fn test_maybe_emit_entry_preserves_multipart_entry_order_before_full_record() { + let provider = Arc::new(KafkaProvider::new("my_topic".to_string())); + let region_id = RegionId::new(1, 1); + let mut buffer = HashMap::new(); + + let first = new_test_record(RecordType::First, 1, region_id.as_u64(), vec![1; 100]); + assert!( + maybe_emit_entry(&provider, first, &mut buffer) + .unwrap() + .is_none() + ); + let last = new_test_record(RecordType::Last, 2, region_id.as_u64(), vec![2; 100]); + let multipart_entry = maybe_emit_entry(&provider, last, &mut buffer) + .unwrap() + .unwrap(); + let full = new_test_record(RecordType::Full, 3, region_id.as_u64(), vec![3; 100]); + let full_entry = maybe_emit_entry(&provider, full, &mut buffer) + .unwrap() + .unwrap(); + + assert_eq!(multipart_entry.entry_id(), 2); + assert_eq!(full_entry.entry_id(), 3); + } + #[test] fn test_meta_size() { let meta = RecordMeta { diff --git a/src/log-store/src/noop/log_store.rs b/src/log-store/src/noop/log_store.rs index f0f056514e8..11b77e79108 100644 --- a/src/log-store/src/noop/log_store.rs +++ b/src/log-store/src/noop/log_store.rs @@ -86,6 +86,10 @@ impl LogStore for NoopLogStore { Ok(()) } + async fn obsolete_all(&self, _provider: &Provider, _region_id: RegionId) -> Result<()> { + Ok(()) + } + fn latest_entry_id(&self, _provider: &Provider) -> Result { Ok(0) } diff --git a/src/log-store/src/raft_engine/log_store.rs b/src/log-store/src/raft_engine/log_store.rs index dbd4d0daa7d..923110858b9 100644 --- a/src/log-store/src/raft_engine/log_store.rs +++ b/src/log-store/src/raft_engine/log_store.rs @@ -484,6 +484,12 @@ impl LogStore for RaftEngineLogStore { Ok(()) } + async fn obsolete_all(&self, provider: &Provider, region_id: RegionId) -> Result<()> { + let latest_entry_id = self.latest_entry_id(provider)?; + self.obsolete(provider, region_id, latest_entry_id).await?; + self.delete_namespace(provider).await + } + fn latest_entry_id(&self, provider: &Provider) -> Result { let ns = provider .as_raft_engine_provider() @@ -576,6 +582,38 @@ mod tests { assert!(logstore.list_namespaces().await.unwrap().is_empty()); } + #[tokio::test] + async fn test_obsolete_all_removes_entries_and_namespace() { + let dir = create_temp_dir("raft-engine-logstore-test"); + let logstore = RaftEngineLogStore::try_new( + dir.path().to_str().unwrap().to_string(), + &RaftEngineConfig::default(), + ) + .await + .unwrap(); + let region_id = RegionId::new(1, 1); + let provider = Provider::raft_engine_provider(region_id.as_u64()); + logstore.create_namespace(&provider).await.unwrap(); + for entry_id in 1..=3 { + logstore + .append( + EntryImpl::create( + entry_id, + region_id.as_u64(), + entry_id.to_string().into_bytes(), + ) + .into(), + ) + .await + .unwrap(); + } + + logstore.obsolete_all(&provider, region_id).await.unwrap(); + + assert_eq!(0, logstore.latest_entry_id(&provider).unwrap()); + assert!(logstore.list_namespaces().await.unwrap().is_empty()); + } + #[tokio::test] async fn test_append_and_read() { let dir = create_temp_dir("raft-engine-logstore-test"); diff --git a/src/meta-client/src/client/heartbeat.rs b/src/meta-client/src/client/heartbeat.rs index c31369dac78..add7bd54de6 100644 --- a/src/meta-client/src/client/heartbeat.rs +++ b/src/meta-client/src/client/heartbeat.rs @@ -39,6 +39,7 @@ use crate::error::{InvalidResponseHeaderSnafu, Result}; pub struct HeartbeatConfig { pub interval: Duration, pub retry_interval: Duration, + pub gc_enabled: bool, } impl Default for HeartbeatConfig { @@ -46,6 +47,7 @@ impl Default for HeartbeatConfig { Self { interval: BASE_HEARTBEAT_INTERVAL, retry_interval: BASE_HEARTBEAT_INTERVAL, + gc_enabled: false, } } } @@ -54,8 +56,8 @@ impl fmt::Display for HeartbeatConfig { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!( f, - "interval={:?}, retry={:?}", - self.interval, self.retry_interval + "interval={:?}, retry={:?}, gc_enabled={}", + self.interval, self.retry_interval, self.gc_enabled ) } } @@ -68,6 +70,7 @@ impl HeartbeatConfig { Self { interval: Duration::from_millis(cfg.heartbeat_interval_ms), retry_interval: Duration::from_millis(cfg.retry_interval_ms), + gc_enabled: cfg.gc_enabled, } } else { let fallback = Self::default(); @@ -259,7 +262,6 @@ impl Inner { .map_err(error::Error::from)? .context(error::CreateHeartbeatStreamSnafu)?; - // Extract heartbeat configuration from handshake response let config = HeartbeatConfig::from_response(&res); info!( diff --git a/src/meta-srv/src/gc.rs b/src/meta-srv/src/gc.rs index f9b3869bc63..f64f52eaeb5 100644 --- a/src/meta-srv/src/gc.rs +++ b/src/meta-srv/src/gc.rs @@ -29,6 +29,8 @@ mod scheduler; mod tracker; mod util; +pub(crate) use ctx::DefaultGcSchedulerCtx; +pub(crate) use options::EXPERIMENTAL_SOFT_DROP_ENABLED; pub use options::GcSchedulerOptions; pub use procedure::BatchGcProcedure; pub use scheduler::{Event, GcJobReport, GcScheduler, GcTickerRef}; diff --git a/src/meta-srv/src/gc/ctx.rs b/src/meta-srv/src/gc/ctx.rs index 6f00ff50a45..6b838e00478 100644 --- a/src/meta-srv/src/gc/ctx.rs +++ b/src/meta-srv/src/gc/ctx.rs @@ -12,13 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; use std::time::Duration; use common_meta::datanode::RegionStat; -use common_meta::key::TableMetadataManagerRef; +use common_meta::ddl_manager::DdlManagerRef; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; +use common_meta::key::{DroppedTableName, TableMetadataManagerRef}; +use common_meta::rpc::ddl::PurgeDroppedTableTask; use common_procedure::{ProcedureManagerRef, ProcedureWithId, watcher}; use common_telemetry::debug; use snafu::{OptionExt as _, ResultExt as _}; @@ -29,6 +33,7 @@ use crate::cluster::MetaPeerClientRef; use crate::error::{self, Result, TableMetadataManagerSnafu}; use crate::gc::Region2Peers; use crate::gc::procedure::BatchGcProcedure; +use crate::metrics::METRIC_META_GC_SOFT_DROP_PURGES_TOTAL; use crate::service::mailbox::MailboxRef; #[async_trait::async_trait] @@ -54,6 +59,77 @@ pub(crate) trait SchedulerCtx: Send + Sync { timeout: Duration, region_routes_override: Region2Peers, ) -> Result; + + async fn list_dropped_tables(&self) -> Result>; + + async fn purge_dropped_table(&self, table_id: TableId) -> Result<()>; + + fn try_reserve_purge( + &self, + table_id: TableId, + max_in_flight: usize, + ) -> Option; + + fn next_purge_scan_start(&self, _table_count: usize) -> usize { + 0 + } +} + +pub(crate) enum PurgeOutcome { + Succeeded, + Failed, +} + +pub(crate) struct PurgeReservation { + table_id: TableId, + in_flight: Arc>>, + outcome_recorded: bool, +} + +impl PurgeReservation { + pub(crate) fn try_new( + in_flight: Arc>>, + table_id: TableId, + max_in_flight: usize, + ) -> Option { + let mut tables = in_flight + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if tables.len() >= max_in_flight || !tables.insert(table_id) { + return None; + } + drop(tables); + Some(Self { + table_id, + in_flight, + outcome_recorded: false, + }) + } + + pub(crate) fn record_outcome(mut self, outcome: PurgeOutcome) { + let status = match outcome { + PurgeOutcome::Succeeded => "succeeded", + PurgeOutcome::Failed => "failed", + }; + METRIC_META_GC_SOFT_DROP_PURGES_TOTAL + .with_label_values(&[status]) + .inc(); + self.outcome_recorded = true; + } +} + +impl Drop for PurgeReservation { + fn drop(&mut self) { + if !self.outcome_recorded { + METRIC_META_GC_SOFT_DROP_PURGES_TOTAL + .with_label_values(&["cancelled"]) + .inc(); + } + self.in_flight + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .remove(&self.table_id); + } } pub(crate) struct DefaultGcSchedulerCtx { @@ -61,6 +137,12 @@ pub(crate) struct DefaultGcSchedulerCtx { pub(crate) table_metadata_manager: TableMetadataManagerRef, /// Procedure manager. pub(crate) procedure_manager: ProcedureManagerRef, + /// DDL manager used to submit the existing purge procedure. + pub(crate) ddl_manager: DdlManagerRef, + /// Process-local reservations for purge procedures submitted by this scheduler. + /// Procedure recovery after a metasrv restart may outlive this set. + in_flight_purges: Arc>>, + purge_scan_cursor: AtomicUsize, /// For getting `RegionStats`. pub(crate) meta_peer_client: MetaPeerClientRef, /// The mailbox to send messages. @@ -73,6 +155,7 @@ impl DefaultGcSchedulerCtx { pub fn try_new( table_metadata_manager: TableMetadataManagerRef, procedure_manager: ProcedureManagerRef, + ddl_manager: DdlManagerRef, meta_peer_client: MetaPeerClientRef, mailbox: MailboxRef, server_addr: String, @@ -80,6 +163,9 @@ impl DefaultGcSchedulerCtx { Ok(Self { table_metadata_manager, procedure_manager, + ddl_manager, + in_flight_purges: Arc::new(Mutex::new(HashSet::new())), + purge_scan_cursor: AtomicUsize::new(0), meta_peer_client, mailbox, server_addr, @@ -154,6 +240,36 @@ impl SchedulerCtx for DefaultGcSchedulerCtx { ) .await } + + async fn list_dropped_tables(&self) -> Result> { + self.table_metadata_manager + .list_dropped_tables() + .await + .context(TableMetadataManagerSnafu) + } + + async fn purge_dropped_table(&self, table_id: TableId) -> Result<()> { + self.ddl_manager + .submit_expired_purge_dropped_table_task(PurgeDroppedTableTask { table_id }) + .await + .context(error::SubmitDdlTaskSnafu)?; + Ok(()) + } + + fn try_reserve_purge( + &self, + table_id: TableId, + max_in_flight: usize, + ) -> Option { + PurgeReservation::try_new(self.in_flight_purges.clone(), table_id, max_in_flight) + } + + fn next_purge_scan_start(&self, table_count: usize) -> usize { + if table_count == 0 { + return 0; + } + self.purge_scan_cursor.fetch_add(1, Ordering::Relaxed) % table_count + } } impl DefaultGcSchedulerCtx { @@ -202,3 +318,48 @@ impl DefaultGcSchedulerCtx { Ok(gc_report) } } + +#[cfg(test)] +mod tests { + use std::panic::{AssertUnwindSafe, catch_unwind}; + + use super::*; + use crate::metrics::METRIC_META_GC_SOFT_DROP_PURGES_TOTAL; + + #[test] + fn test_purge_reservation_releases_slot_on_panic() { + let in_flight = Arc::new(std::sync::Mutex::new(HashSet::new())); + let cancelled = METRIC_META_GC_SOFT_DROP_PURGES_TOTAL.with_label_values(&["cancelled"]); + let before = cancelled.get(); + let reservation = PurgeReservation::try_new(in_flight.clone(), 1, 1).unwrap(); + + let result = catch_unwind(AssertUnwindSafe(|| { + let _reservation = reservation; + panic!("mock purge panic"); + })); + + assert!(result.is_err()); + assert!(PurgeReservation::try_new(in_flight, 2, 1).is_some()); + assert!(cancelled.get() > before); + } + + #[tokio::test] + async fn test_purge_reservation_releases_slot_on_task_abort() { + let in_flight = Arc::new(std::sync::Mutex::new(HashSet::new())); + let cancelled = METRIC_META_GC_SOFT_DROP_PURGES_TOTAL.with_label_values(&["cancelled"]); + let before = cancelled.get(); + let reservation = PurgeReservation::try_new(in_flight.clone(), 1, 1).unwrap(); + let handle = tokio::spawn(async move { + let _reservation = reservation; + std::future::pending::<()>().await; + }); + tokio::task::yield_now().await; + + handle.abort(); + let error = handle.await.unwrap_err(); + + assert!(error.is_cancelled()); + assert!(PurgeReservation::try_new(in_flight, 2, 1).is_some()); + assert!(cancelled.get() > before); + } +} diff --git a/src/meta-srv/src/gc/mock.rs b/src/meta-srv/src/gc/mock.rs index 0686d9a4653..f801df2707e 100644 --- a/src/meta-srv/src/gc/mock.rs +++ b/src/meta-srv/src/gc/mock.rs @@ -326,6 +326,26 @@ impl SchedulerCtx for MockSchedulerCtx { // Return the report with need_retry_regions populated - let the caller handle retry logic Ok(final_report) } + + async fn list_dropped_tables(&self) -> Result> { + Ok(vec![]) + } + + async fn purge_dropped_table(&self, _table_id: TableId) -> Result<()> { + Ok(()) + } + + fn try_reserve_purge( + &self, + table_id: TableId, + max_in_flight: usize, + ) -> Option { + crate::gc::ctx::PurgeReservation::try_new( + Arc::new(std::sync::Mutex::new(HashSet::new())), + table_id, + max_in_flight, + ) + } } pub struct TestEnv { diff --git a/src/meta-srv/src/gc/options.rs b/src/meta-srv/src/gc/options.rs index f6da6b850aa..4d1b93b9d95 100644 --- a/src/meta-srv/src/gc/options.rs +++ b/src/meta-srv/src/gc/options.rs @@ -23,6 +23,29 @@ use crate::error::{self, Result}; #[allow(unused)] pub(crate) const TICKER_INTERVAL: Duration = Duration::from_secs(60 * 5); +/// Soft drop remains disconnected from user configuration until it is ready. +pub(crate) const EXPERIMENTAL_SOFT_DROP_ENABLED: bool = false; + +/// Reserved configuration for garbage collecting soft-dropped tables. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(default)] +pub struct ExperimentalSoftDropGcOptions { + /// Reserved. This option is currently ignored. + pub enable: bool, + /// Reserved retention duration. This option is currently ignored. + #[serde(with = "humantime_serde")] + pub retention: Duration, +} + +impl Default for ExperimentalSoftDropGcOptions { + fn default() -> Self { + Self { + enable: false, + retention: Duration::from_days(7), + } + } +} + /// Configuration for GC operations. /// /// TODO(discord9): not expose most config to users for now, until GC scheduler is fully stable. @@ -33,6 +56,8 @@ pub struct GcSchedulerOptions { /// If set to false, no GC will be performed, and potentially some /// files from datanodes will never be deleted. pub enable: bool, + /// Reserved experimental soft-drop garbage collection options. + pub experimental_soft_drop: ExperimentalSoftDropGcOptions, /// Maximum number of tables to process concurrently. pub max_concurrent_tables: usize, /// Maximum number of retries per region when GC fails. @@ -73,6 +98,7 @@ impl Default for GcSchedulerOptions { fn default() -> Self { Self { enable: false, + experimental_soft_drop: ExperimentalSoftDropGcOptions::default(), max_concurrent_tables: 10, max_retries_per_region: 3, retry_backoff_duration: Duration::from_secs(5), @@ -174,3 +200,46 @@ impl GcSchedulerOptions { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_experimental_soft_drop_defaults() { + let options = GcSchedulerOptions::default(); + + assert!(!options.experimental_soft_drop.enable); + assert_eq!( + Duration::from_days(7), + options.experimental_soft_drop.retention + ); + } + + #[test] + fn test_experimental_soft_drop_valid_when_gc_is_enabled() { + let options = GcSchedulerOptions { + enable: true, + experimental_soft_drop: ExperimentalSoftDropGcOptions { + enable: true, + retention: Duration::from_days(1), + }, + ..Default::default() + }; + + assert!(options.validate().is_ok()); + } + + #[test] + fn test_experimental_soft_drop_options_are_ignored_by_validation() { + let options = GcSchedulerOptions { + experimental_soft_drop: ExperimentalSoftDropGcOptions { + enable: true, + retention: Duration::ZERO, + }, + ..Default::default() + }; + + assert!(options.validate().is_ok()); + } +} diff --git a/src/meta-srv/src/gc/scheduler.rs b/src/meta-srv/src/gc/scheduler.rs index 4f732f52bc0..e3e6a5f96ec 100644 --- a/src/meta-srv/src/gc/scheduler.rs +++ b/src/meta-srv/src/gc/scheduler.rs @@ -17,9 +17,7 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use common_meta::DatanodeId; -use common_meta::key::TableMetadataManagerRef; use common_meta::key::runtime_switch::RuntimeSwitchManagerRef; -use common_procedure::ProcedureManagerRef; use common_telemetry::tracing::Instrument as _; use common_telemetry::{error, info}; use snafu::ResultExt; @@ -27,18 +25,19 @@ use store_api::storage::{GcReport, RegionId}; use tokio::sync::mpsc::{Receiver, Sender}; use tokio::sync::{Mutex, oneshot}; -use crate::cluster::MetaPeerClientRef; use crate::define_ticker; use crate::error::{self, Error, Result}; -use crate::gc::Region2Peers; -use crate::gc::ctx::{DefaultGcSchedulerCtx, SchedulerCtx}; +#[cfg(test)] +use crate::gc::ctx::PurgeReservation; +use crate::gc::ctx::{PurgeOutcome, SchedulerCtx}; use crate::gc::dropped::DroppedRegionCollector; use crate::gc::options::{GcSchedulerOptions, TICKER_INTERVAL}; use crate::gc::tracker::RegionGcTracker; +use crate::gc::{EXPERIMENTAL_SOFT_DROP_ENABLED, Region2Peers}; use crate::metrics::{ METRIC_META_GC_SCHEDULER_CYCLES_TOTAL, METRIC_META_GC_SCHEDULER_DURATION_SECONDS, + METRIC_META_GC_SOFT_DROP_PURGES_TOTAL, }; -use crate::service::mailbox::MailboxRef; /// Report for a GC job. #[derive(Debug)] @@ -132,11 +131,7 @@ pub struct GcScheduler { impl GcScheduler { /// Creates a new [`GcScheduler`] with custom configuration. pub(crate) fn new_with_config( - table_metadata_manager: TableMetadataManagerRef, - procedure_manager: ProcedureManagerRef, - meta_peer_client: MetaPeerClientRef, - mailbox: MailboxRef, - server_addr: String, + ctx: impl SchedulerCtx + 'static, runtime_switch_manager: RuntimeSwitchManagerRef, config: GcSchedulerOptions, ) -> Result<(Self, GcTicker)> { @@ -146,13 +141,7 @@ impl GcScheduler { let (tx, rx) = Self::channel(); let gc_ticker = GcTicker::new(TICKER_INTERVAL, tx); let gc_trigger = Self { - ctx: Arc::new(DefaultGcSchedulerCtx::try_new( - table_metadata_manager, - procedure_manager, - meta_peer_client, - mailbox, - server_addr, - )?), + ctx: Arc::new(ctx), runtime_switch_manager, receiver: rx, config, @@ -218,6 +207,10 @@ impl GcScheduler { info!("Skip gc trigger because maintenance mode is enabled"); return Ok(GcJobReport::default()); } + if EXPERIMENTAL_SOFT_DROP_ENABLED && self.config.experimental_soft_drop.enable { + self.purge_expired_soft_dropped_tables(common_time::util::current_time_millis()) + .await; + } let span = common_telemetry::tracing::info_span!("meta_gc_handle_tick"); let report = self.trigger_gc().instrument(span).await?; @@ -229,6 +222,60 @@ impl GcScheduler { Ok(report) } + async fn purge_expired_soft_dropped_tables(&self, now_millis: i64) { + // The scheduler is only constructed after GcSchedulerOptions::validate(). + let retention_millis = + i64::try_from(self.config.experimental_soft_drop.retention.as_millis()) + .unwrap_or(i64::MAX); + let dropped_tables = match self.ctx.list_dropped_tables().await { + Ok(dropped_tables) => dropped_tables, + Err(error) => { + error!(error; "Failed to list soft-dropped tables for GC"); + return; + } + }; + let table_count = dropped_tables.len(); + let scan_start = self.ctx.next_purge_scan_start(table_count); + + for table in dropped_tables + .iter() + .cycle() + .skip(scan_start) + .take(table_count) + .filter(|table| { + table + .retention_expires_at + .or_else(|| { + table + .dropped_at + .and_then(|dropped_at| dropped_at.checked_add(retention_millis)) + }) + .is_some_and(|expires_at| expires_at <= now_millis) + }) + { + let table_id = table.table_id; + let Some(reservation) = self + .ctx + .try_reserve_purge(table_id, self.config.max_concurrent_tables) + else { + continue; + }; + METRIC_META_GC_SOFT_DROP_PURGES_TOTAL + .with_label_values(&["submitted"]) + .inc(); + let ctx = self.ctx.clone(); + common_runtime::spawn_global(async move { + match ctx.purge_dropped_table(table_id).await { + Ok(()) => reservation.record_outcome(PurgeOutcome::Succeeded), + Err(error) => { + reservation.record_outcome(PurgeOutcome::Failed); + error!(error; "Failed to purge expired soft-dropped table {}", table_id); + } + } + }); + } + } + /// Handles a manual GC request with optional specific parameters. /// /// If `region_ids` is specified, GC will be performed only on those regions. @@ -332,14 +379,17 @@ pub(crate) fn new_test_runtime_switch_manager() -> RuntimeSwitchManagerRef { #[cfg(test)] mod tests { use std::collections::HashMap; + use std::sync::Mutex as StdMutex; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use common_meta::datanode::RegionStat; + use common_meta::key::DroppedTableName; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; use store_api::storage::RegionId; use table::metadata::TableId; + use table::table_name::TableName; use super::*; @@ -348,6 +398,8 @@ mod tests { get_table_to_region_stats_calls: AtomicUsize, get_table_reparts_calls: AtomicUsize, gc_regions_calls: AtomicUsize, + list_dropped_tables_calls: AtomicUsize, + purge_dropped_table_calls: AtomicUsize, } impl CountingSchedulerCtx { @@ -367,6 +419,16 @@ mod tests { self.gc_regions_calls.load(Ordering::Relaxed), "gc_regions should not be called" ); + assert_eq!( + 0, + self.list_dropped_tables_calls.load(Ordering::Relaxed), + "list_dropped_tables should not be called" + ); + assert_eq!( + 0, + self.purge_dropped_table_calls.load(Ordering::Relaxed), + "purge_dropped_table should not be called" + ); } } @@ -407,6 +469,353 @@ mod tests { self.gc_regions_calls.fetch_add(1, Ordering::Relaxed); panic!("gc_regions should not be called in maintenance mode") } + + async fn list_dropped_tables(&self) -> Result> { + self.list_dropped_tables_calls + .fetch_add(1, Ordering::Relaxed); + panic!("list_dropped_tables should not be called in maintenance mode") + } + + async fn purge_dropped_table(&self, _table_id: TableId) -> Result<()> { + self.purge_dropped_table_calls + .fetch_add(1, Ordering::Relaxed); + panic!("purge_dropped_table should not be called in maintenance mode") + } + + fn try_reserve_purge( + &self, + _table_id: TableId, + _max_in_flight: usize, + ) -> Option { + panic!("try_reserve_purge should not be called in maintenance mode") + } + } + + #[derive(Default)] + struct SoftDropSchedulerCtx { + dropped_tables: StdMutex>, + purge_attempts: StdMutex>, + failed_table: StdMutex>, + never_complete_tables: StdMutex>, + in_flight_purges: Arc>>, + purge_scan_cursor: AtomicUsize, + region_gc_calls: AtomicUsize, + } + + #[async_trait::async_trait] + impl SchedulerCtx for SoftDropSchedulerCtx { + async fn get_table_to_region_stats(&self) -> Result>> { + self.region_gc_calls.fetch_add(1, Ordering::Relaxed); + Ok(HashMap::new()) + } + + async fn get_table_reparts(&self) -> Result> { + Ok(vec![]) + } + + async fn get_table_route( + &self, + _table_id: TableId, + ) -> Result<(TableId, PhysicalTableRouteValue)> { + unreachable!() + } + + async fn batch_get_table_route( + &self, + _table_ids: &[TableId], + ) -> Result> { + Ok(HashMap::new()) + } + + async fn gc_regions( + &self, + _region_ids: &[RegionId], + _full_file_listing: bool, + _timeout: Duration, + _region_routes_override: Region2Peers, + ) -> Result { + Ok(GcReport::default()) + } + + async fn list_dropped_tables(&self) -> Result> { + Ok(self.dropped_tables.lock().unwrap().clone()) + } + + async fn purge_dropped_table(&self, table_id: TableId) -> Result<()> { + self.purge_attempts.lock().unwrap().push(table_id); + if self + .never_complete_tables + .lock() + .unwrap() + .contains(&table_id) + { + std::future::pending().await + } + if *self.failed_table.lock().unwrap() == Some(table_id) { + return crate::error::UnexpectedSnafu { + violated: format!("mock purge failure for table {table_id}"), + } + .fail(); + } + Ok(()) + } + + fn try_reserve_purge( + &self, + table_id: TableId, + max_in_flight: usize, + ) -> Option { + PurgeReservation::try_new(self.in_flight_purges.clone(), table_id, max_in_flight) + } + + fn next_purge_scan_start(&self, table_count: usize) -> usize { + if table_count == 0 { + return 0; + } + self.purge_scan_cursor.fetch_add(1, Ordering::Relaxed) % table_count + } + } + + fn dropped_table(table_id: TableId, dropped_at: Option) -> DroppedTableName { + DroppedTableName { + table_id, + table_name: TableName::new("greptime", "public", format!("table_{table_id}")), + dropped_at, + retention_expires_at: None, + drop_generation: None, + } + } + + fn dropped_table_with_deadline(table_id: TableId, deadline: i64) -> DroppedTableName { + DroppedTableName { + retention_expires_at: Some(deadline), + ..dropped_table(table_id, Some(i64::MAX)) + } + } + + fn soft_drop_scheduler(ctx: Arc) -> GcScheduler { + let (tx, rx) = GcScheduler::channel(); + drop(tx); + GcScheduler { + ctx, + runtime_switch_manager: new_test_runtime_switch_manager(), + receiver: rx, + config: GcSchedulerOptions { + enable: true, + experimental_soft_drop: crate::gc::options::ExperimentalSoftDropGcOptions { + enable: true, + retention: Duration::from_millis(100), + }, + ..Default::default() + }, + region_gc_tracker: Arc::new(Mutex::new(HashMap::new())), + last_tracker_cleanup: Arc::new(Mutex::new(Instant::now())), + } + } + + async fn wait_for_purge_attempts(ctx: &SoftDropSchedulerCtx, expected: usize) { + tokio::time::timeout(Duration::from_secs(1), async { + while ctx.purge_attempts.lock().unwrap().len() < expected { + tokio::task::yield_now().await; + } + }) + .await + .expect("purge tasks should be scheduled"); + } + + #[tokio::test] + async fn test_purge_expired_soft_dropped_tables_filters_by_retention_and_timestamp() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![ + dropped_table(1, Some(899)), + dropped_table(2, Some(900)), + dropped_table(3, Some(901)), + dropped_table(4, None), + dropped_table(5, Some(i64::MAX)), + ]; + let scheduler = soft_drop_scheduler(ctx.clone()); + + scheduler.purge_expired_soft_dropped_tables(1_000).await; + wait_for_purge_attempts(&ctx, 2).await; + + let mut attempts = ctx.purge_attempts.lock().unwrap().clone(); + attempts.sort_unstable(); + assert_eq!(vec![1, 2], attempts); + } + + #[tokio::test] + async fn test_purge_uses_fixed_deadline_before_legacy_retention_fallback() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![ + dropped_table_with_deadline(1, 1_000), + dropped_table_with_deadline(2, 1_001), + dropped_table(3, Some(900)), + ]; + let scheduler = soft_drop_scheduler(ctx.clone()); + + scheduler.purge_expired_soft_dropped_tables(1_000).await; + wait_for_purge_attempts(&ctx, 2).await; + + let mut attempts = ctx.purge_attempts.lock().unwrap().clone(); + attempts.sort_unstable(); + assert_eq!(vec![1, 3], attempts); + } + + #[tokio::test] + async fn test_purge_saturates_legacy_retention_that_exceeds_i64() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![dropped_table(1, Some(0))]; + ctx.never_complete_tables.lock().unwrap().insert(1); + let mut scheduler = soft_drop_scheduler(ctx.clone()); + scheduler.config.experimental_soft_drop.retention = + Duration::from_millis(i64::MAX as u64 + 1); + + scheduler.purge_expired_soft_dropped_tables(0).await; + + assert!(!ctx.in_flight_purges.lock().unwrap().contains(&1)); + } + + #[tokio::test] + async fn test_purge_failure_does_not_prevent_other_purges() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![ + dropped_table(1, Some(0)), + dropped_table(2, Some(0)), + dropped_table(3, Some(0)), + ]; + *ctx.failed_table.lock().unwrap() = Some(2); + let scheduler = soft_drop_scheduler(ctx.clone()); + + scheduler.purge_expired_soft_dropped_tables(1_000).await; + wait_for_purge_attempts(&ctx, 3).await; + + let mut attempts = ctx.purge_attempts.lock().unwrap().clone(); + attempts.sort_unstable(); + assert_eq!(vec![1, 2, 3], attempts); + } + + #[tokio::test] + async fn test_purge_scans_rotate_after_failure() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = + vec![dropped_table(1, Some(0)), dropped_table(2, Some(0))]; + *ctx.failed_table.lock().unwrap() = Some(1); + let mut scheduler = soft_drop_scheduler(ctx.clone()); + scheduler.config.max_concurrent_tables = 1; + + scheduler.purge_expired_soft_dropped_tables(1_000).await; + wait_for_purge_attempts(&ctx, 1).await; + tokio::time::timeout(Duration::from_secs(1), async { + while !ctx.in_flight_purges.lock().unwrap().is_empty() { + tokio::task::yield_now().await; + } + }) + .await + .expect("failed purge should release its in-flight slot"); + + scheduler.purge_expired_soft_dropped_tables(1_000).await; + wait_for_purge_attempts(&ctx, 2).await; + + assert_eq!(vec![1, 2], *ctx.purge_attempts.lock().unwrap()); + } + + #[tokio::test] + async fn test_tick_ignores_experimental_soft_drop_config_and_still_runs_region_gc() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![dropped_table(1, Some(i64::MIN))]; + let scheduler = soft_drop_scheduler(ctx.clone()); + + scheduler.handle_tick().await.unwrap(); + + assert!(ctx.purge_attempts.lock().unwrap().is_empty()); + assert_eq!(1, ctx.region_gc_calls.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn test_tick_with_no_tombstones_makes_no_purge_submissions() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + let scheduler = soft_drop_scheduler(ctx.clone()); + + scheduler.handle_tick().await.unwrap(); + + assert!(ctx.purge_attempts.lock().unwrap().is_empty()); + assert_eq!(1, ctx.region_gc_calls.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn test_multiple_purge_scans_do_not_duplicate_in_flight_purges() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![ + dropped_table(1, Some(i64::MIN)), + dropped_table(2, Some(i64::MIN)), + ]; + ctx.never_complete_tables.lock().unwrap().insert(1); + *ctx.failed_table.lock().unwrap() = Some(2); + let mut scheduler = soft_drop_scheduler(ctx.clone()); + scheduler.config.max_concurrent_tables = 2; + + tokio::time::timeout( + Duration::from_millis(100), + scheduler.purge_expired_soft_dropped_tables(1_000), + ) + .await + .expect("first scan should not wait for purge completion"); + wait_for_purge_attempts(&ctx, 2).await; + tokio::time::timeout(Duration::from_secs(1), async { + while ctx.in_flight_purges.lock().unwrap().contains(&2) { + tokio::task::yield_now().await; + } + }) + .await + .expect("terminal failure should release its in-flight slot"); + tokio::time::timeout( + Duration::from_millis(100), + scheduler.purge_expired_soft_dropped_tables(1_000), + ) + .await + .expect("second scan should not wait for purge completion"); + wait_for_purge_attempts(&ctx, 3).await; + + let attempts = ctx.purge_attempts.lock().unwrap().clone(); + assert_eq!( + 1, + attempts.iter().filter(|&&table_id| table_id == 1).count() + ); + assert_eq!( + 2, + attempts.iter().filter(|&&table_id| table_id == 2).count() + ); + assert_eq!(0, ctx.region_gc_calls.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn test_purge_scan_caps_submissions() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![ + dropped_table(1, Some(i64::MIN)), + dropped_table(2, Some(i64::MIN)), + dropped_table(3, Some(i64::MIN)), + ]; + let mut scheduler = soft_drop_scheduler(ctx.clone()); + scheduler.config.max_concurrent_tables = 2; + ctx.never_complete_tables.lock().unwrap().extend([1, 2]); + + scheduler.purge_expired_soft_dropped_tables(1_000).await; + wait_for_purge_attempts(&ctx, 2).await; + + assert_eq!(2, ctx.purge_attempts.lock().unwrap().len()); + assert_eq!(0, ctx.region_gc_calls.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn test_manual_gc_does_not_purge_soft_dropped_tables() { + let ctx = Arc::new(SoftDropSchedulerCtx::default()); + *ctx.dropped_tables.lock().unwrap() = vec![dropped_table(1, Some(i64::MIN))]; + let scheduler = soft_drop_scheduler(ctx.clone()); + + scheduler.handle_manual_gc(None, None, None).await.unwrap(); + + assert!(ctx.purge_attempts.lock().unwrap().is_empty()); } struct ErrorMockSchedulerCtx; @@ -447,6 +856,26 @@ mod tests { } .fail() } + + async fn list_dropped_tables(&self) -> Result> { + Ok(vec![]) + } + + async fn purge_dropped_table(&self, _table_id: TableId) -> Result<()> { + Ok(()) + } + + fn try_reserve_purge( + &self, + table_id: TableId, + max_in_flight: usize, + ) -> Option { + PurgeReservation::try_new( + Arc::new(StdMutex::new(HashSet::new())), + table_id, + max_in_flight, + ) + } } #[tokio::test] diff --git a/src/meta-srv/src/handler.rs b/src/meta-srv/src/handler.rs index 9cfd4e6079d..72d5fcdab8a 100644 --- a/src/meta-srv/src/handler.rs +++ b/src/meta-srv/src/handler.rs @@ -21,8 +21,8 @@ use std::time::{Duration, Instant}; use api::v1::meta::mailbox_message::Payload; use api::v1::meta::{ - HeartbeatRequest, HeartbeatResponse, MailboxMessage, PROTOCOL_VERSION, RegionLease, - ResponseHeader, Role, + HeartbeatConfig, HeartbeatRequest, HeartbeatResponse, MailboxMessage, PROTOCOL_VERSION, + RegionLease, ResponseHeader, Role, }; use check_leader_handler::CheckLeaderHandler; use collect_cluster_info_handler::{ @@ -387,7 +387,8 @@ impl HeartbeatHandlerGroup { // Populate heartbeat_config during handshake let heartbeat_config = if is_handshake { - let config = ctx.heartbeat_options_for(role).into(); + let mut config: HeartbeatConfig = ctx.heartbeat_options_for(role).into(); + config.gc_enabled = ctx.gc_enabled; info!( "Handshake with {:?} node, sending config: {:?}", diff --git a/src/meta-srv/src/handler/test_utils.rs b/src/meta-srv/src/handler/test_utils.rs index 7bfdfbee792..d8241c785ea 100644 --- a/src/meta-srv/src/handler/test_utils.rs +++ b/src/meta-srv/src/handler/test_utils.rs @@ -92,6 +92,7 @@ impl TestEnv { leader_region_registry: self.leader_region_registry.clone(), topic_stats_registry: self.topic_stats_registry.clone(), heartbeat_interval: BASE_HEARTBEAT_INTERVAL, + gc_enabled: false, is_handshake: false, } } diff --git a/src/meta-srv/src/metasrv.rs b/src/meta-srv/src/metasrv.rs index f7f5bbf77da..7480b8cb500 100644 --- a/src/meta-srv/src/metasrv.rs +++ b/src/meta-srv/src/metasrv.rs @@ -175,6 +175,7 @@ impl From for HeartbeatConfig { Self { heartbeat_interval_ms: opts.interval.as_millis() as u64, retry_interval_ms: opts.retry_interval.as_millis() as u64, + gc_enabled: false, } } } @@ -438,6 +439,7 @@ pub struct Context { pub leader_region_registry: LeaderRegionRegistryRef, pub topic_stats_registry: TopicStatsRegistryRef, pub heartbeat_interval: Duration, + pub gc_enabled: bool, pub is_handshake: bool, } @@ -932,6 +934,7 @@ impl Metasrv { leader_region_registry, topic_stats_registry, heartbeat_interval: self.options().heartbeat_interval, + gc_enabled: self.options().gc.enable, is_handshake: false, } } diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index 90ce54142f1..fd7be6ae6d4 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -28,7 +28,9 @@ use common_meta::ddl::table_meta::{TableMetadataAllocator, TableMetadataAllocato use common_meta::ddl::{ DdlContext, NoopRegionFailureDetectorControl, RegionFailureDetectorControllerRef, }; -use common_meta::ddl_manager::{DdlManager, DdlManagerConfiguratorRef}; +use common_meta::ddl_manager::{ + DdlManager, DdlManagerConfiguratorRef, RepartitionProcedureFactoryRef, +}; use common_meta::distributed_time_constants::default_distributed_time_constants; use common_meta::key::TableMetadataManager; use common_meta::key::flow::FlowMetadataManager; @@ -56,7 +58,7 @@ use crate::cache_invalidator::MetasrvCacheInvalidator; use crate::cluster::MetaPeerClientRef; use crate::error::{self, BuildWalProviderSnafu, OtherSnafu, Result}; use crate::events::EventHandlerImpl; -use crate::gc::GcScheduler; +use crate::gc::{DefaultGcSchedulerCtx, EXPERIMENTAL_SOFT_DROP_ENABLED, GcScheduler}; use crate::greptimedb_telemetry::get_greptimedb_telemetry_task; use crate::handler::failure_handler::RegionFailureHandler; use crate::handler::flow_state_handler::FlowStateHandler; @@ -70,7 +72,9 @@ use crate::metasrv::{ use crate::peer::MetasrvPeerAllocator; use crate::procedure::region_migration::DefaultContextFactory; use crate::procedure::region_migration::manager::RegionMigrationManager; -use crate::procedure::repartition::DefaultRepartitionProcedureFactory; +use crate::procedure::repartition::{ + DefaultRepartitionProcedureFactory, GcDisabledRepartitionProcedureFactory, +}; use crate::procedure::wal_prune::Context as WalPruneContext; use crate::procedure::wal_prune::manager::{WalPruneManager, WalPruneTicker}; use crate::region::flush_trigger::RegionFlushTrigger; @@ -426,12 +430,20 @@ impl MetasrvBuilder { flow_metadata_allocator: flow_metadata_allocator.clone(), region_failure_detector_controller, soft_drop_enabled: ddl_soft_drop_enabled(&options), + soft_drop_retention: ddl_soft_drop_retention(&options), }; let procedure_manager_c = procedure_manager.clone(); - let repartition_procedure_factory = Arc::new(DefaultRepartitionProcedureFactory::new( - mailbox.clone(), - options.grpc.server_addr.clone(), - )); + let repartition_procedure_factory: RepartitionProcedureFactoryRef = if options.gc.enable { + Arc::new(DefaultRepartitionProcedureFactory::new( + mailbox.clone(), + options.grpc.server_addr.clone(), + )) + } else { + Arc::new(GcDisabledRepartitionProcedureFactory::new( + mailbox.clone(), + options.grpc.server_addr.clone(), + )) + }; let ddl_manager = DdlManager::try_new( ddl_context, procedure_manager_c, @@ -510,12 +522,16 @@ impl MetasrvBuilder { }; let gc_ticker = if options.gc.enable { - let (gc_scheduler, gc_ticker) = GcScheduler::new_with_config( + let gc_scheduler_ctx = DefaultGcSchedulerCtx::try_new( table_metadata_manager.clone(), procedure_manager.clone(), + ddl_manager.clone(), meta_peer_client.clone(), mailbox.clone(), options.grpc.server_addr.clone(), + )?; + let (gc_scheduler, gc_ticker) = GcScheduler::new_with_config( + gc_scheduler_ctx, runtime_switch_manager.clone(), options.gc.clone(), )?; @@ -676,10 +692,13 @@ fn build_procedure_manager( } /// Resolves if soft-drop is enabled from metasrv options. -fn ddl_soft_drop_enabled(_options: &MetasrvOptions) -> bool { - // TODO(hl): add a dedicated soft-drop cluster config - // when wiring the user-facing option. - false +fn ddl_soft_drop_enabled(options: &MetasrvOptions) -> bool { + EXPERIMENTAL_SOFT_DROP_ENABLED && options.gc.experimental_soft_drop.enable +} + +/// Returns soft-drop retention for recovering persisted procedures. +fn ddl_soft_drop_retention(options: &MetasrvOptions) -> Option { + Some(options.gc.experimental_soft_drop.retention) } impl Default for MetasrvBuilder { @@ -693,3 +712,49 @@ pub struct DdlManagerConfigureContext { pub kv_backend: KvBackendRef, pub meta_peer_client: MetaPeerClientRef, } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_ddl_soft_drop_gate_preserves_retention_for_recovery() { + let mut options = MetasrvOptions::default(); + options.gc.enable = true; + options.gc.experimental_soft_drop.enable = true; + options.gc.experimental_soft_drop.retention = Duration::from_secs(123); + + assert!(!ddl_soft_drop_enabled(&options)); + assert_eq!( + Some(Duration::from_secs(123)), + ddl_soft_drop_retention(&options) + ); + } + + #[test] + fn test_ddl_soft_drop_is_disabled_by_default() { + assert!(!ddl_soft_drop_enabled(&MetasrvOptions::default())); + } + + #[test] + fn test_ignored_soft_drop_options_are_not_validated() { + let mut options = MetasrvOptions::default(); + options.gc.experimental_soft_drop.enable = true; + options.gc.experimental_soft_drop.retention = Duration::ZERO; + + assert!(options.gc.validate().is_ok()); + } + + #[tokio::test] + async fn test_builder_skips_gc_validation_when_gc_is_disabled() { + let mut options = MetasrvOptions::default(); + options.gc.enable = false; + options.gc.max_concurrent_tables = 0; + + MetasrvBuilder::new() + .options(options) + .build() + .await + .unwrap(); + } +} diff --git a/src/meta-srv/src/metrics.rs b/src/meta-srv/src/metrics.rs index 879b1de2065..76283495c06 100644 --- a/src/meta-srv/src/metrics.rs +++ b/src/meta-srv/src/metrics.rs @@ -127,4 +127,12 @@ lazy_static! { "Total regions that failed GC", ) .unwrap(); + + /// Soft-drop purge submissions by status. + pub static ref METRIC_META_GC_SOFT_DROP_PURGES_TOTAL: IntCounterVec = register_int_counter_vec!( + "greptime_metasrv_gc_soft_drop_purges_total", + "Soft-drop purge submissions by status", + &["status"], + ) + .unwrap(); } diff --git a/src/meta-srv/src/procedure/repartition.rs b/src/meta-srv/src/procedure/repartition.rs index 53e4bacf1ec..2aa8f18799c 100644 --- a/src/meta-srv/src/procedure/repartition.rs +++ b/src/meta-srv/src/procedure/repartition.rs @@ -808,6 +808,50 @@ impl DefaultRepartitionProcedureFactory { } } +/// Rejects new repartition requests when metasrv GC is disabled. +/// +/// Procedure loaders are still delegated to the enabled factory so procedures +/// persisted before a metasrv restart can be recovered. +pub struct GcDisabledRepartitionProcedureFactory { + enabled_factory: DefaultRepartitionProcedureFactory, +} + +impl GcDisabledRepartitionProcedureFactory { + pub fn new(mailbox: MailboxRef, server_addr: String) -> Self { + Self { + enabled_factory: DefaultRepartitionProcedureFactory::new(mailbox, server_addr), + } + } +} + +impl RepartitionProcedureFactory for GcDisabledRepartitionProcedureFactory { + fn create( + &self, + _ddl_ctx: &DdlContext, + _table_name: TableName, + _table_id: TableId, + _source: RepartitionSource, + _to_exprs: Vec, + _timeout: Option, + ) -> std::result::Result { + Err(BoxedError::new( + error::InvalidArgumentsSnafu { + err_msg: "Repartition requires metasrv GC to be enabled".to_string(), + } + .build(), + )) + } + + fn register_loaders( + &self, + ddl_ctx: &DdlContext, + procedure_manager: &ProcedureManagerRef, + ) -> std::result::Result<(), BoxedError> { + self.enabled_factory + .register_loaders(ddl_ctx, procedure_manager) + } +} + impl RepartitionProcedureFactory for DefaultRepartitionProcedureFactory { fn create( &self, @@ -919,7 +963,7 @@ mod tests { use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; - use common_error::ext::BoxedError; + use common_error::ext::{BoxedError, ErrorExt}; use common_error::mock::MockError; use common_error::status_code::StatusCode; use common_meta::ddl::test_util::datanode_handler::{ @@ -1050,6 +1094,37 @@ mod tests { ) } + #[test] + fn test_gc_disabled_factory_rejects_repartition() { + let env = TestingEnv::new(); + let node_manager = Arc::new(MockDatanodeManager::new(UnexpectedErrorDatanodeHandler)); + let ddl_ctx = env.ddl_context(node_manager); + let factory = GcDisabledRepartitionProcedureFactory::new( + env.mailbox_ctx.mailbox().clone(), + env.server_addr.clone(), + ); + + let err = factory + .create( + &ddl_ctx, + TableName::new("test_catalog", "test_schema", "test_table"), + 1024, + RepartitionSource::Unpartitioned { + partition_columns: vec![], + }, + vec![], + None, + ) + .err() + .expect("GC-disabled factory must reject repartition"); + + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + assert_eq!( + "Invalid arguments: Repartition requires metasrv GC to be enabled", + err.to_string() + ); + } + #[test] fn test_filter_allocated_region_routes() { let table_id = 1024; diff --git a/src/meta-srv/src/procedure/utils.rs b/src/meta-srv/src/procedure/utils.rs index bcae558773a..81dfe56e670 100644 --- a/src/meta-srv/src/procedure/utils.rs +++ b/src/meta-srv/src/procedure/utils.rs @@ -547,6 +547,7 @@ pub mod test_data { leader_region_registry: Arc::new(LeaderRegionRegistry::default()), region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl), soft_drop_enabled: false, + soft_drop_retention: None, } } } diff --git a/src/metric-engine/src/config.rs b/src/metric-engine/src/config.rs index 9acee3c2b6a..3eaaf20aa78 100644 --- a/src/metric-engine/src/config.rs +++ b/src/metric-engine/src/config.rs @@ -23,9 +23,6 @@ pub(crate) const DEFAULT_FLUSH_METADATA_REGION_INTERVAL: Duration = Duration::fr /// Configuration for the metric engine. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct EngineConfig { - /// Whether to use sparse primary key encoding. - #[serde(default = "EngineConfig::default_sparse_primary_key_encoding")] - pub sparse_primary_key_encoding: bool, /// Whether to store exact integral metric values in internal Int64 columns. #[serde(default)] pub experimental_enable_metric_value_split: bool, @@ -41,7 +38,6 @@ impl Default for EngineConfig { fn default() -> Self { Self { flush_metadata_region_interval: DEFAULT_FLUSH_METADATA_REGION_INTERVAL, - sparse_primary_key_encoding: Self::default_sparse_primary_key_encoding(), experimental_enable_metric_value_split: false, } } @@ -52,10 +48,6 @@ impl EngineConfig { DEFAULT_FLUSH_METADATA_REGION_INTERVAL } - fn default_sparse_primary_key_encoding() -> bool { - true - } - /// Sanitizes the configuration. pub fn sanitize(&mut self) { if self.flush_metadata_region_interval.is_zero() { diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index 3cc18f5f153..a57344525a4 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -237,6 +237,9 @@ impl RegionEngine for MetricEngine { } RegionRequest::Drop(drop) => self.inner.drop_region(region_id, drop).await, RegionRequest::Open(open) => self.inner.open_region(region_id, open).await, + RegionRequest::CleanUp(clean_up) => { + self.inner.clean_up_region(region_id, clean_up).await + } RegionRequest::Close(close) => self.inner.close_region(region_id, close).await, RegionRequest::Alter(alter) => { self.inner @@ -471,13 +474,14 @@ impl MetricEngine { let state = Arc::new(RwLock::default()); config.sanitize(); let flush_interval = config.flush_metadata_region_interval; + let enable_metric_value_split = config.experimental_enable_metric_value_split; let inner = Arc::new(MetricEngineInner { mito: mito.clone(), metadata_region, data_region, state: state.clone(), - config, row_modifier: RowModifier::default(), + enable_metric_value_split, flush_task: RepeatedTask::new( flush_interval, Box::new(FlushMetadataRegionTask { @@ -556,11 +560,6 @@ impl MetricEngine { ) -> Result { self.inner.scan_to_stream(region_id, request).await } - - /// Returns the configuration of the engine. - pub fn config(&self) -> &EngineConfig { - &self.inner.config - } } struct MetricEngineInner { @@ -568,8 +567,8 @@ struct MetricEngineInner { metadata_region: MetadataRegion, data_region: DataRegion, state: Arc>, - config: EngineConfig, row_modifier: RowModifier, + enable_metric_value_split: bool, flush_task: RepeatedTask, } @@ -585,8 +584,8 @@ mod test { use store_api::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY; use store_api::mito_engine_options::WAL_OPTIONS_KEY; use store_api::region_request::{ - PathType, RegionCloseRequest, RegionDropRequest, RegionFlushRequest, RegionOpenRequest, - RegionRequest, + PathType, RegionCleanUpRequest, RegionCloseRequest, RegionDropRequest, RegionFlushRequest, + RegionOpenRequest, RegionRequest, }; use super::*; @@ -604,7 +603,7 @@ mod test { engine .handle_request( physical_region_id, - RegionRequest::Close(RegionCloseRequest {}), + RegionRequest::Close(RegionCloseRequest::default()), ) .await .unwrap(); @@ -632,7 +631,7 @@ mod test { engine .handle_request( nonexistent_region_id, - RegionRequest::Close(RegionCloseRequest {}), + RegionRequest::Close(RegionCloseRequest::default()), ) .await .unwrap(); @@ -656,6 +655,58 @@ mod test { .unwrap(); } + #[tokio::test] + async fn test_offline_cleanup_physical_region() { + let env = TestEnv::new().await; + env.init_metric_region().await; + let engine = env.metric(); + let mito = env.mito(); + let physical_region_id = env.default_physical_region_id(); + let metadata_region_id = crate::utils::to_metadata_region_id(physical_region_id); + let data_region_id = crate::utils::to_data_region_id(physical_region_id); + + engine + .handle_request( + physical_region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + + let object_store = env.get_object_store().unwrap(); + let metadata_region_dir = region_dir_from_table_dir( + &TestEnv::default_table_dir(), + metadata_region_id, + PathType::Metadata, + ); + let data_region_dir = region_dir_from_table_dir( + &TestEnv::default_table_dir(), + data_region_id, + PathType::Data, + ); + assert!(object_store.exists(&metadata_region_dir).await.unwrap()); + assert!(object_store.exists(&data_region_dir).await.unwrap()); + + let physical_region_option = [(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new())] + .into_iter() + .collect(); + let clean_up_request = RegionCleanUpRequest { + engine: METRIC_ENGINE_NAME.to_string(), + table_dir: TestEnv::default_table_dir(), + path_type: PathType::Bare, + options: physical_region_option, + }; + engine + .handle_request(physical_region_id, RegionRequest::CleanUp(clean_up_request)) + .await + .unwrap(); + + assert!(!mito.is_region_exists(metadata_region_id)); + assert!(!mito.is_region_exists(data_region_id)); + assert!(!object_store.exists(&metadata_region_dir).await.unwrap()); + assert!(!object_store.exists(&data_region_dir).await.unwrap()); + } + #[tokio::test] async fn test_role() { let env = TestEnv::new().await; @@ -736,7 +787,7 @@ mod test { metric_engine .handle_request( physical_region_id, - RegionRequest::Close(RegionCloseRequest {}), + RegionRequest::Close(RegionCloseRequest::default()), ) .await .unwrap(); @@ -828,7 +879,10 @@ mod test { // Closes all regions for region_id in logical_region_ids.iter().chain(physical_region_ids.iter()) { metric_engine - .handle_request(*region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + *region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); } diff --git a/src/metric-engine/src/engine/close.rs b/src/metric-engine/src/engine/close.rs index 1af507cf367..fdf2685c0d1 100644 --- a/src/metric-engine/src/engine/close.rs +++ b/src/metric-engine/src/engine/close.rs @@ -29,7 +29,7 @@ impl MetricEngineInner { pub async fn close_region( &self, region_id: RegionId, - _req: RegionCloseRequest, + req: RegionCloseRequest, ) -> Result { let data_region_id = utils::to_data_region_id(region_id); if self @@ -38,7 +38,8 @@ impl MetricEngineInner { .unwrap() .exist_physical_region(data_region_id) { - self.close_physical_region(data_region_id).await?; + self.close_physical_region(data_region_id, req.flush_on_close) + .await?; self.state .write() .unwrap() @@ -59,18 +60,25 @@ impl MetricEngineInner { } } - pub(crate) async fn close_physical_region(&self, region_id: RegionId) -> Result { + pub(crate) async fn close_physical_region( + &self, + region_id: RegionId, + flush_on_close: bool, + ) -> Result { let data_region_id = utils::to_data_region_id(region_id); let metadata_region_id = utils::to_metadata_region_id(region_id); self.mito - .handle_request(data_region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + data_region_id, + RegionRequest::Close(RegionCloseRequest { flush_on_close }), + ) .await .with_context(|_| CloseMitoRegionSnafu { region_id })?; self.mito .handle_request( metadata_region_id, - RegionRequest::Close(RegionCloseRequest {}), + RegionRequest::Close(RegionCloseRequest { flush_on_close }), ) .await .with_context(|_| CloseMitoRegionSnafu { region_id })?; diff --git a/src/metric-engine/src/engine/create.rs b/src/metric-engine/src/engine/create.rs index 64ca199599a..ac293353404 100644 --- a/src/metric-engine/src/engine/create.rs +++ b/src/metric-engine/src/engine/create.rs @@ -551,7 +551,7 @@ impl MetricEngineInner { let table_id_col_def = request.column_metadatas.iter().any(is_metric_name_col); let tsid_col_def = request.column_metadatas.iter().any(is_tsid_col); - let value_split_enabled = self.config.experimental_enable_metric_value_split + let value_split_enabled = self.enable_metric_value_split && request .options .get(MERGE_MODE_KEY) @@ -585,10 +585,7 @@ impl MetricEngineInner { data_region_request.primary_key = primary_key; // set data region options - set_data_region_options( - &mut data_region_request.options, - self.config.sparse_primary_key_encoding, - ); + set_data_region_options(&mut data_region_request.options); data_region_request } diff --git a/src/metric-engine/src/engine/open.rs b/src/metric-engine/src/engine/open.rs index 8fcdfcd8211..85b538d05b0 100644 --- a/src/metric-engine/src/engine/open.rs +++ b/src/metric-engine/src/engine/open.rs @@ -22,14 +22,17 @@ use datafusion::common::HashMap; use mito2::engine::MITO_ENGINE_NAME; use snafu::{OptionExt, ResultExt}; use store_api::region_engine::{BatchResponses, RegionEngine}; -use store_api::region_request::{AffectedRows, PathType, RegionOpenRequest, ReplayCheckpoint}; +use store_api::region_request::{ + AffectedRows, PathType, RegionCleanUpRequest, RegionOpenRequest, RegionRequest, + ReplayCheckpoint, +}; use store_api::storage::RegionId; use crate::engine::MetricEngineInner; use crate::engine::create::region_options_for_metadata_region; use crate::engine::options::{PhysicalRegionOptions, set_data_region_options}; use crate::error::{ - BatchOpenMitoRegionSnafu, NoOpenRegionResultSnafu, OpenMitoRegionSnafu, + BatchOpenMitoRegionSnafu, CleanUpMitoRegionSnafu, NoOpenRegionResultSnafu, OpenMitoRegionSnafu, PhysicalRegionNotFoundSnafu, Result, }; use crate::metrics::{LOGICAL_REGION_COUNT, PHYSICAL_REGION_COUNT}; @@ -105,7 +108,7 @@ impl MetricEngineInner { utils::to_metadata_region_id(physical_region_id), utils::to_data_region_id(physical_region_id) ); - if let Err(err) = self.close_physical_region(physical_region_id).await { + if let Err(err) = self.close_physical_region(physical_region_id, false).await { error!(err; "Failed to close physical region {}", physical_region_id); } } @@ -200,6 +203,91 @@ impl MetricEngineInner { } } + pub async fn clean_up_region( + &self, + region_id: RegionId, + request: RegionCleanUpRequest, + ) -> Result { + if request.is_physical_table() { + self.cleanup_physical_region_offline(region_id, request) + .await + } else { + Ok(0) + } + } + + async fn cleanup_physical_region_offline( + &self, + region_id: RegionId, + request: RegionCleanUpRequest, + ) -> Result { + let metadata_region_id = utils::to_metadata_region_id(region_id); + let data_region_id = utils::to_data_region_id(region_id); + let (clean_up_metadata_region_request, clean_up_data_region_request) = + self.transform_clean_up_physical_region_request(request); + let _ = self + .mito + .handle_request( + metadata_region_id, + RegionRequest::CleanUp(clean_up_metadata_region_request), + ) + .await + .context(CleanUpMitoRegionSnafu { + region_type: "metadata", + })?; + let data_region_response = self + .mito + .handle_request( + data_region_id, + RegionRequest::CleanUp(clean_up_data_region_request), + ) + .await + .context(CleanUpMitoRegionSnafu { + region_type: "data", + })?; + + if self.state.read().unwrap().exist_physical_region(region_id) { + self.state + .write() + .unwrap() + .remove_physical_region(region_id)?; + PHYSICAL_REGION_COUNT.dec(); + } + + Ok(data_region_response.affected_rows) + } + + /// Transform the cleanup request to metadata region and data region cleanup requests. + /// + /// Returns: + /// - The cleanup request for metadata region. + /// - The cleanup request for data region. + fn transform_clean_up_physical_region_request( + &self, + request: RegionCleanUpRequest, + ) -> (RegionCleanUpRequest, RegionCleanUpRequest) { + let clean_up_metadata_region_request = RegionCleanUpRequest { + table_dir: request.table_dir.clone(), + path_type: PathType::Metadata, + options: region_options_for_metadata_region(&request.options), + engine: MITO_ENGINE_NAME.to_string(), + }; + + let mut data_region_options = request.options; + set_data_region_options(&mut data_region_options); + let clean_up_data_region_request = RegionCleanUpRequest { + table_dir: request.table_dir, + path_type: PathType::Data, + options: data_region_options, + engine: MITO_ENGINE_NAME.to_string(), + }; + + ( + clean_up_metadata_region_request, + clean_up_data_region_request, + ) + } + /// Transform the open request to open metadata region and data region. /// /// Returns: @@ -226,10 +314,7 @@ impl MetricEngineInner { }; let mut data_region_options = request.options; - set_data_region_options( - &mut data_region_options, - self.config.sparse_primary_key_encoding, - ); + set_data_region_options(&mut data_region_options); let open_data_region_request = RegionOpenRequest { table_dir: request.table_dir.clone(), path_type: PathType::Data, diff --git a/src/metric-engine/src/engine/options.rs b/src/metric-engine/src/engine/options.rs index 7008325bc54..a952ffab82d 100644 --- a/src/metric-engine/src/engine/options.rs +++ b/src/metric-engine/src/engine/options.rs @@ -60,10 +60,7 @@ pub enum IndexOptions { } /// Sets data region specific options. -pub fn set_data_region_options( - options: &mut HashMap, - sparse_primary_key_encoding_if_absent: bool, -) { +pub fn set_data_region_options(options: &mut HashMap) { options.remove(METRIC_ENGINE_INDEX_TYPE_OPTION); options.remove(METRIC_ENGINE_INDEX_SKIPPING_INDEX_GRANULARITY_OPTION); options.remove(METRIC_ENGINE_INDEX_SKIPPING_INDEX_FALSE_POSITIVE_RATE_OPTION); @@ -83,11 +80,11 @@ pub fn set_data_region_options( options.insert("sst_format".to_string(), "flat".to_string()); // Decide the top-level primary key encoding: caller-supplied top-level key wins, - // then extracted legacy value, then the `sparse` default if requested. + // then extracted legacy value, then the `sparse` default. if !options.contains_key(PRIMARY_KEY_ENCODING) { if let Some(encoding) = legacy_encoding { options.insert(PRIMARY_KEY_ENCODING.to_string(), encoding); - } else if sparse_primary_key_encoding_if_absent { + } else { options.insert(PRIMARY_KEY_ENCODING.to_string(), "sparse".to_string()); } } @@ -175,7 +172,7 @@ mod tests { METRIC_ENGINE_INDEX_SKIPPING_INDEX_FALSE_POSITIVE_RATE_OPTION.to_string(), "0.01".to_string(), ); - set_data_region_options(&mut options, false); + set_data_region_options(&mut options); for key in [ METRIC_ENGINE_INDEX_TYPE_OPTION, @@ -227,7 +224,7 @@ mod tests { fn test_set_data_region_options_default_compaction_time_window() { // Test that default time window is set when not specified let mut options = HashMap::new(); - set_data_region_options(&mut options, false); + set_data_region_options(&mut options); assert_eq!(options.get("memtable.type"), Some(&"bulk".to_string())); assert_eq!(options.get("sst_format"), Some(&"flat".to_string())); @@ -241,7 +238,7 @@ mod tests { #[test] fn test_set_data_region_options_sparse_primary_key_encoding() { let mut options = HashMap::new(); - set_data_region_options(&mut options, true); + set_data_region_options(&mut options); assert_eq!(options.get("memtable.type"), Some(&"bulk".to_string())); assert_eq!(options.get("sst_format"), Some(&"flat".to_string())); @@ -264,7 +261,7 @@ mod tests { "memtable.partition_tree.index_max_keys_per_shard".to_string(), "2048".to_string(), ); - set_data_region_options(&mut options, false); + set_data_region_options(&mut options); assert_eq!(options.get("memtable.type"), Some(&"bulk".to_string())); assert_eq!(options.get("sst_format"), Some(&"flat".to_string())); @@ -282,7 +279,7 @@ mod tests { let mut options = HashMap::new(); options.insert(PRIMARY_KEY_ENCODING.to_string(), "dense".to_string()); // Sparse flag is on but caller already specified dense. - set_data_region_options(&mut options, true); + set_data_region_options(&mut options); assert_eq!( options.get(PRIMARY_KEY_ENCODING), @@ -296,7 +293,7 @@ mod tests { let mut options = HashMap::new(); options.insert(TWCS_TIME_WINDOW.to_string(), "2h".to_string()); options.insert(COMPACTION_TYPE.to_string(), "twcs".to_string()); - set_data_region_options(&mut options, false); + set_data_region_options(&mut options); // User's time window should be preserved assert_eq!(options.get(TWCS_TIME_WINDOW), Some(&"2h".to_string())); diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index 33f92ec1ed2..90007621519 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/src/metric-engine/src/engine/put.rs @@ -152,7 +152,7 @@ impl MetricEngineInner { PrimaryKeyEncoding::Sparse => self.merge_sparse_batch(physical_region_id, requests)?, PrimaryKeyEncoding::Dense => self.merge_dense_batch(data_region_id, requests)?, }; - if self.config.experimental_enable_metric_value_split { + if self.enable_metric_value_split { split_metric_values(&mut merged_request.rows); } @@ -454,7 +454,7 @@ impl MetricEngineInner { &mut request.rows, primary_key_encoding, )?; - if self.config.experimental_enable_metric_value_split { + if self.enable_metric_value_split { split_metric_values(&mut request.rows); } if primary_key_encoding == PrimaryKeyEncoding::Sparse { diff --git a/src/metric-engine/src/engine/sync/region.rs b/src/metric-engine/src/engine/sync/region.rs index d1f92bef64b..c7b3ac8682c 100644 --- a/src/metric-engine/src/engine/sync/region.rs +++ b/src/metric-engine/src/engine/sync/region.rs @@ -304,7 +304,7 @@ mod tests { metric_engine .handle_request( target_physical_region_id, - RegionRequest::Close(RegionCloseRequest {}), + RegionRequest::Close(RegionCloseRequest::default()), ) .await .unwrap(); diff --git a/src/metric-engine/src/error.rs b/src/metric-engine/src/error.rs index 59aee77b67b..316588e6b31 100644 --- a/src/metric-engine/src/error.rs +++ b/src/metric-engine/src/error.rs @@ -43,6 +43,14 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to clean up mito region, region type: {}", region_type))] + CleanUpMitoRegion { + region_type: String, + source: BoxedError, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to batch open mito region"))] BatchOpenMitoRegion { source: BoxedError, @@ -438,6 +446,7 @@ impl ErrorExt for Error { CreateMitoRegion { source, .. } | OpenMitoRegion { source, .. } + | CleanUpMitoRegion { source, .. } | CloseMitoRegion { source, .. } | MitoReadOperation { source, .. } | MitoWriteOperation { source, .. } @@ -471,6 +480,7 @@ impl ErrorExt for Error { match self { CreateMitoRegion { source, .. } | OpenMitoRegion { source, .. } + | CleanUpMitoRegion { source, .. } | BatchOpenMitoRegion { source, .. } | BatchCatchupMitoRegion { source, .. } | CloseMitoRegion { source, .. } diff --git a/src/metric-engine/src/repeated_task.rs b/src/metric-engine/src/repeated_task.rs index d0299035d95..f2cfa4e0b9c 100644 --- a/src/metric-engine/src/repeated_task.rs +++ b/src/metric-engine/src/repeated_task.rs @@ -145,18 +145,13 @@ mod tests { ) } - #[tokio::test] - async fn test_flush_metadata_region_sanitize() { - let env = TestEnv::with_prefix_and_config( - "test_flush_metadata_region_sanitize", - EngineConfig { - flush_metadata_region_interval: Duration::from_secs(0), - ..Default::default() - }, - ) - .await; - let metric = env.metric(); - let config = metric.config(); + #[test] + fn test_flush_metadata_region_sanitize() { + let mut config = EngineConfig { + flush_metadata_region_interval: Duration::from_secs(0), + ..Default::default() + }; + config.sanitize(); assert_eq!( config.flush_metadata_region_interval, DEFAULT_FLUSH_METADATA_REGION_INTERVAL diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 7ae38da9f2a..e48a5758eab 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -93,6 +93,7 @@ tokio-stream.workspace = true tokio-util.workspace = true tracing.workspace = true uuid.workspace = true +zstd.workspace = true [dev-dependencies] common-function.workspace = true diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index b6adb6eb595..5622acacefa 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -16,11 +16,13 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use async_stream::try_stream; +use common_telemetry::warn; use common_time::Timestamp; use futures::{Stream, TryStreamExt}; use object_store::services::Fs; use object_store::util::{join_dir, with_instrument_layers}; use object_store::{ATOMIC_WRITE_DIR, ErrorKind, OLD_ATOMIC_WRITE_DIR, ObjectStore}; +use parquet::file::metadata::PageIndexPolicy; use smallvec::SmallVec; use snafu::ResultExt; use store_api::metadata::RegionMetadataRef; @@ -28,9 +30,9 @@ use store_api::region_request::PathType; use store_api::sst_entry::StorageSstEntry; use store_api::storage::{FileId, RegionId, SequenceNumber}; -use crate::cache::CacheManagerRef; use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey}; use crate::cache::write_cache::SstUploadRequest; +use crate::cache::{CacheManagerRef, SstMetaPreparation, prepare_sst_meta_sync}; use crate::config::{BloomFilterConfig, FulltextIndexConfig, IndexConfig, InvertedIndexConfig}; use crate::error::{ CleanDirSnafu, DeleteIndexSnafu, DeleteIndexesSnafu, DeleteSstsSnafu, OpenDalSnafu, Result, @@ -410,14 +412,46 @@ impl AccessLayer { }; // Put parquet metadata to cache manager. - if !sst_info.is_empty() { + if !sst_info.is_empty() && cache_manager.sst_meta_cache_enabled() { for sst in &sst_info { if let Some(parquet_metadata) = &sst.file_metadata { - cache_manager.put_parquet_meta_data( - RegionFileId::new(region_id, sst.file_id), - parquet_metadata.clone(), - Some(region_metadata.clone()), - ) + let file_id = RegionFileId::new(region_id, sst.file_id); + let file_path = format!( + "region_id={}, file_id={}", + file_id.region_id(), + file_id.file_id() + ); + let page_index_policy = if parquet_metadata.offset_index().is_some() { + PageIndexPolicy::Optional + } else { + PageIndexPolicy::Skip + }; + let parquet_metadata = parquet_metadata.clone(); + let region_metadata = region_metadata.clone(); + let cache_manager = cache_manager.clone(); + // Compact cache preparation is best-effort. Run the entire operation in one + // detached blocking task so it neither blocks an async worker nor delays the + // SST write. + common_runtime::spawn_blocking_global(move || { + match prepare_sst_meta_sync( + &file_path, + Arc::unwrap_or_clone(parquet_metadata), + Some(region_metadata), + page_index_policy, + ) { + Ok(SstMetaPreparation::Prepared(metadata)) => { + cache_manager.put_prepared_sst_meta(file_id, metadata, true); + } + Ok(SstMetaPreparation::DecodedOnly { encoding_error, .. }) => warn!( + encoding_error; + "Failed to encode parquet metadata for cache, file: {}", + file_path + ), + Err(err) => { + warn!(err; "Failed to cache parquet metadata for {}", file_path); + } + } + }); } } } diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 5915aa32083..27db7daa5ab 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -26,10 +26,11 @@ pub(crate) mod write_cache; use std::collections::{BTreeMap, HashMap}; use std::mem; use std::ops::Range; -use std::sync::{Arc, RwLock}; +use std::sync::{Arc, RwLock, Weak}; use bytes::Bytes; use common_base::readable_size::ReadableSize; +use common_datasource::compression::CompressionType; use common_telemetry::warn; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::record_batch::RecordBatch; @@ -41,11 +42,13 @@ use moka::notification::RemovalCause; use moka::sync::Cache; use object_store::ObjectStore; use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; -use parquet::file::metadata::{FileMetaData, PageIndexPolicy, ParquetMetaData}; +use parquet::file::metadata::{ + FileMetaData, PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader, ParquetMetaDataWriter, +}; use puffin::puffin_manager::cache::{PuffinMetadataCache, PuffinMetadataCacheRef}; use smallvec::SmallVec; use snafu::{OptionExt, ResultExt}; -use store_api::metadata::RegionMetadataRef; +use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::storage::{ConcreteDataType, FileId, RegionId, TimeSeriesRowSelector}; use crate::cache::cache_size::parquet_meta_size; @@ -54,7 +57,10 @@ use crate::cache::index::inverted_index::{InvertedIndexCache, InvertedIndexCache #[cfg(feature = "vector_index")] use crate::cache::index::vector_index::{VectorIndexCache, VectorIndexCacheRef}; use crate::cache::write_cache::WriteCacheRef; -use crate::error::{InvalidMetadataSnafu, InvalidParquetSnafu, Result, UnexpectedSnafu}; +use crate::error::{ + CompressObjectSnafu, DecompressObjectSnafu, InvalidMetadataSnafu, InvalidParquetSnafu, + JoinSnafu, ReadParquetSnafu, Result, UnexpectedSnafu, WriteParquetSnafu, +}; use crate::memtable::record_batch_estimated_size; use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, CACHE_HIT, CACHE_MISS}; use crate::read::Batch; @@ -66,6 +72,8 @@ use crate::sst::parquet::reader::MetadataCacheMetrics; /// Metrics type key for sst meta. const SST_META_TYPE: &str = "sst_meta"; +/// Metrics type key for the optional decoded SST metadata acceleration tier. +const SST_META_DECODED_TYPE: &str = "sst_meta_decoded"; /// Metrics type key for vector. const VECTOR_TYPE: &str = "vector"; /// Metrics type key for pages. @@ -152,11 +160,235 @@ impl RangeResultMemoryLimiter { #[derive(Debug)] pub(crate) struct CachedSstMeta { parquet_metadata: Arc, + parquet_metadata_size: usize, region_metadata: RegionMetadataRef, region_metadata_weight: usize, page_index_policy: PageIndexPolicy, } +/// Compact, authoritative form of one SST's metadata. +/// +/// The entry contains a zstd-compressed, self-contained Parquet metadata stream. The decoded +/// representation is held weakly to coalesce concurrent decodes without retaining memory outside +/// the cache capacity. +#[derive(Debug)] +pub(crate) struct CompactSstMeta { + encoded_metadata: Bytes, + decoded_size: usize, + region_metadata: Weak, + page_index_policy: PageIndexPolicy, + decoded: tokio::sync::Mutex>, +} + +/// Both forms produced by decoding metadata after a cache miss. +#[derive(Debug)] +pub(crate) struct PreparedSstMeta { + compact: Arc, + decoded: Arc, +} + +impl PreparedSstMeta { + pub(crate) fn decoded(&self) -> Arc { + self.decoded.clone() + } +} + +/// Result of decoding SST metadata and attempting to encode its compact cache entry. +#[derive(Debug)] +pub(crate) enum SstMetaPreparation { + /// Both cache representations are ready for admission. + Prepared(PreparedSstMeta), + /// The metadata is usable by the reader, but its compact cache encoding failed. + DecodedOnly { + decoded: Arc, + encoding_error: crate::error::Error, + }, +} + +impl SstMetaPreparation { + pub(crate) fn decoded(&self) -> Arc { + match self { + Self::Prepared(metadata) => metadata.decoded(), + Self::DecodedOnly { decoded, .. } => decoded.clone(), + } + } +} + +impl CompactSstMeta { + async fn decode(&self) -> Result> { + let mut decoded_guard = self.decoded.lock().await; + if let Some(decoded) = decoded_guard.upgrade() { + return Ok(decoded); + } + + let encoded_metadata = self.encoded_metadata.clone(); + let decoded_size = self.decoded_size; + let region_metadata = self.region_metadata.upgrade(); + let page_index_policy = self.page_index_policy; + let decoded = common_runtime::spawn_blocking_global(move || { + let bytes = zstd::bulk::decompress(&encoded_metadata, decoded_size).context( + DecompressObjectSnafu { + compress_type: CompressionType::Zstd, + path: "cached SST metadata", + }, + )?; + let bytes = Bytes::from(bytes); + let mut reader = ParquetMetaDataReader::new() + .with_column_index_policy(PageIndexPolicy::Skip) + .with_offset_index_policy(page_index_policy); + reader.try_parse(&bytes).context(ReadParquetSnafu { + path: "cached SST metadata", + })?; + let metadata = reader.finish().context(ReadParquetSnafu { + path: "cached SST metadata", + })?; + CachedSstMeta::try_new_with_page_index_policy( + "cached SST metadata", + metadata, + region_metadata, + page_index_policy, + ) + .map(Arc::new) + }) + .await + .context(JoinSnafu)??; + + *decoded_guard = Arc::downgrade(&decoded); + Ok(decoded) + } + + fn satisfies_page_index_policy(&self, requested: PageIndexPolicy) -> bool { + satisfies_page_index_policy(self.page_index_policy, requested) + } +} + +/// Decodes SST metadata without preparing a compact cache entry. +pub(crate) async fn decode_sst_meta( + file_path: &str, + parquet_metadata: ParquetMetaData, + region_metadata: Option, + page_index_policy: PageIndexPolicy, +) -> Result> { + let file_path = file_path.to_string(); + common_runtime::spawn_blocking_global(move || { + let parquet_metadata = strip_column_indexes(parquet_metadata); + CachedSstMeta::try_new_with_page_index_policy( + &file_path, + parquet_metadata, + region_metadata, + page_index_policy, + ) + .map(Arc::new) + }) + .await + .context(JoinSnafu)? +} + +/// Decodes SST metadata and attempts to encode both cache representations on the blocking runtime. +pub(crate) async fn prepare_sst_meta( + file_path: &str, + parquet_metadata: ParquetMetaData, + region_metadata: Option, + page_index_policy: PageIndexPolicy, +) -> Result { + let file_path = file_path.to_string(); + common_runtime::spawn_blocking_global(move || { + prepare_sst_meta_sync( + &file_path, + parquet_metadata, + region_metadata, + page_index_policy, + ) + }) + .await + .context(JoinSnafu)? +} + +/// Synchronously prepares SST metadata. Callers must run this on a blocking runtime. +pub(crate) fn prepare_sst_meta_sync( + file_path: &str, + parquet_metadata: ParquetMetaData, + region_metadata: Option, + page_index_policy: PageIndexPolicy, +) -> Result { + let parquet_metadata = strip_column_indexes(parquet_metadata); + let cache_encoding = encode_compact_sst_meta(file_path, &parquet_metadata); + finish_sst_meta_preparation( + file_path, + parquet_metadata, + region_metadata, + page_index_policy, + cache_encoding, + ) +} + +fn strip_column_indexes(parquet_metadata: ParquetMetaData) -> ParquetMetaData { + // Defensively discard column indexes supplied by external metadata producers. Mito only + // consumes offset indexes. + let mut builder = parquet_metadata.into_builder(); + builder.take_column_index(); + builder.build() +} + +fn encode_compact_sst_meta( + file_path: &str, + parquet_metadata: &ParquetMetaData, +) -> Result<(Bytes, usize)> { + let mut encoded = Vec::new(); + ParquetMetaDataWriter::new(&mut encoded, parquet_metadata) + .finish() + .context(WriteParquetSnafu)?; + let decoded_size = encoded.len(); + let encoded_metadata = zstd::bulk::compress(&encoded, 3).context(CompressObjectSnafu { + compress_type: CompressionType::Zstd, + path: file_path, + })?; + + // `zstd::bulk::compress` allocates for the compression upper bound and leaves the excess + // capacity in its `Vec`. Convert through a boxed slice so the retained allocation matches the + // cache weight. + Ok(( + Bytes::from(encoded_metadata.into_boxed_slice()), + decoded_size, + )) +} + +fn finish_sst_meta_preparation( + file_path: &str, + parquet_metadata: ParquetMetaData, + region_metadata: Option, + page_index_policy: PageIndexPolicy, + cache_encoding: Result<(Bytes, usize)>, +) -> Result { + let decoded = Arc::new(CachedSstMeta::try_new_with_page_index_policy( + file_path, + parquet_metadata, + region_metadata, + page_index_policy, + )?); + let (encoded_metadata, decoded_size) = match cache_encoding { + Ok(encoded) => encoded, + Err(encoding_error) => { + return Ok(SstMetaPreparation::DecodedOnly { + decoded, + encoding_error, + }); + } + }; + let compact = Arc::new(CompactSstMeta { + encoded_metadata, + decoded_size, + region_metadata: Arc::downgrade(&decoded.region_metadata), + page_index_policy, + decoded: tokio::sync::Mutex::new(Arc::downgrade(&decoded)), + }); + + Ok(SstMetaPreparation::Prepared(PreparedSstMeta { + compact, + decoded, + })) +} + impl CachedSstMeta { #[cfg(test)] pub(crate) fn try_new(file_path: &str, parquet_metadata: ParquetMetaData) -> Result { @@ -215,9 +447,11 @@ impl CachedSstMeta { // Keep the previous JSON-byte floor and charge the decoded structures as well. let region_metadata_weight = region_metadata.estimated_size().max(json.len()); let parquet_metadata = Arc::new(strip_region_metadata_from_parquet(parquet_metadata)); + let parquet_metadata_size = parquet_meta_size(&parquet_metadata); Ok(Self { parquet_metadata, + parquet_metadata_size, region_metadata, region_metadata_weight, page_index_policy, @@ -228,21 +462,30 @@ impl CachedSstMeta { self.parquet_metadata.clone() } + /// Returns the immutable parquet metadata size computed when it was decoded. + pub(crate) fn parquet_metadata_size(&self) -> usize { + self.parquet_metadata_size + } + pub(crate) fn region_metadata(&self) -> RegionMetadataRef { self.region_metadata.clone() } fn satisfies_page_index_policy(&self, requested: PageIndexPolicy) -> bool { - match requested { - PageIndexPolicy::Skip => true, - PageIndexPolicy::Optional => self.page_index_policy != PageIndexPolicy::Skip, - PageIndexPolicy::Required => self.page_index_policy == PageIndexPolicy::Required, - } + satisfies_page_index_policy(self.page_index_policy, requested) + } +} + +fn satisfies_page_index_policy(cached: PageIndexPolicy, requested: PageIndexPolicy) -> bool { + match requested { + PageIndexPolicy::Skip => true, + PageIndexPolicy::Optional => cached != PageIndexPolicy::Skip, + PageIndexPolicy::Required => cached == PageIndexPolicy::Required, } } fn infer_loaded_page_index_policy(parquet_metadata: &ParquetMetaData) -> PageIndexPolicy { - if parquet_metadata.column_index().is_some() || parquet_metadata.offset_index().is_some() { + if parquet_metadata.offset_index().is_some() { PageIndexPolicy::Optional } else { PageIndexPolicy::Skip @@ -270,12 +513,10 @@ fn strip_region_metadata_from_parquet(parquet_metadata: ParquetMetaData) -> Parq let mut builder = parquet_metadata.into_builder(); let row_groups = builder.take_row_groups(); - let column_index = builder.take_column_index(); let offset_index = builder.take_offset_index(); parquet::file::metadata::ParquetMetaDataBuilder::new(stripped_file_metadata) .set_row_groups(row_groups) - .set_column_index(column_index) .set_offset_index(offset_index) .build() } @@ -407,6 +648,16 @@ pub enum CacheStrategy { } impl CacheStrategy { + /// Returns whether the SST metadata cache is enabled for this strategy. + pub(crate) fn sst_meta_cache_enabled(&self) -> bool { + match self { + CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { + cache_manager.sst_meta_cache_enabled() + } + CacheStrategy::Disabled => false, + } + } + /// Gets fused SST metadata with cache metrics tracking. pub(crate) async fn get_sst_meta_data( &self, @@ -450,11 +701,16 @@ impl CacheStrategy { .map(|metadata| metadata.parquet_metadata()) } - /// Calls [CacheManager::put_sst_meta_data()]. - pub(crate) fn put_sst_meta_data(&self, file_id: RegionFileId, metadata: Arc) { + /// Puts compact and decoded forms of SST metadata into the shared cache capacity. + pub(crate) fn put_prepared_sst_meta( + &self, + file_id: RegionFileId, + metadata: PreparedSstMeta, + retain_decoded: bool, + ) { match self { CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { - cache_manager.put_sst_meta_data(file_id, metadata); + cache_manager.put_prepared_sst_meta(file_id, metadata, retain_decoded); } CacheStrategy::Disabled => {} } @@ -730,8 +986,10 @@ impl CacheStrategy { /// All caches are disabled by default. #[derive(Default)] pub struct CacheManager { - /// Cache for SST metadata. + /// Cache for compact, authoritative SST metadata. sst_meta_cache: Option, + /// Cache for decoded SST metadata, used only as an acceleration tier. + sst_decoded_meta_cache: Option, /// Cache for vectors. vector_cache: Option, /// Cache for SST byte ranges. @@ -777,40 +1035,82 @@ impl CacheManager { metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, ) -> Option> { - if let Some(metadata) = self.get_sst_meta_data_from_mem_cache(file_id, page_index_policy) { + let cache_key = SstMetaKey(file_id.region_id(), file_id.file_id()); + let compact = self + .get_compact_sst_meta(&cache_key) + .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); + let decoded = self + .get_decoded_sst_meta(&cache_key) + .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); + + if let Some(compact) = compact { + CACHE_HIT.with_label_values(&[SST_META_TYPE]).inc(); metrics.mem_cache_hit += 1; - return Some(metadata); + if let Some(decoded) = decoded { + CACHE_HIT.with_label_values(&[SST_META_DECODED_TYPE]).inc(); + return Some(decoded); + } + + CACHE_MISS.with_label_values(&[SST_META_DECODED_TYPE]).inc(); + match compact.decode().await { + Ok(decoded) => { + self.put_sst_meta_data(file_id, decoded.clone()); + return Some(decoded); + } + Err(err) => { + warn!(err; "Failed to decode compact SST metadata, region_id: {}, file_id: {}", file_id.region_id(), file_id.file_id()); + self.remove_parquet_meta_data(file_id); + } + } + } else if let Some(decoded) = decoded { + // Metadata produced by a new SST writer can enter the decoded tier before its compact + // representation is prepared. + CACHE_HIT.with_label_values(&[SST_META_TYPE]).inc(); + CACHE_HIT.with_label_values(&[SST_META_DECODED_TYPE]).inc(); + metrics.mem_cache_hit += 1; + return Some(decoded); + } else { + CACHE_MISS.with_label_values(&[SST_META_TYPE]).inc(); } let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet); - if let Some(write_cache) = &self.write_cache - && let Some(metadata) = write_cache - .file_cache() - .get_sst_meta_data(key, metrics, page_index_policy) + if let Some(write_cache) = &self.write_cache { + let file_cache = write_cache.file_cache(); + if self.sst_meta_cache_enabled() { + if let Some(metadata) = file_cache + .get_sst_meta_data(key, metrics, page_index_policy) + .await + { + metrics.file_cache_hit += 1; + let decoded = metadata.decoded(); + match metadata { + SstMetaPreparation::Prepared(metadata) => { + self.put_prepared_sst_meta(file_id, metadata, true); + } + SstMetaPreparation::DecodedOnly { encoding_error, .. } => { + warn!( + encoding_error; + "Failed to encode file-cached SST metadata for memory cache, region_id: {}, file_id: {}", + file_id.region_id(), + file_id.file_id() + ); + } + } + return Some(decoded); + } + } else if let Some(decoded) = file_cache + .get_decoded_sst_meta_data(key, metrics, page_index_policy) .await - { - metrics.file_cache_hit += 1; - self.put_sst_meta_data(file_id, metadata.clone()); - return Some(metadata); + { + metrics.file_cache_hit += 1; + return Some(decoded); + } } metrics.cache_miss += 1; None } - /// Gets cached [ParquetMetaData] with metrics tracking. - /// Tries in-memory cache first, then file cache, updating metrics accordingly. - pub(crate) async fn get_parquet_meta_data( - &self, - file_id: RegionFileId, - metrics: &mut MetadataCacheMetrics, - page_index_policy: PageIndexPolicy, - ) -> Option> { - self.get_sst_meta_data(file_id, metrics, page_index_policy) - .await - .map(|metadata| metadata.parquet_metadata()) - } - /// Gets cached fused SST metadata from in-memory cache. /// This method does not perform I/O. pub(crate) fn get_sst_meta_data_from_mem_cache( @@ -818,12 +1118,11 @@ impl CacheManager { file_id: RegionFileId, page_index_policy: PageIndexPolicy, ) -> Option> { - self.sst_meta_cache.as_ref().and_then(|sst_meta_cache| { - let value = sst_meta_cache.get(&SstMetaKey(file_id.region_id(), file_id.file_id())); - let value = - value.filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); - update_hit_miss(value, SST_META_TYPE) - }) + let key = SstMetaKey(file_id.region_id(), file_id.file_id()); + let value = self + .get_decoded_sst_meta(&key) + .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); + update_hit_miss(value, SST_META_DECODED_TYPE) } /// Gets cached [ParquetMetaData] from in-memory cache. @@ -838,15 +1137,65 @@ impl CacheManager { /// Puts fused SST metadata into the cache. pub(crate) fn put_sst_meta_data(&self, file_id: RegionFileId, metadata: Arc) { - if let Some(cache) = &self.sst_meta_cache { + if let Some(cache) = &self.sst_decoded_meta_cache { let key = SstMetaKey(file_id.region_id(), file_id.file_id()); CACHE_BYTES - .with_label_values(&[SST_META_TYPE]) - .add(meta_cache_weight(&key, &metadata).into()); + .with_label_values(&[SST_META_DECODED_TYPE]) + .add(decoded_meta_cache_weight(&key, &metadata).into()); cache.insert(key, metadata); } } + /// Puts a compact metadata entry and optionally retains its decoded acceleration entry. + pub(crate) fn put_prepared_sst_meta( + &self, + file_id: RegionFileId, + metadata: PreparedSstMeta, + retain_decoded: bool, + ) { + let key = SstMetaKey(file_id.region_id(), file_id.file_id()); + if let Some(cache) = &self.sst_meta_cache { + CACHE_BYTES + .with_label_values(&[SST_META_TYPE]) + .add(meta_cache_weight(&key, &metadata.compact).into()); + cache.insert(key.clone(), metadata.compact); + } + if retain_decoded { + self.put_sst_meta_data(file_id, metadata.decoded); + } + } + + fn get_compact_sst_meta(&self, key: &SstMetaKey) -> Option> { + self.sst_meta_cache.as_ref()?.get(key) + } + + fn get_decoded_sst_meta(&self, key: &SstMetaKey) -> Option> { + self.sst_decoded_meta_cache.as_ref()?.get(key) + } + + /// Gets and decodes only the compact tier without promoting into the decoded cache. + /// + /// This is used by startup preloading so inspecting already-cached metadata cannot consume the + /// decoded reservation and prematurely stop compact preloading. + pub(crate) async fn get_compact_sst_meta_data( + &self, + file_id: RegionFileId, + page_index_policy: PageIndexPolicy, + ) -> Option> { + let key = SstMetaKey(file_id.region_id(), file_id.file_id()); + let compact = self + .get_compact_sst_meta(&key) + .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy))?; + match compact.decode().await { + Ok(metadata) => Some(metadata), + Err(err) => { + warn!(err; "Failed to decode compact SST metadata, region_id: {}, file_id: {}", file_id.region_id(), file_id.file_id()); + self.remove_parquet_meta_data(file_id); + None + } + } + } + /// Puts [ParquetMetaData] into the cache. pub fn put_parquet_meta_data( &self, @@ -854,7 +1203,7 @@ impl CacheManager { metadata: Arc, region_metadata: Option, ) { - if self.sst_meta_cache.is_some() { + if self.sst_decoded_meta_cache.is_some() { let file_path = format!( "region_id={}, file_id={}", file_id.region_id(), @@ -877,17 +1226,24 @@ impl CacheManager { /// Removes [ParquetMetaData] from the cache. pub fn remove_parquet_meta_data(&self, file_id: RegionFileId) { + let key = SstMetaKey(file_id.region_id(), file_id.file_id()); if let Some(cache) = &self.sst_meta_cache { - cache.remove(&SstMetaKey(file_id.region_id(), file_id.file_id())); + cache.remove(&key); + } + if let Some(cache) = &self.sst_decoded_meta_cache { + cache.remove(&key); } } - /// Returns the total weighted size of the in-memory SST meta cache. - pub(crate) fn sst_meta_cache_weighted_size(&self) -> u64 { - self.sst_meta_cache - .as_ref() - .map(|cache| cache.weighted_size()) - .unwrap_or(0) + /// Returns whether the authoritative SST metadata tier has reached its reservation. + pub(crate) fn sst_meta_cache_is_full(&self) -> bool { + let Some(cache) = &self.sst_meta_cache else { + return true; + }; + cache + .policy() + .max_capacity() + .is_some_and(|capacity| cache.weighted_size() >= capacity) } /// Returns true if the in-memory SST meta cache is enabled. @@ -1190,9 +1546,14 @@ impl CacheManagerBuilder { /// Builds the [CacheManager]. pub fn build(self) -> CacheManager { - let sst_meta_cache = (self.sst_meta_cache_size != 0).then(|| { + // Reserve half the configured capacity for the compact authoritative tier. This prevents + // decoded acceleration entries from evicting metadata that would require storage I/O to + // recover. Both reservations together equal the configured limit. + let compact_meta_capacity = self.sst_meta_cache_size.div_ceil(2); + let decoded_meta_capacity = self.sst_meta_cache_size / 2; + let sst_meta_cache = (compact_meta_capacity != 0).then(|| { Cache::builder() - .max_capacity(self.sst_meta_cache_size) + .max_capacity(compact_meta_capacity) .weigher(meta_cache_weight) .eviction_listener(|k, v, cause| { let size = meta_cache_weight(&k, &v); @@ -1205,6 +1566,21 @@ impl CacheManagerBuilder { }) .build() }); + let sst_decoded_meta_cache = (decoded_meta_capacity != 0).then(|| { + Cache::builder() + .max_capacity(decoded_meta_capacity) + .weigher(decoded_meta_cache_weight) + .eviction_listener(|k, v, cause| { + let size = decoded_meta_cache_weight(&k, &v); + CACHE_BYTES + .with_label_values(&[SST_META_DECODED_TYPE]) + .sub(size.into()); + CACHE_EVICTION + .with_label_values(&[SST_META_DECODED_TYPE, removal_cause_str(cause)]) + .inc(); + }) + .build() + }); let vector_cache = (self.vector_cache_size != 0).then(|| { Cache::builder() .max_capacity(self.vector_cache_size) @@ -1274,6 +1650,7 @@ impl CacheManagerBuilder { }); CacheManager { sst_meta_cache, + sst_decoded_meta_cache, vector_cache, page_cache, write_cache: self.write_cache, @@ -1295,10 +1672,15 @@ impl CacheManagerBuilder { } } -fn meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { - // We ignore the size of `Arc`. - let size = - k.estimated_size() + parquet_meta_size(&v.parquet_metadata) + v.region_metadata_weight; +fn meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { + // We ignore the size of `Arc`. Region metadata is already present in the compressed Parquet + // stream and is materialized only in the decoded acceleration tier. + let size = k.estimated_size() + mem::size_of::() + v.encoded_metadata.len(); + u32::try_from(size).unwrap_or(u32::MAX) +} + +fn decoded_meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { + let size = k.estimated_size() + v.parquet_metadata_size + v.region_metadata_weight; u32::try_from(size).unwrap_or(u32::MAX) } @@ -1677,7 +2059,9 @@ impl SelectorResultValue { } /// Maps (region id, file id) to fused SST metadata. -type SstMetaCache = Cache>; +type SstMetaCache = Cache>; +/// Maps (region id, file id) to decoded SST metadata retained as an acceleration tier. +type SstDecodedMetaCache = Cache>; /// Maps [Value] to a vector that holds this value repeatedly. /// /// e.g. `"hello" => ["hello", "hello", "hello"]` @@ -1695,6 +2079,7 @@ mod tests { use api::v1::index::{BloomFilterMeta, InvertedIndexMetas}; use datatypes::schema::ColumnSchema; use datatypes::vectors::Int64Vector; + use parquet::file::page_index::offset_index::{OffsetIndexMetaData, PageLocation}; use puffin::file_metadata::FileMetadata; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::storage::ColumnId; @@ -1715,6 +2100,7 @@ mod tests { async fn test_disable_cache() { let cache = CacheManager::default(); assert!(cache.sst_meta_cache.is_none()); + assert!(cache.sst_decoded_meta_cache.is_none()); assert!(cache.vector_cache.is_none()); assert!(cache.page_cache.is_none()); @@ -1725,7 +2111,7 @@ mod tests { cache.put_parquet_meta_data(file_id, metadata, None); assert!( cache - .get_parquet_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data(file_id, &mut metrics, Default::default()) .await .is_none() ); @@ -1754,6 +2140,30 @@ mod tests { assert!(cache.write_cache().is_none()); } + #[test] + fn test_sst_meta_cache_splits_capacity() { + let cache = CacheManager::builder().sst_meta_cache_size(101).build(); + + assert_eq!( + Some(51), + cache + .sst_meta_cache + .as_ref() + .unwrap() + .policy() + .max_capacity() + ); + assert_eq!( + Some(50), + cache + .sst_decoded_meta_cache + .as_ref() + .unwrap() + .policy() + .max_capacity() + ); + } + #[tokio::test] async fn test_parquet_meta_cache() { let cache = CacheManager::builder().sst_meta_cache_size(2000).build(); @@ -1762,7 +2172,7 @@ mod tests { let file_id = RegionFileId::new(region_id, FileId::random()); assert!( cache - .get_parquet_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data(file_id, &mut metrics, Default::default()) .await .is_none() ); @@ -1787,7 +2197,7 @@ mod tests { cache.remove_parquet_meta_data(file_id); assert!( cache - .get_parquet_meta_data(file_id, &mut metrics, Default::default()) + .get_sst_meta_data(file_id, &mut metrics, Default::default()) .await .is_none() ); @@ -1810,6 +2220,109 @@ mod tests { assert!(Arc::ptr_eq(®ion_metadata, &cached.region_metadata())); } + #[tokio::test] + async fn test_compact_sst_meta_round_trip() { + let cache = CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) + .build(); + let region_id = RegionId::new(1, 1); + let file_id = RegionFileId::new(region_id, FileId::random()); + let (metadata, expected_region_metadata) = sst_parquet_meta(); + let metadata = Arc::unwrap_or_clone(metadata); + let offset_indexes = metadata + .row_groups() + .iter() + .map(|row_group| { + row_group + .columns() + .iter() + .map(|_| OffsetIndexMetaData { + page_locations: vec![PageLocation { + offset: 42, + compressed_page_size: 128, + first_row_index: 0, + }], + unencoded_byte_array_data_bytes: None, + }) + .collect() + }) + .collect(); + let metadata = metadata + .into_builder() + .set_offset_index(Some(offset_indexes)) + .build(); + let expected_rows = metadata.file_metadata().num_rows(); + let prepared = prepare_sst_meta("test.parquet", metadata, None, PageIndexPolicy::Required) + .await + .unwrap(); + let SstMetaPreparation::Prepared(prepared) = prepared else { + panic!("valid metadata should produce a compact cache entry"); + }; + + // Retain only the compact form so the lookup must exercise decompression and decoding. + cache.put_prepared_sst_meta(file_id, prepared, false); + let mut metrics = MetadataCacheMetrics::default(); + let cached = cache + .get_sst_meta_data(file_id, &mut metrics, PageIndexPolicy::Required) + .await + .unwrap(); + + assert_eq!(1, metrics.mem_cache_hit); + assert_eq!(0, metrics.cache_miss); + assert_eq!( + expected_rows, + cached.parquet_metadata.file_metadata().num_rows() + ); + assert!(cached.parquet_metadata.offset_index().is_some()); + assert_eq!(expected_region_metadata, cached.region_metadata()); + assert!( + cached + .parquet_metadata + .file_metadata() + .key_value_metadata() + .is_none_or(|key_values| key_values + .iter() + .all(|key_value| key_value.key != PARQUET_METADATA_KEY)) + ); + } + + #[test] + fn test_cache_encoding_failure_preserves_decoded_metadata() { + let (metadata, expected_region_metadata) = sst_parquet_meta(); + let expected_rows = metadata.file_metadata().num_rows(); + let cache_encoding: Result<(Bytes, usize)> = Err(UnexpectedSnafu { + reason: "injected compact metadata encoding failure", + } + .build()); + + let preparation = finish_sst_meta_preparation( + "test.parquet", + Arc::unwrap_or_clone(metadata), + None, + PageIndexPolicy::Skip, + cache_encoding, + ) + .unwrap(); + let SstMetaPreparation::DecodedOnly { + decoded, + encoding_error, + } = preparation + else { + panic!("cache encoding failure should return decoded-only metadata"); + }; + + assert_eq!( + expected_rows, + decoded.parquet_metadata.file_metadata().num_rows() + ); + assert_eq!(expected_region_metadata, decoded.region_metadata()); + assert!( + encoding_error + .to_string() + .contains("injected compact metadata encoding failure") + ); + } + #[tokio::test] async fn test_parquet_meta_cache_respects_page_index_policy() { let cache = CacheManager::builder().sst_meta_cache_size(2000).build(); @@ -1868,7 +2381,7 @@ mod tests { } #[test] - fn test_meta_cache_weight_accounts_for_decoded_region_metadata() { + fn test_decoded_meta_cache_weight_accounts_for_region_metadata() { let region_metadata = Arc::new(wide_region_metadata(128)); let json_len = region_metadata.to_json().unwrap().len(); let metadata = sst_parquet_meta_with_region_metadata(region_metadata.clone()); @@ -1879,15 +2392,17 @@ mod tests { assert!(cached.region_metadata_weight > json_len); assert_eq!( - meta_cache_weight(&key, &cached) as usize, - key.estimated_size() - + parquet_meta_size(&cached.parquet_metadata) - + cached.region_metadata_weight + decoded_meta_cache_weight(&key, &cached) as usize, + key.estimated_size() + cached.parquet_metadata_size + cached.region_metadata_weight + ); + assert_eq!( + cached.parquet_metadata_size, + parquet_meta_size(&cached.parquet_metadata) ); } #[test] - fn test_meta_cache_weight_saturates_on_overflow() { + fn test_decoded_meta_cache_weight_saturates_on_overflow() { let region_metadata = Arc::new(wide_region_metadata(1)); let metadata = sst_parquet_meta_with_region_metadata(region_metadata.clone()); let mut cached = @@ -1896,7 +2411,7 @@ mod tests { let cached = Arc::new(cached); let key = SstMetaKey(region_metadata.region_id, FileId::random()); - assert_eq!(u32::MAX, meta_cache_weight(&key, &cached)); + assert_eq!(u32::MAX, decoded_meta_cache_weight(&key, &cached)); } #[test] diff --git a/src/mito2/src/cache/file_cache.rs b/src/mito2/src/cache/file_cache.rs index 681b2f90e3f..64f2fad6f1b 100644 --- a/src/mito2/src/cache/file_cache.rs +++ b/src/mito2/src/cache/file_cache.rs @@ -34,7 +34,9 @@ use store_api::storage::{FileId, RegionId}; use tokio::sync::mpsc::{Sender, UnboundedReceiver}; use crate::access_layer::TempFileCleaner; -use crate::cache::{CachedSstMeta, FILE_TYPE, INDEX_TYPE}; +use crate::cache::{ + CachedSstMeta, FILE_TYPE, INDEX_TYPE, SstMetaPreparation, decode_sst_meta, prepare_sst_meta, +}; use crate::error::{self, OpenDalSnafu, Result}; use crate::metrics::{ CACHE_BYTES, CACHE_HIT, CACHE_MISS, WRITE_CACHE_DOWNLOAD_BYTES_TOTAL, @@ -619,35 +621,56 @@ impl FileCache { /// Get fused SST metadata from the file cache. /// If the file is not in the cache, or metadata loading/decoding fails, return None. + /// Compact cache encoding failures return decoded-only metadata to the caller. pub(crate) async fn get_sst_meta_data( &self, key: IndexKey, cache_metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, + ) -> Option { + let file_path = self.inner.cache_file_path(key); + let metadata = self + .get_parquet_meta_data(key, cache_metrics, page_index_policy) + .await?; + match prepare_sst_meta(&file_path, metadata, None, page_index_policy).await { + Ok(metadata) => Some(metadata), + Err(err) => { + CACHE_MISS + .with_label_values(&[key.file_type.metric_label()]) + .inc(); + warn!( + err; "Failed to prepare cached parquet metadata for key {:?}", + key + ); + None + } + } + } + + /// Gets decoded SST metadata without preparing an in-memory cache entry. + pub(crate) async fn get_decoded_sst_meta_data( + &self, + key: IndexKey, + cache_metrics: &mut MetadataCacheMetrics, + page_index_policy: PageIndexPolicy, ) -> Option> { let file_path = self.inner.cache_file_path(key); - self.get_parquet_meta_data(key, cache_metrics, page_index_policy) - .await - .and_then(|metadata| { - match CachedSstMeta::try_new_with_page_index_policy( - &file_path, - metadata, - None, - page_index_policy, - ) { - Ok(metadata) => Some(Arc::new(metadata)), - Err(err) => { - CACHE_MISS - .with_label_values(&[key.file_type.metric_label()]) - .inc(); - warn!( - err; "Failed to decode cached parquet metadata for key {:?}", - key - ); - None - } - } - }) + let metadata = self + .get_parquet_meta_data(key, cache_metrics, page_index_policy) + .await?; + match decode_sst_meta(&file_path, metadata, None, page_index_policy).await { + Ok(metadata) => Some(metadata), + Err(err) => { + CACHE_MISS + .with_label_values(&[key.file_type.metric_label()]) + .inc(); + warn!( + err; "Failed to decode cached parquet metadata for key {:?}", + key + ); + None + } + } } async fn get_reader(&self, file_path: &str) -> object_store::Result> { diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index 6b02a92811a..d186307c9b1 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -638,6 +638,7 @@ mod tests { .write_cache(Some(write_cache.clone())) .build(), ); + assert!(!cache_manager.sst_meta_cache_enabled()); // Create source let metadata = Arc::new(sst_region_metadata()); @@ -682,7 +683,7 @@ mod tests { let sst_info = sst_infos.remove(0); let write_parquet_metadata = sst_info.file_metadata.unwrap(); - // Read metadata from write cache + // Read metadata from write cache without preparing an in-memory metadata cache entry. let handle = sst_file_handle_with_file_id(sst_info.file_id, 0, 1000); let builder = ParquetReaderBuilder::new( data_home, diff --git a/src/mito2/src/compaction.rs b/src/mito2/src/compaction.rs index 10900b5540c..9c0c0520458 100644 --- a/src/mito2/src/compaction.rs +++ b/src/mito2/src/compaction.rs @@ -64,7 +64,9 @@ use crate::error::{ RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, RemoteCompactionSnafu, Result, TimeRangePredicateOverflowSnafu, TimeoutSnafu, }; -use crate::metrics::{COMPACTION_STAGE_ELAPSED, INFLIGHT_COMPACTION_COUNT}; +use crate::metrics::{ + COMPACTION_MEMORY_REJECTED, COMPACTION_STAGE_ELAPSED, INFLIGHT_COMPACTION_COUNT, +}; use crate::read::FlatSource; use crate::read::flat_projection::FlatProjectionMapper; use crate::read::read_columns::ReadColumns; @@ -613,8 +615,21 @@ impl CompactionScheduler { waiters }; - // Create a local compaction task. + // Check whether this local compaction can ever fit before submitting it. let estimated_bytes = estimate_compaction_bytes(&picker_output); + if let Some(limit_bytes) = self.exceeds_compaction_memory_limit(estimated_bytes) { + COMPACTION_MEMORY_REJECTED + .with_label_values(&["oversized"]) + .inc(); + warn!( + "Skip compaction for region {} because estimated memory {} bytes exceeds compaction memory limit {} bytes", + region_id, estimated_bytes, limit_bytes, + ); + for waiter in waiters { + waiter.send(Ok(0)); + } + return Ok(None); + } let cancel_handle = Arc::new(CancellationHandle::default()); let state = LocalCompactionState::new(cancel_handle.clone()); @@ -652,6 +667,15 @@ impl CompactionScheduler { ) } + fn exceeds_compaction_memory_limit(&self, estimated_bytes: u64) -> Option { + let limit_bytes = self.memory_manager.limit_bytes(); + if limit_bytes > 0 && estimated_bytes > limit_bytes { + Some(limit_bytes) + } else { + None + } + } + fn remove_region_on_failure(&mut self, region_id: RegionId, err: Arc) { // Remove this region. let Some(status) = self.region_status.remove(®ion_id) else { @@ -1551,6 +1575,65 @@ mod tests { assert!(scheduler.region_status.contains_key(®ion_id)); } + #[tokio::test] + async fn test_schedule_compaction_skips_task_exceeding_memory_limit() { + let job_scheduler = Arc::new(VecScheduler::default()); + let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone()); + let (tx, _rx) = mpsc::channel(4); + let mut scheduler = env.mock_compaction_scheduler(tx); + scheduler.memory_manager = Arc::new(new_compaction_memory_manager(1024 * 1024)); + + let mut builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let end = 1000 * 1000; + let version_control = Arc::new( + builder + .push_l0_file_with_max_row_group_size(0, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(10, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(50, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(80, end, 1024 * 1024) + .push_l0_file_with_max_row_group_size(90, end, 1024 * 1024) + .build(), + ); + let manifest_ctx = env + .mock_manifest_context(version_control.current().version.metadata.clone()) + .await; + let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager(); + schema_metadata_manager + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + kv_backend, + ) + .await; + let (output_tx, output_rx) = oneshot::channel(); + let rejected = COMPACTION_MEMORY_REJECTED.with_label_values(&["oversized"]); + let rejected_before = rejected.get(); + + let scheduled = scheduler + .schedule_compaction( + region_id, + Options::Regular(Default::default()), + &version_control, + &env.access_layer, + OptionOutputTx::from(output_tx), + &manifest_ctx, + schema_metadata_manager, + 1, + ) + .await + .unwrap(); + + assert!(!scheduled); + assert_eq!(output_rx.await.unwrap().unwrap(), 0); + assert_eq!(rejected_before + 1, rejected.get()); + assert_eq!(0, job_scheduler.num_jobs()); + assert!(!scheduler.region_status.contains_key(®ion_id)); + } + #[tokio::test] async fn test_schedule_on_finished() { common_telemetry::init_default_ut_logging(); diff --git a/src/mito2/src/compaction/window.rs b/src/mito2/src/compaction/window.rs index 49e43f41fd4..f802601df30 100644 --- a/src/mito2/src/compaction/window.rs +++ b/src/mito2/src/compaction/window.rs @@ -254,6 +254,7 @@ mod tests { merge_mode: None, sst_format: None, primary_key_encoding: None, + write_buffer_size: None, }, compaction_time_window: None, } diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index f9aaf8bb2ee..f7489709ee3 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -38,7 +38,11 @@ pub(crate) const DEFAULT_MAX_CONCURRENT_SCAN_FILES: usize = 384; // Use `1/GLOBAL_WRITE_BUFFER_SIZE_FACTOR` of OS memory as global write buffer size in default mode const GLOBAL_WRITE_BUFFER_SIZE_FACTOR: u64 = 8; /// Use `1/SST_META_CACHE_SIZE_FACTOR` of OS memory size as SST meta cache size in default mode -const SST_META_CACHE_SIZE_FACTOR: u64 = 32; +const SST_META_CACHE_SIZE_FACTOR: u64 = 8; +/// Use `1/PREFILTER_RESULT_CACHE_SIZE_FACTOR` of OS memory size as prefilter result cache size in default mode +const PREFILTER_RESULT_CACHE_SIZE_FACTOR: u64 = 32; +/// Use `1/INDEX_METADATA_CACHE_SIZE_FACTOR` of OS memory size as index metadata cache size in default mode +const INDEX_METADATA_CACHE_SIZE_FACTOR: u64 = 32; /// Use `1/MEM_CACHE_SIZE_FACTOR` of OS memory size as mem cache size in default mode const MEM_CACHE_SIZE_FACTOR: u64 = 16; /// Use `1/PAGE_CACHE_SIZE_FACTOR` of OS memory size as page cache size in default mode @@ -105,6 +109,10 @@ pub struct MitoConfig { pub global_write_buffer_size: ReadableSize, /// Global write buffer size threshold to reject write requests. pub global_write_buffer_reject_size: ReadableSize, + /// Default write buffer size for each region. Regions stall at this size and + /// reject writes at twice this size. Setting it to 0 disables both limits + /// unless the table specifies `write_buffer_size`. + pub default_region_write_buffer_size: ReadableSize, // Cache configs: /// Cache size for SST metadata. Setting it to 0 to disable the cache. @@ -202,6 +210,7 @@ impl Default for MitoConfig { auto_flush_interval: Duration::from_secs(30 * 60), global_write_buffer_size: ReadableSize::gb(1), global_write_buffer_reject_size: ReadableSize::gb(2), + default_region_write_buffer_size: ReadableSize::mb(0), sst_meta_cache_size: ReadableSize::mb(128), vector_cache_size: ReadableSize::mb(512), page_cache_size: ReadableSize::mb(512), @@ -321,9 +330,14 @@ impl MitoConfig { ); // Use 2x of global write buffer size as global write buffer reject size. let global_write_buffer_reject_size = global_write_buffer_size * 2; - // shouldn't be greater than 128MB in default mode. + // Page-index-bearing SST metadata can be much larger than footers alone. + // Keep the auto-sized default bounded, but allow a larger warm working set. let sst_meta_cache_size = cmp::min( sys_memory / SST_META_CACHE_SIZE_FACTOR, + ReadableSize::mb(512), + ); + let prefilter_result_cache_size = cmp::min( + sys_memory / PREFILTER_RESULT_CACHE_SIZE_FACTOR, ReadableSize::mb(128), ); // shouldn't be greater than 512MB in default mode. @@ -338,7 +352,7 @@ impl MitoConfig { self.selector_result_cache_size = mem_cache_size; self.range_result_cache_size = mem_cache_size; // Use a smaller cache size because prefilter result usually should be small. - self.prefilter_result_cache_size = sst_meta_cache_size; + self.prefilter_result_cache_size = prefilter_result_cache_size; self.index.adjust_buffer_and_cache_size(sys_memory); } @@ -359,6 +373,24 @@ impl MitoConfig { } } +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_adjust_sst_metadata_and_prefilter_cache_caps_independently() { + let mut config = MitoConfig::default(); + + config.adjust_buffer_and_cache_size(ReadableSize::gb(1)); + assert_eq!(ReadableSize::mb(128), config.sst_meta_cache_size); + assert_eq!(ReadableSize::mb(32), config.prefilter_result_cache_size); + + config.adjust_buffer_and_cache_size(ReadableSize::gb(64)); + assert_eq!(ReadableSize::mb(512), config.sst_meta_cache_size); + assert_eq!(ReadableSize::mb(128), config.prefilter_result_cache_size); + } +} + /// Index build mode. #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq, Default)] #[serde(rename_all = "snake_case")] @@ -466,7 +498,7 @@ impl IndexConfig { self.content_cache_size = cmp::min(self.content_cache_size, cache_size); let metadata_cache_size = cmp::min( - sys_memory / SST_META_CACHE_SIZE_FACTOR, + sys_memory / INDEX_METADATA_CACHE_SIZE_FACTOR, ReadableSize::mb(64), ); self.metadata_cache_size = cmp::min(self.metadata_cache_size, metadata_cache_size); diff --git a/src/mito2/src/engine/catchup_test.rs b/src/mito2/src/engine/catchup_test.rs index 54326cecefd..5c8e357c056 100644 --- a/src/mito2/src/engine/catchup_test.rs +++ b/src/mito2/src/engine/catchup_test.rs @@ -498,7 +498,10 @@ async fn test_catchup_with_manifest_update(factory: Option) { async fn close_region(engine: &MitoEngine, region_id: RegionId) { engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); } diff --git a/src/mito2/src/engine/close_test.rs b/src/mito2/src/engine/close_test.rs index 4c06583b0b8..c3ace2f5c42 100644 --- a/src/mito2/src/engine/close_test.rs +++ b/src/mito2/src/engine/close_test.rs @@ -12,12 +12,23 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; +use std::sync::atomic::Ordering; + +use api::v1::Rows; +use common_base::Plugins; +use common_recordbatch::RecordBatches; use store_api::region_engine::RegionEngine; use store_api::region_request::{RegionCloseRequest, RegionRequest}; -use store_api::storage::RegionId; +use store_api::storage::{RegionId, ScanRequest}; use crate::config::MitoConfig; -use crate::test_util::{CreateRequestBuilder, TestEnv}; +use crate::engine::flush_test::MockRegionHook; +use crate::engine::listener::AlterFlushListener; +use crate::engine::region_hook::RegionHookRef; +use crate::test_util::{ + CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, rows_schema, +}; #[tokio::test] async fn test_engine_close_region() { @@ -37,7 +48,10 @@ async fn test_engine_close_region_with_format(flat_format: bool) { let region_id = RegionId::new(1, 1); // It's okay to close a region doesn't exist. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -49,14 +63,217 @@ async fn test_engine_close_region_with_format(flat_format: bool) { // Close the created region. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); assert!(!engine.is_region_exists(region_id)); // It's okay to close this region again. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); } + +#[tokio::test] +async fn test_region_hook_on_close() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + + let hook = Arc::new(MockRegionHook::new()); + let plugins = Plugins::new(); + plugins.insert(hook.clone() as RegionHookRef); + + let engine = env + .create_engine_with_plugins(MitoConfig::default(), plugins) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Sanity: no lifecycle events before closing besides the create open. + assert_eq!(hook.opened_count.load(Ordering::Relaxed), 1); + assert_eq!(hook.closed_count.load(Ordering::Relaxed), 0); + assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0); + + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + + // Closing fires on_region_closed exactly once. + assert_eq!(hook.closed_count.load(Ordering::Relaxed), 1); + // Close must not be confused with drop. + assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0); + assert_eq!(hook.files_removed_count.load(Ordering::Relaxed), 0); +} + +#[tokio::test] +async fn test_engine_close_region_flush_on_close() { + let mut env = TestEnv::with_prefix("close-flush-on-close").await; + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + + put_rows( + &engine, + region_id, + Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }, + ) + .await; + assert!( + !engine + .get_region(region_id) + .unwrap() + .version() + .memtables + .is_empty() + ); + + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest { + flush_on_close: true, + }), + ) + .await + .unwrap(); + assert!(!engine.is_region_exists(region_id)); + + engine + .handle_request( + region_id, + RegionRequest::Open(store_api::region_request::RegionOpenRequest { + engine: String::new(), + table_dir: request.table_dir.clone(), + path_type: store_api::region_request::PathType::Bare, + options: request.options.clone(), + skip_wal_replay: true, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!(3, batches.iter().map(|b| b.num_rows()).sum::()); +} + +#[tokio::test] +async fn test_engine_close_region_flush_on_close_while_flushing() { + let mut env = TestEnv::with_prefix("close-flush-on-close-while-flushing").await; + let listener = Arc::new(AlterFlushListener::default()); + let engine = env + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + + put_rows( + &engine, + region_id, + Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }, + ) + .await; + + let flush_engine = engine.clone(); + let flush_task = tokio::spawn(async move { + flush_region(&flush_engine, region_id, None).await; + }); + listener.wait_flush_begin().await; + + put_rows( + &engine, + region_id, + Rows { + schema: rows_schema(&request), + rows: build_rows(3, 6), + }, + ) + .await; + + let request_count = listener.request_count(); + let close_engine = engine.clone(); + let close_task = tokio::spawn(async move { + close_engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest { + flush_on_close: true, + }), + ) + .await + .unwrap(); + }); + listener.wait_request_count(request_count + 1).await; + + let second_flush_listener = listener.clone(); + let second_flush_task = tokio::spawn(async move { + second_flush_listener.wait_flush_begin().await; + second_flush_listener.wake_flush(); + }); + listener.wake_flush(); + + flush_task.await.unwrap(); + close_task.await.unwrap(); + second_flush_task.abort(); + assert!(!engine.is_region_exists(region_id)); + + engine + .handle_request( + region_id, + RegionRequest::Open(store_api::region_request::RegionOpenRequest { + engine: String::new(), + table_dir: request.table_dir, + path_type: store_api::region_request::PathType::Bare, + options: request.options, + skip_wal_replay: true, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!(6, batches.iter().map(|b| b.num_rows()).sum::()); +} diff --git a/src/mito2/src/engine/create_test.rs b/src/mito2/src/engine/create_test.rs index 8c4fe25b727..25b6995190e 100644 --- a/src/mito2/src/engine/create_test.rs +++ b/src/mito2/src/engine/create_test.rs @@ -105,7 +105,10 @@ async fn test_engine_create_close_create_region_with_format(flat_format: bool) { .unwrap(); // Close the region. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); // Create the same region id again. diff --git a/src/mito2/src/engine/drop_test.rs b/src/mito2/src/engine/drop_test.rs index a34a5d11727..90afd66ba32 100644 --- a/src/mito2/src/engine/drop_test.rs +++ b/src/mito2/src/engine/drop_test.rs @@ -13,9 +13,11 @@ // limitations under the License. use std::sync::Arc; +use std::sync::atomic::Ordering; use std::time::Duration; use api::v1::Rows; +use common_base::Plugins; use common_meta::key::SchemaMetadataManager; use common_meta::kv_backend::KvBackendRef; use object_store::util::join_path; @@ -25,7 +27,9 @@ use store_api::storage::RegionId; use crate::config::MitoConfig; use crate::engine::MitoEngine; +use crate::engine::flush_test::MockRegionHook; use crate::engine::listener::DropListener; +use crate::engine::region_hook::RegionHookRef; use crate::test_util::{ CreateRequestBuilder, TestEnv, build_rows_for_key, flush_region, put_rows, rows_schema, }; @@ -302,3 +306,75 @@ async fn test_engine_drop_region_for_custom_store_with_format(flat_format: bool) .unwrap() ); } + +#[tokio::test] +async fn test_region_hook_on_drop() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("drop_hook").await; + + let hook = Arc::new(MockRegionHook::new()); + let plugins = Plugins::new(); + plugins.insert(hook.clone() as RegionHookRef); + + let engine = env + .create_engine_with_plugins(MitoConfig::default(), plugins) + .await; + + let region_id = RegionId::new(1, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Put some rows so the region has active data (kept in memtable; we + // intentionally do not flush so the drop GC worker finds no parquet files + // and removes the directory immediately, avoiding the multi-minute GC wait). + let rows = Rows { + schema: column_schemas, + rows: build_rows_for_key("a", 0, 2, 0), + }; + put_rows(&engine, region_id, rows).await; + + // Sanity: nothing fired yet besides the open from create. + assert_eq!(hook.opened_count.load(Ordering::Relaxed), 1); + assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0); + assert_eq!(hook.closed_count.load(Ordering::Relaxed), 0); + + // full drop (partial_drop = false) so the GC worker physically deletes the + // directory and fires on_region_files_removed. + engine + .handle_request( + region_id, + RegionRequest::Drop(RegionDropRequest { + fast_path: false, + force: false, + partial_drop: false, + }), + ) + .await + .unwrap(); + + // on_region_dropped fires inline while handling the drop request. + hook.wait_for_dropped().await; + assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 1); + // Drop must not be confused with close. + assert_eq!(hook.closed_count.load(Ordering::Relaxed), 0); + + // on_region_files_removed fires from the background GC worker once the + // directory is physically deleted. + hook.wait_for_files_removed().await; + assert_eq!(hook.files_removed_count.load(Ordering::Relaxed), 1); +} diff --git a/src/mito2/src/engine/edit_region_test.rs b/src/mito2/src/engine/edit_region_test.rs index 106156a42c0..5aa514de436 100644 --- a/src/mito2/src/engine/edit_region_test.rs +++ b/src/mito2/src/engine/edit_region_test.rs @@ -419,7 +419,10 @@ async fn test_stalled_write_fails_fast_if_region_closed_during_editing() { let close_engine = engine.clone(); let close_task = tokio::spawn(async move { close_engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await }); diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs index 4870e8c2bbe..ccac8982970 100644 --- a/src/mito2/src/engine/flush_test.rs +++ b/src/mito2/src/engine/flush_test.rs @@ -14,6 +14,7 @@ //! Flush tests for mito engine. +use std::assert_matches; use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering}; @@ -22,6 +23,7 @@ use std::time::Duration; use api::v1::Rows; use async_trait::async_trait; use common_base::Plugins; +use common_base::readable_size::ReadableSize; use common_recordbatch::RecordBatches; use common_time::util::current_time_millis; use common_wal::options::{KafkaWalOptions, WAL_OPTIONS_KEY, WalOptions}; @@ -29,17 +31,20 @@ use rstest::rstest; use rstest_reuse::{self, apply}; use store_api::ManifestVersion; use store_api::metadata::RegionMetadataRef; -use store_api::region_engine::RegionEngine; +use store_api::mito_engine_options::WRITE_BUFFER_SIZE_KEY; +use store_api::region_engine::{RegionEngine, RegionRole}; use store_api::region_request::{ - PathType, RegionCloseRequest, RegionFlushRequest, RegionOpenRequest, RegionRequest, - ReplayCheckpoint, + AlterKind, PathType, RegionAlterRequest, RegionCloseRequest, RegionFlushRequest, + RegionOpenRequest, RegionPutRequest, RegionRequest, ReplayCheckpoint, SetRegionOption, }; use store_api::storage::{RegionId, ScanRequest}; use tokio::sync::Notify; use crate::config::MitoConfig; -use crate::engine::listener::{FlushListener, StallListener}; +use crate::engine::MitoEngine; +use crate::engine::listener::{EventListener, FlushListener, StallListener}; use crate::engine::region_hook::{RegionHook, RegionHookRef, SstFileInfo}; +use crate::error::Error; use crate::manifest::action::RegionMetaActionList; use crate::test_util::{ CreateRequestBuilder, LogStoreFactory, MockWriteBufferManager, TestEnv, build_rows, @@ -50,6 +55,73 @@ use crate::test_util::{ use crate::time_provider::TimeProvider; use crate::worker::MAX_INITIAL_CHECK_DELAY_SECS; +async fn set_write_buffer_size_to_current_usage(engine: &MitoEngine, region_id: RegionId) { + let version = engine.get_region(region_id).unwrap().version(); + let memory_usage = version.memtables.mutable_usage() + version.memtables.immutables_usage(); + assert!(memory_usage > 0); + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::SetRegionOptions { + options: vec![SetRegionOption::WriteBufferSize(Some(ReadableSize( + memory_usage as u64, + )))], + }, + }), + ) + .await + .unwrap(); +} + +struct RegionFlushGate { + blocked_region_id: RegionId, + observed_region_id: RegionId, + blocked_flush_started: Notify, + resume_blocked_flush: Notify, + observed_flush_finished: Notify, +} + +impl RegionFlushGate { + fn new(blocked_region_id: RegionId, observed_region_id: RegionId) -> Self { + Self { + blocked_region_id, + observed_region_id, + blocked_flush_started: Notify::new(), + resume_blocked_flush: Notify::new(), + observed_flush_finished: Notify::new(), + } + } + + async fn wait_blocked_flush_started(&self) { + self.blocked_flush_started.notified().await; + } + + fn resume_blocked_flush(&self) { + self.resume_blocked_flush.notify_one(); + } + + async fn wait_observed_flush_finished(&self) { + self.observed_flush_finished.notified().await; + } +} + +#[async_trait] +impl EventListener for RegionFlushGate { + fn on_flush_success(&self, region_id: RegionId) { + if region_id == self.observed_region_id { + self.observed_flush_finished.notify_one(); + } + } + + async fn on_flush_begin(&self, region_id: RegionId) { + if region_id == self.blocked_region_id { + self.blocked_flush_started.notify_one(); + self.resume_blocked_flush.notified().await; + } + } +} + #[tokio::test] async fn test_manual_flush() { test_manual_flush_with_format(false).await; @@ -273,6 +345,503 @@ async fn test_write_stall_with_format(flat_format: bool) { assert_eq!(expected, batches.pretty_print().unwrap()); } +#[tokio::test] +async fn test_region_write_buffer_limit_flushes_hot_region() { + let mut env = TestEnv::new().await; + let listener = Arc::new(FlushListener::default()); + let engine = env + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) + .await; + + let hot_region_id = RegionId::new(1, 1); + let normal_region_id = RegionId::new(2, 1); + let idle_region_id = RegionId::new(3, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + hot_region_id.table_id(), + "hot_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + idle_region_id.table_id(), + "idle_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + normal_region_id.table_id(), + "normal_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + // The hot region sets its limit to its memory usage after the initial write. + let hot_request = CreateRequestBuilder::new().build(); + let hot_schema = rows_schema(&hot_request); + engine + .handle_request(hot_region_id, RegionRequest::Create(hot_request)) + .await + .unwrap(); + + // An explicit table option takes precedence over the engine-level default. + let normal_request = CreateRequestBuilder::new() + .table_dir("normal") + .insert_option(WRITE_BUFFER_SIZE_KEY, "1GiB") + .build(); + let normal_schema = rows_schema(&normal_request); + engine + .handle_request(normal_region_id, RegionRequest::Create(normal_request)) + .await + .unwrap(); + + let idle_request = CreateRequestBuilder::new() + .table_dir("idle") + .insert_option(WRITE_BUFFER_SIZE_KEY, "1B") + .build(); + let idle_schema = rows_schema(&idle_request); + engine + .handle_request(idle_region_id, RegionRequest::Create(idle_request)) + .await + .unwrap(); + + put_rows( + &engine, + hot_region_id, + Rows { + schema: hot_schema.clone(), + rows: build_rows_for_key("hot", 0, 2, 0), + }, + ) + .await; + set_write_buffer_size_to_current_usage(&engine, hot_region_id).await; + put_rows( + &engine, + normal_region_id, + Rows { + schema: normal_schema, + rows: build_rows_for_key("normal", 0, 2, 0), + }, + ) + .await; + put_rows( + &engine, + idle_region_id, + Rows { + schema: idle_schema, + rows: build_rows_for_key("idle", 0, 2, 0), + }, + ) + .await; + + // The next write triggers a flush after the first write crosses the threshold. + put_rows( + &engine, + hot_region_id, + Rows { + schema: hot_schema, + rows: build_rows_for_key("hot", 2, 2, 0), + }, + ) + .await; + listener.wait().await; + + let scanner = engine + .scanner(hot_region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!(1, scanner.num_files()); + + let scanner = engine + .scanner(normal_region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!(0, scanner.num_files()); + assert_eq!(1, scanner.num_memtables()); + + let scanner = engine + .scanner(idle_region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!(0, scanner.num_files()); + assert_eq!(1, scanner.num_memtables()); +} + +#[tokio::test] +async fn test_region_write_buffer_stall_does_not_block_other_region() { + let mut env = TestEnv::new().await; + let listener = Arc::new(StallListener::default()); + let engine = env + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) + .await; + + let stalled_region_id = RegionId::new(1, 1); + let normal_region_id = RegionId::new(2, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + stalled_region_id.table_id(), + "stalled_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + normal_region_id.table_id(), + "normal_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let stalled_request = CreateRequestBuilder::new().build(); + let stalled_schema = rows_schema(&stalled_request); + engine + .handle_request(stalled_region_id, RegionRequest::Create(stalled_request)) + .await + .unwrap(); + + let normal_request = CreateRequestBuilder::new() + .table_dir("normal") + .insert_option(WRITE_BUFFER_SIZE_KEY, "1GiB") + .build(); + let normal_schema = rows_schema(&normal_request); + engine + .handle_request(normal_region_id, RegionRequest::Create(normal_request)) + .await + .unwrap(); + + put_rows( + &engine, + stalled_region_id, + Rows { + schema: stalled_schema.clone(), + rows: build_rows_for_key("stalled", 0, 2, 0), + }, + ) + .await; + set_write_buffer_size_to_current_usage(&engine, stalled_region_id).await; + + let engine_cloned = engine.clone(); + let stalled_write = tokio::spawn(async move { + put_rows( + &engine_cloned, + stalled_region_id, + Rows { + schema: stalled_schema, + rows: build_rows_for_key("stalled", 2, 2, 0), + }, + ) + .await; + }); + + listener.wait().await; + + put_rows( + &engine, + normal_region_id, + Rows { + schema: normal_schema, + rows: build_rows_for_key("normal", 0, 2, 0), + }, + ) + .await; + + flush_region(&engine, stalled_region_id, None).await; + stalled_write.await.unwrap(); + + let scanner = engine + .scanner(normal_region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!(0, scanner.num_files()); + assert_eq!(1, scanner.num_memtables()); +} + +#[tokio::test] +async fn test_region_write_buffer_does_not_stall_follower_write() { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_region_write_buffer_size: ReadableSize(1), + ..Default::default() + }) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: build_rows_for_key("follower", 0, 2, 0), + }, + ) + .await; + + engine + .set_region_role(region_id, RegionRole::Follower) + .unwrap(); + + let err = tokio::time::timeout( + Duration::from_secs(1), + engine.handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema, + rows: build_rows_for_key("follower", 2, 4, 0), + }, + hint: None, + partition_expr_version: None, + }), + ), + ) + .await + .expect("write to follower should not remain stalled") + .unwrap_err(); + assert_matches!( + err.into_inner().as_any().downcast_ref::().unwrap(), + Error::RegionState { .. } + ); +} + +#[tokio::test] +async fn test_unrelated_flush_does_not_resume_region_stalled_write() { + let mut env = TestEnv::new().await; + let stalled_region_id = RegionId::new(1, 1); + let normal_region_id = RegionId::new(2, 1); + let listener = Arc::new(RegionFlushGate::new(stalled_region_id, normal_region_id)); + let engine = env + .create_engine_with( + MitoConfig { + max_background_flushes: 2, + ..Default::default() + }, + None, + Some(listener.clone()), + None, + ) + .await; + + env.get_schema_metadata_manager() + .register_region_table_info( + stalled_region_id.table_id(), + "stalled_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + normal_region_id.table_id(), + "normal_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let stalled_request = CreateRequestBuilder::new().build(); + let stalled_schema = rows_schema(&stalled_request); + engine + .handle_request(stalled_region_id, RegionRequest::Create(stalled_request)) + .await + .unwrap(); + + let normal_request = CreateRequestBuilder::new() + .table_dir("normal") + .insert_option(WRITE_BUFFER_SIZE_KEY, "1GiB") + .build(); + let normal_schema = rows_schema(&normal_request); + engine + .handle_request(normal_region_id, RegionRequest::Create(normal_request)) + .await + .unwrap(); + + put_rows( + &engine, + stalled_region_id, + Rows { + schema: stalled_schema.clone(), + rows: build_rows_for_key("stalled", 0, 2, 0), + }, + ) + .await; + set_write_buffer_size_to_current_usage(&engine, stalled_region_id).await; + put_rows( + &engine, + normal_region_id, + Rows { + schema: normal_schema, + rows: build_rows_for_key("normal", 0, 2, 0), + }, + ) + .await; + + let engine_cloned = engine.clone(); + let mut stalled_write = tokio::spawn(async move { + put_rows( + &engine_cloned, + stalled_region_id, + Rows { + schema: stalled_schema, + rows: build_rows_for_key("stalled", 2, 2, 0), + }, + ) + .await; + }); + + listener.wait_blocked_flush_started().await; + flush_region(&engine, normal_region_id, None).await; + listener.wait_observed_flush_finished().await; + + assert!( + tokio::time::timeout(Duration::from_secs(1), &mut stalled_write) + .await + .is_err(), + "an unrelated region flush resumed the region-stalled write" + ); + + listener.resume_blocked_flush(); + tokio::time::timeout(Duration::from_secs(10), &mut stalled_write) + .await + .expect("the stalled write should resume after its region flushes") + .unwrap(); +} + +#[tokio::test] +async fn test_region_write_buffer_rejects_only_full_region_queue() { + let mut env = TestEnv::new().await; + let hot_region_id = RegionId::new(1, 1); + let normal_region_id = RegionId::new(2, 1); + let listener = Arc::new(RegionFlushGate::new(hot_region_id, hot_region_id)); + let engine = env + .create_engine_with( + MitoConfig { + max_background_flushes: 2, + ..Default::default() + }, + None, + Some(listener.clone()), + None, + ) + .await; + + env.get_schema_metadata_manager() + .register_region_table_info( + hot_region_id.table_id(), + "hot_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + env.get_schema_metadata_manager() + .register_region_table_info( + normal_region_id.table_id(), + "normal_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let hot_request = CreateRequestBuilder::new().build(); + let hot_schema = rows_schema(&hot_request); + engine + .handle_request(hot_region_id, RegionRequest::Create(hot_request)) + .await + .unwrap(); + let normal_request = CreateRequestBuilder::new().table_dir("normal").build(); + let normal_schema = rows_schema(&normal_request); + engine + .handle_request(normal_region_id, RegionRequest::Create(normal_request)) + .await + .unwrap(); + + put_rows( + &engine, + hot_region_id, + Rows { + schema: hot_schema.clone(), + rows: build_rows_for_key("hot", 0, 2, 0), + }, + ) + .await; + set_write_buffer_size_to_current_usage(&engine, hot_region_id).await; + + let engine_cloned = engine.clone(); + let stalled_schema = hot_schema.clone(); + let stalled_write = tokio::spawn(async move { + engine_cloned + .handle_request( + hot_region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: stalled_schema, + rows: build_rows_for_key("hot", 2, 1026, 2), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + }); + + listener.wait_blocked_flush_started().await; + + put_rows( + &engine, + normal_region_id, + Rows { + schema: normal_schema, + rows: build_rows_for_key("normal", 0, 2, 0), + }, + ) + .await; + + listener.resume_blocked_flush(); + listener.wait_observed_flush_finished().await; + + let err = tokio::time::timeout(Duration::from_secs(1), stalled_write) + .await + .expect("the stalled queue should be rejected at the region hard limit") + .unwrap() + .unwrap_err(); + assert_matches!( + err.into_inner().as_any().downcast_ref::().unwrap(), + Error::RejectWrite { .. } + ); +} + #[tokio::test] async fn test_flush_empty() { test_flush_empty_with_format(false).await; @@ -447,7 +1016,10 @@ async fn test_skip_remote_wal_replay_sets_topic_latest_entry_id(factory: Option< assert_eq!(1, region.topic_latest_entry_id.load(Ordering::Relaxed)); engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -589,7 +1161,10 @@ async fn test_remote_wal_open_without_replayed_memtable_sets_topic_latest_entry_ // The empty flush updates `topic_latest_entry_id` from KafkaLogStore's topic stats. flush_region(&engine, region_id, None).await; engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -895,24 +1470,46 @@ async fn test_update_topic_latest_entry_id(factory: Option) { } #[derive(Debug)] -struct MockRegionHook { - sst_written_count: AtomicUsize, - manifest_updated_count: AtomicUsize, +pub(super) struct MockRegionHook { + pub(super) sst_written_count: AtomicUsize, + pub(super) manifest_updated_count: AtomicUsize, + pub(super) opened_count: AtomicUsize, + pub(super) closed_count: AtomicUsize, + pub(super) dropped_count: AtomicUsize, + pub(super) files_removed_count: AtomicUsize, notify: Notify, + dropped_notify: Notify, + files_removed_notify: Notify, } impl MockRegionHook { - fn new() -> Self { + pub(super) fn new() -> Self { Self { sst_written_count: AtomicUsize::new(0), manifest_updated_count: AtomicUsize::new(0), + opened_count: AtomicUsize::new(0), + closed_count: AtomicUsize::new(0), + dropped_count: AtomicUsize::new(0), + files_removed_count: AtomicUsize::new(0), notify: Notify::new(), + dropped_notify: Notify::new(), + files_removed_notify: Notify::new(), } } async fn wait_for_manifest_update(&self) { self.notify.notified().await; } + + /// Waits until `on_region_dropped` has fired at least once. + pub(super) async fn wait_for_dropped(&self) { + self.dropped_notify.notified().await; + } + + /// Waits until `on_region_files_removed` has fired at least once. + pub(super) async fn wait_for_files_removed(&self) { + self.files_removed_notify.notified().await; + } } #[async_trait] @@ -966,6 +1563,35 @@ impl RegionHook for MockRegionHook { ); self.notify.notify_one(); } + + async fn on_region_opened(&self, region_id: RegionId, _region_metadata: &RegionMetadataRef) { + self.opened_count.fetch_add(1, Ordering::Relaxed); + common_telemetry::info!("MockRegionHook::on_region_opened: region={}", region_id); + } + + async fn on_region_closed(&self, region_id: RegionId, _region_metadata: &RegionMetadataRef) { + self.closed_count.fetch_add(1, Ordering::Relaxed); + common_telemetry::info!("MockRegionHook::on_region_closed: region={}", region_id); + } + + async fn on_region_dropped(&self, region_id: RegionId, _region_metadata: &RegionMetadataRef) { + self.dropped_count.fetch_add(1, Ordering::Relaxed); + common_telemetry::info!("MockRegionHook::on_region_dropped: region={}", region_id); + self.dropped_notify.notify_one(); + } + + async fn on_region_files_removed( + &self, + region_id: RegionId, + _region_metadata: &RegionMetadataRef, + ) { + self.files_removed_count.fetch_add(1, Ordering::Relaxed); + common_telemetry::info!( + "MockRegionHook::on_region_files_removed: region={}", + region_id + ); + self.files_removed_notify.notify_one(); + } } #[tokio::test] diff --git a/src/mito2/src/engine/open_test.rs b/src/mito2/src/engine/open_test.rs index cf4bfb1ff1d..f5ee4dd19f7 100644 --- a/src/mito2/src/engine/open_test.rs +++ b/src/mito2/src/engine/open_test.rs @@ -23,20 +23,27 @@ use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_recordbatch::RecordBatches; use either::Either; +use store_api::logstore::LogStore; +use store_api::logstore::provider::Provider; use store_api::region_engine::{RegionEngine, RegionRole, SettableRegionRoleState}; use store_api::region_request::{ - PathType, RegionCloseRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, + PathType, RegionCleanUpRequest, RegionCloseRequest, RegionOpenRequest, RegionPutRequest, + RegionRequest, }; use store_api::storage::{RegionId, ScanRequest}; use tokio::sync::oneshot; use crate::compaction::compactor::{OpenCompactionRegionRequest, open_compaction_region}; use crate::config::MitoConfig; +use crate::engine::flush_test::MockRegionHook; +use crate::engine::region_hook::RegionHookRef; use crate::error; use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef}; use crate::region::options::RegionOptions; +use crate::sst::location::region_dir_from_table_dir; use crate::test_util::{ - CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, reopen_region, rows_schema, + CreateRequestBuilder, LogStoreImpl, TestEnv, build_rows, flush_region, put_rows, reopen_region, + rows_schema, }; #[tokio::test] @@ -79,6 +86,103 @@ async fn test_engine_open_empty_with_format(flat_format: bool) { assert_eq!(role, None); } +#[tokio::test] +async fn test_engine_offline_cleanup_closed_region() { + let mut env = TestEnv::with_prefix("offline-cleanup").await; + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + let path_type = request.path_type; + let options = request.options.clone(); + let region_dir = region_dir_from_table_dir(&table_dir, region_id, path_type); + let object_store = env.get_object_store().unwrap(); + + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + assert!(object_store.exists(®ion_dir).await.unwrap()); + + put_rows( + &engine, + region_id, + Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }, + ) + .await; + let Some(LogStoreImpl::RaftEngine(log_store)) = env.get_log_store() else { + unreachable!() + }; + let provider = Provider::raft_engine_provider(region_id.as_u64()); + assert!(log_store.latest_entry_id(&provider).unwrap() > 0); + + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + assert!(!engine.is_region_exists(region_id)); + + let cleanup_request = RegionCleanUpRequest { + engine: String::new(), + table_dir, + path_type, + options, + }; + for _ in 0..2 { + engine + .handle_request(region_id, RegionRequest::CleanUp(cleanup_request.clone())) + .await + .unwrap(); + } + + assert!(!engine.is_region_exists(region_id)); + assert!(!object_store.exists(®ion_dir).await.unwrap()); + assert_eq!(0, log_store.latest_entry_id(&provider).unwrap()); +} + +#[tokio::test] +async fn test_engine_offline_cleanup_rejects_opened_region() { + let mut env = TestEnv::with_prefix("offline-cleanup-opened").await; + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + let path_type = request.path_type; + let options = request.options.clone(); + let region_dir = region_dir_from_table_dir(&table_dir, region_id, path_type); + let object_store = env.get_object_store().unwrap(); + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + let err = engine + .handle_request( + region_id, + RegionRequest::CleanUp(RegionCleanUpRequest { + engine: String::new(), + table_dir, + path_type, + options, + }), + ) + .await + .unwrap_err(); + + assert_eq!(StatusCode::RegionBusy, err.status_code()); + assert!(engine.is_region_exists(region_id)); + assert!(object_store.exists(®ion_dir).await.unwrap()); +} + #[tokio::test] async fn test_engine_open_existing() { test_engine_open_existing_with_format(false).await; @@ -243,7 +347,10 @@ async fn test_engine_region_open_with_options_with_format(flat_format: bool) { // Close the region. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -304,7 +411,10 @@ async fn test_engine_region_open_with_custom_store_with_format(flat_format: bool // Close the custom region. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -626,7 +736,10 @@ async fn test_open_compaction_region_with_format(flat_format: bool) { // Close the region. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -706,7 +819,10 @@ async fn test_open_backfills_partition_expr_with_fetcher() { // close and reopen to trigger backfill in opener engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); engine @@ -740,7 +856,10 @@ async fn test_open_backfills_partition_expr_with_fetcher() { // reopen again to ensure no further changes and still Some engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); let engine = env.reopen_engine(engine, MitoConfig::default()).await; @@ -782,7 +901,10 @@ async fn test_open_keeps_none_without_fetcher() { assert!(meta.partition_expr.is_none()); engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); let engine = env.reopen_engine(engine, MitoConfig::default()).await; @@ -805,3 +927,36 @@ async fn test_open_keeps_none_without_fetcher() { let meta = engine.get_region(region_id).unwrap().metadata(); assert!(meta.partition_expr.is_none()); } + +#[tokio::test] +async fn test_region_hook_on_open() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("open_hook").await; + + let hook = Arc::new(MockRegionHook::new()); + let plugins = Plugins::new(); + plugins.insert(hook.clone() as RegionHookRef); + + let engine = env + .create_engine_with_plugins(MitoConfig::default(), plugins) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Creating a region fires on_region_opened exactly once. + assert_eq!(hook.opened_count.load(Ordering::Relaxed), 1); + assert_eq!(hook.closed_count.load(Ordering::Relaxed), 0); + + // `reopen_region` closes then opens; the open must fire the hook again. + reopen_region(&engine, region_id, table_dir, false, Default::default()).await; + + assert_eq!(hook.opened_count.load(Ordering::Relaxed), 2); + assert_eq!(hook.closed_count.load(Ordering::Relaxed), 1); + assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0); +} diff --git a/src/mito2/src/engine/region_hook.rs b/src/mito2/src/engine/region_hook.rs index d1b28f6ebae..9c5a7dfc95d 100644 --- a/src/mito2/src/engine/region_hook.rs +++ b/src/mito2/src/engine/region_hook.rs @@ -16,15 +16,24 @@ //! //! ## Design //! -//! The [`RegionHook`] trait provides two methods with clear separation of concerns: +//! The [`RegionHook`] trait observes region activity through two categories of +//! callbacks — manifest/file observation and region lifecycle: //! //! - [`on_sst_files_written`]: Fires when mito2 physically writes SST **data files**. //! Provides per-file [`SstInfo`] + [`FileMeta`]; metadata richness varies by path //! (see [`SstFileInfo`] and the coverage footnote). //! -//! - [`on_manifest_updated`]: Fires after **any** manifest write is successfully committed. -//! Receives the full [`RegionMetaActionList`] so consumers can inspect what changed -//! (file additions/removals, schema changes, truncation, partition expression changes, etc.). +//! - [`on_manifest_updated`]: Fires after a manifest write is committed to the **live** +//! (normal) manifest directory. Writes to the staging directory (enter staging, +//! operations during staging, the intermediate apply-staging edit) are suppressed — +//! their effects are accumulated and delivered in a single notification when the +//! staged actions are promoted to the live manifest. Receives the full +//! [`RegionMetaActionList`] so consumers can inspect what changed (file additions/ +//! removals, schema changes, truncation, partition expression changes, etc.). +//! +//! - [`on_region_opened`] / [`on_region_closed`] / [`on_region_dropped`] / [`on_region_files_removed`]: +//! Region **lifecycle** callbacks for open, close, logical drop, and physical file removal. +//! See [Region lifecycle](#region-lifecycle) below. //! //! Hook implementations are registered via the [`Plugins`](common_base::Plugins) system: //! ```ignore @@ -33,6 +42,12 @@ //! //! ## Coverage //! +//! Only manifest writes to the **normal** (live) manifest directory trigger +//! `on_manifest_updated`. Writes to the staging manifest directory (operations +//! that happen while the region is in staging mode) are intentionally suppressed: +//! their effects are accumulated and delivered in a single notification when +//! the staged actions are promoted to the live manifest via `exit_staging_on_success`. +//! //! | Scenario | `on_sst_files_written` | `on_manifest_updated` | //! |------------------------------|:----------------------:|:---------------------:| //! | Flush (memtable → SST) | ✅ Yes | ✅ Yes | @@ -40,23 +55,57 @@ //! | Remote compaction | ✅ (compactor node) ¹ | ✅ (compactor node) ¹ | //! | RegionEdit / bulk ingestion | ❌ (files pre-written) | ✅ Yes | //! | Copy region | ❌ (object-store copy) | ✅ Yes | -//! | Apply staging | ❌ (delegates to edit) | ✅ Yes ² | +//! | Apply staging (promote) | ❌ (delegates to edit) | ✅ Yes ² | //! | Alter (schema change) | ❌ (no SST files) | ✅ Yes | //! | Truncate | ❌ (removes files) | ✅ Yes | -//! | Enter staging | ❌ (no SST files) | ✅ Yes | +//! | Enter staging | ❌ (no SST files) | ❌ (staging dir) | +//! | Operations during staging | N/A | ❌ (staging dir) | //! | Async index build | ❌ (index files only) | ✅ Yes | //! //! ¹ Remote compaction runs on a dedicated compactor node via `open_compaction_region()`; //! pass plugins via `OpenCompactionRegionRequest` to enable hooks there. `sst_infos` is //! `#[serde(skip)]` over the wire, so each [`SstInfo`] is rebuilt from [`FileMeta`] with //! empty footer/index — see [`SstFileInfo`] for field-level detail. -//! ² Apply staging fires `on_manifest_updated` twice: once when the staging SST files are -//! committed via `RegionEdit`, and once when `exit_staging_on_success` merges all staged -//! manifest actions into the live manifest. +//! ² Apply staging fires `on_manifest_updated` once when `exit_staging_on_success` promotes +//! all staged manifest actions (including the SST file additions) into the live manifest. +//! The intermediate staging `RegionEdit` is written to the staging directory and does not +//! fire the hook — its file list is included in the promote notification. //! //! The following paths do **not** trigger any hook: //! - Follower region sync / catchup (manifest read-only; followers don't author changes) -//! - GC / checkpoint / drop / remap (internal bookkeeping, not logical state changes) +//! - GC / checkpoint / remap (internal bookkeeping, not logical state changes) +//! +//! An explicit region **drop** does fire lifecycle hooks — see +//! [Region lifecycle](#region-lifecycle). +//! +//! ## Region lifecycle +//! +//! Beyond manifest/SST observation, the hook observes the high-level lifecycle of an +//! active region: +//! +//! | Event | Method | When | +//! |-------|--------|------| +//! | Open | [`on_region_opened`] | A create or open request registers the region as active (the counterpart to close/drop). Does not fire for the compactor's transient regions or the catch-up reopen. | +//! | Close | [`on_region_closed`] | A close request (or a close-after-flush) removes the region from the active set. Data files, manifest and WAL state are **preserved**; the region may be reopened. | +//! | Logical drop | [`on_region_dropped`] | A drop request has been handled: the region leaves the active set and its WAL entries are marked obsolete. Data files are **not yet deleted**. | +//! | Physical file removal | [`on_region_files_removed`] | The drop GC worker has deleted the region directory. Terminal file-lifecycle event. | +//! +//! Notes: +//! - `on_region_closed` / `on_region_dropped` run **inline in the region worker loop**, +//! so implementations must be fast (same contract as `on_manifest_updated`). +//! - `on_region_opened` runs inline in the worker loop on the **create** path, but on the +//! **open** path it fires inside the spawned open task (`common_runtime::spawn_global`), +//! i.e. concurrently with the worker loop — after WAL replay, before the region is +//! registered and its open request is acknowledged. Implementations must still be fast +//! and must not assume worker-loop-thread affinity or strict ordering against concurrent +//! requests to other regions. +//! - `on_region_files_removed` runs on the background drop GC task, outside the worker loop. +//! - When global GC is enabled and a normal table region is dropped with `partial_drop`, its +//! directory is left for global reclamation and `on_region_files_removed` is **not** fired +//! by the drop worker (observe it via the global GC path instead). +//! - Logical file removal (compaction, region edit, truncate) is already observable via +//! [`on_manifest_updated`] (`Edit.files_to_remove` / `Truncate` action); only the drop +//! worker's physical directory deletion needs a dedicated file hook. //! //! ## Invocation points //! @@ -77,11 +126,18 @@ //! //! ## Future work //! -//! A future `on_files_removed` hook may be added to observe file lifecycle end -//! (GC, drop, truncate, compaction removal). This is not yet implemented. +//! `on_region_files_removed` currently covers only the **drop** GC worker's physical +//! directory removal. A broader per-file `on_files_removed` hook covering compaction +//! removal, truncate, and the global GC reclamation path is not yet implemented +//! (though logical file removal is already observable via `on_manifest_updated`). +//! Role/leadership transitions (`on_region_role_changed`) are also not hooked. //! //! [`on_sst_files_written`]: RegionHook::on_sst_files_written //! [`on_manifest_updated`]: RegionHook::on_manifest_updated +//! [`on_region_opened`]: RegionHook::on_region_opened +//! [`on_region_closed`]: RegionHook::on_region_closed +//! [`on_region_dropped`]: RegionHook::on_region_dropped +//! [`on_region_files_removed`]: RegionHook::on_region_files_removed //! [`RegionManifestManager::update`]: crate::manifest::manager::RegionManifestManager::update //! [`ManifestContext::update_locked`]: crate::region::ManifestContext::update_locked //! [`ManifestContext::update_manifest`]: crate::region::ManifestContext::update_manifest @@ -103,6 +159,13 @@ use crate::sst::parquet::SstInfo; /// /// Must be [`fire`](Self::fire)d **after** the manifest write lock is released /// (the hook may read the manifest). `#[must_use]` so a forgotten receipt warns. +/// +/// ## Staging suppression +/// +/// When `is_staging` is `true`, [`fire`](Self::fire) is a no-op — the write went to +/// the staging manifest directory. The hook only observes writes to the live (normal) +/// manifest directory. Staging actions are accumulated and delivered in a single +/// notification when `exit_staging_on_success` promotes them (`is_staging = false`). #[must_use = "the region hook must be fired after releasing the manifest write lock"] pub(crate) struct PendingManifestHook { region_id: RegionId, @@ -110,6 +173,9 @@ pub(crate) struct PendingManifestHook { action_list: Option, version: ManifestVersion, hook: Option, + /// Whether the manifest write went to the staging directory. + /// When `true`, `fire()` is suppressed — the hook only observes live manifest writes. + is_staging: bool, } impl PendingManifestHook { @@ -118,12 +184,14 @@ impl PendingManifestHook { action_list: Option, version: ManifestVersion, hook: Option, + is_staging: bool, ) -> Self { Self { region_id, action_list, version, hook, + is_staging, } } @@ -132,9 +200,13 @@ impl PendingManifestHook { self.version } - /// Fires the hook if one is registered. Safe to call unconditionally: it is - /// a no-op when no hook is registered. + /// Fires the hook if one is registered, **unless** the write went to the staging + /// manifest directory (`is_staging = true`). Safe to call unconditionally: + /// it is a no-op when no hook is registered or when the write is staging-only. pub(crate) async fn fire(self) { + if self.is_staging { + return; + } if let (Some(hook), Some(action_list)) = (self.hook, self.action_list) { hook.on_manifest_updated(self.region_id, &action_list, self.version) .await; @@ -146,6 +218,9 @@ impl PendingManifestHook { /// keeps `self`'s actions followed by `other`'s, and the *later* manifest /// version wins. Used when a sequence of writes (e.g. staging-exit followed /// by metadata backfill) should notify the hook exactly once. + /// + /// `is_staging` is `true` only if **both** sides are staging; if either side + /// is a live write, the merged result is also live. pub(crate) fn merge(self, other: PendingManifestHook) -> PendingManifestHook { debug_assert_eq!( self.region_id, other.region_id, @@ -164,6 +239,7 @@ impl PendingManifestHook { }, version: self.version.max(other.version), hook: self.hook.or(other.hook), + is_staging: self.is_staging && other.is_staging, } } } @@ -211,10 +287,14 @@ pub trait RegionHook: Send + Sync + Debug { let _ = (region_id, region_metadata, files); } - /// Called after the region manifest is successfully committed. + /// Called after the region manifest is successfully committed to the **live** + /// (normal) manifest directory. /// - /// Fires for **all** manifest write paths: flush, compaction, region edit, - /// copy region, alter, truncate, enter staging, index build, etc. + /// Fires for: flush, compaction, region edit, copy region, alter, truncate, + /// async index build, and apply-staging promote. Does **not** fire for writes + /// to the staging manifest directory (enter staging, operations during staging, + /// the intermediate apply-staging edit) — those are suppressed because their + /// effects are accumulated and delivered in a single promote notification. /// /// Does **not** fire for: /// - Manifest reads / follower sync (no write) @@ -228,6 +308,80 @@ pub trait RegionHook: Send + Sync + Debug { ) { let _ = (region_id, action_list, manifest_version); } + + /// Called once a region **open** or **create** succeeds, but **before** the + /// region is registered in the engine's active set (`insert_region` runs + /// immediately afterwards). + /// + /// Fires once when a region becomes active via a create or open request — + /// the natural counterpart to [`on_region_closed`] / [`on_region_dropped`]. + /// It does **not** fire for the compactor's transient compaction regions + /// (`open_compaction_region`), nor for the internal reopen performed during + /// follower catch-up / leadership promotion. + /// + /// On the **create** path it runs inline in the region worker loop; on the + /// **open** path it runs inside the spawned open task + /// (`common_runtime::spawn_global`), concurrently with the worker loop + /// (after WAL replay, before the region is registered/acknowledged). + /// Implementations must be fast and must **not** assume worker-loop-thread + /// affinity or strict ordering against concurrent requests to other regions. + /// + /// [`on_region_closed`]: RegionHook::on_region_closed + /// [`on_region_dropped`]: RegionHook::on_region_dropped + async fn on_region_opened(&self, region_id: RegionId, region_metadata: &RegionMetadataRef) { + let _ = (region_id, region_metadata); + } + + /// Called after a region is **closed** via a close request. + /// + /// The region is removed from the engine's active set, but its data files, + /// manifest, and WAL state are **preserved**; the region may be reopened + /// later. Fires once per successful close, after the region's background + /// tasks (flush/compaction) have been stopped. + /// + /// Fires for a region of **any** role (leader or follower) that is closed. + /// Does **not** fire when a region is dropped (see [`on_region_dropped`]). + /// + /// Runs inline in the region worker loop; implementations should be fast. + /// + /// [`on_region_dropped`]: RegionHook::on_region_dropped + async fn on_region_closed(&self, region_id: RegionId, region_metadata: &RegionMetadataRef) { + let _ = (region_id, region_metadata); + } + + /// Called after a region is **logically dropped** (a drop request has been + /// handled). + /// + /// The region is removed from the active set and its WAL entries are marked + /// obsolete. Its data files are **not yet deleted** — they are scheduled for + /// asynchronous removal by the GC worker. Observe physical deletion via + /// [`on_region_files_removed`]. + /// + /// Runs inline in the region worker loop; implementations should be fast. + /// + /// [`on_region_files_removed`]: RegionHook::on_region_files_removed + async fn on_region_dropped(&self, region_id: RegionId, region_metadata: &RegionMetadataRef) { + let _ = (region_id, region_metadata); + } + + /// Called after a dropped region's data files are **physically removed** by + /// the drop GC worker (the region directory has been deleted). + /// + /// This is the terminal event in a region's file lifecycle; no further + /// callbacks fire for this region id afterwards. Fires only when the drop + /// worker itself deletes the directory. When global GC is enabled and the + /// region is a normal table region dropped with `partial_drop`, the + /// directory is left for global reclamation and this hook is **not** fired + /// by the drop worker. + /// + /// Runs on a background task, outside the region worker loop. + async fn on_region_files_removed( + &self, + region_id: RegionId, + region_metadata: &RegionMetadataRef, + ) { + let _ = (region_id, region_metadata); + } } pub type RegionHookRef = Arc; diff --git a/src/mito2/src/engine/skip_wal_test.rs b/src/mito2/src/engine/skip_wal_test.rs index eca19f22d77..2309d91569b 100644 --- a/src/mito2/src/engine/skip_wal_test.rs +++ b/src/mito2/src/engine/skip_wal_test.rs @@ -78,7 +78,10 @@ async fn test_close_region_skip_wal(insert: bool) { // Close the region. This should trigger a flush. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -146,7 +149,10 @@ async fn test_close_follower_region_skip_wal() { // Close the region. This should trigger a flush. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -214,7 +220,10 @@ async fn test_close_follower_region_skip_wal_with_pending_data() { assert!(region.is_follower()); engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); @@ -268,7 +277,10 @@ async fn test_close_region_skip_wal_while_flush_in_flight_closes_region() { let engine_cloned = engine.clone(); let close_job = tokio::spawn(async move { engine_cloned - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); }); @@ -344,7 +356,10 @@ async fn test_close_region_skip_wal_rejects_writes_queued_after_close() { let engine_cloned = engine.clone(); let close_job = tokio::spawn(async move { engine_cloned - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); }); @@ -445,7 +460,10 @@ async fn test_concurrent_close_region_skip_wal_while_flush_in_flight_succeeds() let engine_cloned = engine.clone(); let first_close = tokio::spawn(async move { engine_cloned - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); }); @@ -455,7 +473,10 @@ async fn test_concurrent_close_region_skip_wal_while_flush_in_flight_succeeds() let engine_cloned = engine.clone(); let second_close = tokio::spawn(async move { engine_cloned - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); }); @@ -518,7 +539,10 @@ async fn test_close_region_after_truncate_skip_wal() { assert!(!region.version().memtables.is_empty()); engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index f087d8a5545..71abc0b0a13 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -211,6 +211,8 @@ impl WriteBufferManager for WriteBufferManagerImpl { pub enum FlushReason { /// Engine reaches flush threshold. EngineFull, + /// Region reaches its write buffer threshold. + RegionFull, /// Manual flush. Manual, /// Flush to alter table. @@ -2107,7 +2109,7 @@ mod tests { scheduler.add_ddl_request_to_pending(SenderDdlRequest { sender: OptionOutputTx::from(sender), region_id: builder.region_id(), - request: DdlRequest::Close(store_api::region_request::RegionCloseRequest {}), + request: DdlRequest::Close(store_api::region_request::RegionCloseRequest::default()), }); let version_data = version_control.current(); diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index 1e173ab62d2..c566dd47c04 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -1236,6 +1236,7 @@ impl ManifestContext { action_list_for_hook, version, self.hook.clone(), + is_staging, )) } diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 2d5cdcb2847..9388f943ceb 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -43,8 +43,8 @@ use store_api::storage::{ColumnId, RegionId}; use tokio::sync::Semaphore; use crate::access_layer::AccessLayer; -use crate::cache::CacheManagerRef; use crate::cache::file_cache::{FileCache, FileType, IndexKey}; +use crate::cache::{CacheManagerRef, SstMetaPreparation, prepare_sst_meta}; use crate::config::MitoConfig; use crate::engine::region_hook::RegionHookRef; use crate::error; @@ -436,31 +436,7 @@ impl RegionOpener { } fn provider(&self, wal_options: &WalOptions) -> Result { - match wal_options { - WalOptions::RaftEngine => { - ensure!( - TypeId::of::() == TypeId::of::() - || TypeId::of::() == TypeId::of::(), - error::IncompatibleWalProviderChangeSnafu { - global: "`kafka`", - region: "`raft_engine`", - } - ); - Ok(Provider::raft_engine_provider(self.region_id.as_u64())) - } - WalOptions::Kafka(options) => { - ensure!( - TypeId::of::() == TypeId::of::() - || TypeId::of::() == TypeId::of::(), - error::IncompatibleWalProviderChangeSnafu { - global: "`raft_engine`", - region: "`kafka`", - } - ); - Ok(Provider::kafka_provider(options.topic.clone())) - } - WalOptions::Noop => Ok(Provider::noop_provider()), - } + provider_from_wal_options::(self.region_id, wal_options) } /// Tries to open the region and returns `None` if the region directory is empty. @@ -656,6 +632,37 @@ impl RegionOpener { } } +pub(crate) fn provider_from_wal_options( + region_id: RegionId, + wal_options: &WalOptions, +) -> Result { + match wal_options { + WalOptions::RaftEngine => { + ensure!( + TypeId::of::() == TypeId::of::() + || TypeId::of::() == TypeId::of::(), + error::IncompatibleWalProviderChangeSnafu { + global: "`kafka`", + region: "`raft_engine`", + } + ); + Ok(Provider::raft_engine_provider(region_id.as_u64())) + } + WalOptions::Kafka(options) => { + ensure!( + TypeId::of::() == TypeId::of::() + || TypeId::of::() == TypeId::of::(), + error::IncompatibleWalProviderChangeSnafu { + global: "`raft_engine`", + region: "`kafka`", + } + ); + Ok(Provider::kafka_provider(options.topic.clone())) + } + WalOptions::Noop => Ok(Provider::noop_provider()), + } +} + #[cfg(not(feature = "test-shared-fs-region-migration"))] fn supports_open_region_object_storage_requirement(object_store: &ObjectStore) -> bool { is_object_storage(object_store) @@ -1085,7 +1092,7 @@ async fn preload_parquet_meta_cache_for_files( ) -> usize { if !cache_manager.sst_meta_cache_enabled() || sst_meta_cache_capacity == 0 - || cache_manager.sst_meta_cache_weighted_size() >= sst_meta_cache_capacity + || cache_manager.sst_meta_cache_is_full() { return 0; } @@ -1097,30 +1104,60 @@ async fn preload_parquet_meta_cache_for_files( let mut loaded = 0usize; for file_handle in files { - // Stop when the shared SST meta cache is full. - if cache_manager.sst_meta_cache_weighted_size() >= sst_meta_cache_capacity { + // Stop when the protected authoritative tier is full. + if cache_manager.sst_meta_cache_is_full() { break; } let file_id = file_handle.file_id(); let mut cache_metrics = MetadataCacheMetrics::default(); if let Some(metadata) = cache_manager - .get_parquet_meta_data(file_id, &mut cache_metrics, PageIndexPolicy::Optional) + .get_compact_sst_meta_data(file_id, PageIndexPolicy::Optional) .await { if file_handle.primary_key_range().is_none() - && let Some(primary_key_range) = - extract_primary_key_range(&metadata, ®ion_metadata) + && let Some(primary_key_range) = extract_primary_key_range( + metadata.parquet_metadata().as_ref(), + ®ion_metadata, + ) { file_handle.set_primary_key_range(primary_key_range); } - // Metadata is either already in memory or loaded from file cache. - if cache_metrics.mem_cache_hit == 0 { - loaded += 1; - } continue; } + if let Some(write_cache) = cache_manager.write_cache() { + let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet); + if let Some(metadata) = write_cache + .file_cache() + .get_sst_meta_data(key, &mut cache_metrics, PageIndexPolicy::Optional) + .await + { + let decoded = metadata.decoded(); + if file_handle.primary_key_range().is_none() + && let Some(primary_key_range) = extract_primary_key_range( + decoded.parquet_metadata().as_ref(), + ®ion_metadata, + ) + { + file_handle.set_primary_key_range(primary_key_range); + } + match metadata { + SstMetaPreparation::Prepared(metadata) => { + cache_manager.put_prepared_sst_meta(file_id, metadata, false); + loaded += 1; + } + SstMetaPreparation::DecodedOnly { encoding_error, .. } => warn!( + encoding_error; + "Failed to encode file-cached SST metadata during preload, region: {}, file: {}", + region_id, + file_id.file_id() + ), + } + continue; + } + } + if !allow_direct_load { continue; } @@ -1136,8 +1173,35 @@ async fn preload_parquet_meta_cache_for_files( { file_handle.set_primary_key_range(primary_key_range); } - cache_manager.put_parquet_meta_data(file_id, Arc::new(metadata), None); - loaded += 1; + match prepare_sst_meta( + &file_path, + metadata, + // The SST can predate schema changes, so decode its embedded region metadata + // instead of substituting the region's current schema. + None, + PageIndexPolicy::Optional, + ) + .await + { + Ok(SstMetaPreparation::Prepared(metadata)) => { + // Preload the compact canonical entry only. Query traffic promotes decoded + // entries into the separate acceleration tier according to actual demand. + cache_manager.put_prepared_sst_meta(file_id, metadata, false); + loaded += 1; + } + Ok(SstMetaPreparation::DecodedOnly { encoding_error, .. }) => warn!( + encoding_error; + "Failed to encode preloaded parquet metadata, region: {}, file: {}", + region_id, + file_path + ), + Err(err) => { + warn!( + err; "Failed to prepare preloaded parquet metadata, region: {}, file: {}", + region_id, file_path + ); + } + } } Err(err) => { // Preloading is best-effort. Failure shouldn't affect region open. @@ -1250,7 +1314,7 @@ mod tests { use object_store::ObjectStore; use object_store::services::{Fs, Memory, S3}; use parquet::arrow::ArrowWriter; - use parquet::file::metadata::KeyValue; + use parquet::file::metadata::{KeyValue, PageIndexPolicy}; use parquet::file::properties::WriterProperties; use store_api::region_request::PathType; use store_api::storage::{FileId, RegionId}; @@ -1483,17 +1547,11 @@ mod tests { assert_eq!(loaded, 1); assert!( cache_manager - .get_parquet_meta_data_from_mem_cache(region_file_id) - .is_some() - ); - // The cached entry must carry the page index so that later `Optional` queries hit - // the in-memory cache instead of reloading metadata on demand. - assert!( - cache_manager - .get_sst_meta_data_from_mem_cache( + .get_compact_sst_meta_data( region_file_id, parquet::file::metadata::PageIndexPolicy::Optional, ) + .await .is_some() ); assert!(file_handle.primary_key_range().is_some()); @@ -1638,7 +1696,8 @@ mod tests { assert_eq!(loaded, 1); assert!( cache_manager - .get_parquet_meta_data_from_mem_cache(region_file_id) + .get_compact_sst_meta_data(region_file_id, PageIndexPolicy::Optional) + .await .is_some() ); } diff --git a/src/mito2/src/region/options.rs b/src/mito2/src/region/options.rs index 92dfd460fe6..2ac881c4072 100644 --- a/src/mito2/src/region/options.rs +++ b/src/mito2/src/region/options.rs @@ -107,6 +107,11 @@ pub struct RegionOptions { /// Internal primary key encoding override used by metric-engine. #[serde(skip_serializing_if = "Option::is_none")] pub primary_key_encoding: Option, + /// Per-region write buffer size. A positive size flushes/stalls this region + /// independently of the global write buffer limit and rejects writes at twice + /// the configured size; zero disables both limits. + #[serde(skip_serializing_if = "Option::is_none")] + pub write_buffer_size: Option, } impl RegionOptions { @@ -256,6 +261,7 @@ impl RegionOptions { merge_mode: options.merge_mode, sst_format, primary_key_encoding, + write_buffer_size: options.write_buffer_size, }; opts.validate()?; @@ -354,6 +360,7 @@ impl Default for TwcsOptions { #[derive(Debug, Deserialize)] #[serde(default)] struct RegionOptionsWithoutEnum { + write_buffer_size: Option, /// Region SST files TTL. ttl: Option, #[serde(with = "humantime_serde")] @@ -371,6 +378,7 @@ impl Default for RegionOptionsWithoutEnum { fn default() -> Self { let options = RegionOptions::default(); RegionOptionsWithoutEnum { + write_buffer_size: options.write_buffer_size, ttl: options.ttl, auto_flush_interval: options.auto_flush_interval, storage: options.storage, @@ -520,6 +528,7 @@ mod tests { use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_wal::options::KafkaWalOptions; + use store_api::mito_engine_options::WRITE_BUFFER_SIZE_KEY; use super::*; @@ -569,6 +578,28 @@ mod tests { ); } + #[test] + fn test_with_write_buffer_size() { + let map = make_map(&[(WRITE_BUFFER_SIZE_KEY, "128MiB")]); + let options = RegionOptions::try_from_options(RegionId::new(0, 0), &map).unwrap(); + let expect = RegionOptions { + write_buffer_size: Some(ReadableSize::mb(128)), + ..Default::default() + }; + assert_eq!(expect, options); + } + + #[test] + fn test_with_zero_write_buffer_size() { + let map = make_map(&[(WRITE_BUFFER_SIZE_KEY, "0")]); + let options = RegionOptions::try_from_options(RegionId::new(0, 0), &map).unwrap(); + let expect = RegionOptions { + write_buffer_size: Some(ReadableSize::mb(0)), + ..Default::default() + }; + assert_eq!(expect, options); + } + #[test] fn test_with_storage() { let map = make_map(&[("storage", "S3")]); @@ -898,6 +929,7 @@ mod tests { merge_mode: Some(MergeMode::LastNonNull), sst_format: Some(FormatType::Flat), primary_key_encoding: None, + write_buffer_size: None, }; assert_eq!(expect, options); } @@ -928,10 +960,18 @@ mod tests { merge_mode: Some(MergeMode::LastNonNull), sst_format: None, primary_key_encoding: None, + write_buffer_size: Some(ReadableSize::mb(128)), }; let region_options_json_str = serde_json::to_string(&options).unwrap(); let got: RegionOptions = serde_json::from_str(®ion_options_json_str).unwrap(); assert_eq!(options, got); + + let old_region_options_json_str = r#"{"ttl":null}"#; + let got: RegionOptions = serde_json::from_str(old_region_options_json_str).unwrap(); + assert_eq!(None, got.write_buffer_size); + + let default_json = serde_json::to_value(RegionOptions::default()).unwrap(); + assert!(default_json.get(WRITE_BUFFER_SIZE_KEY).is_none()); } #[test] @@ -985,6 +1025,7 @@ mod tests { merge_mode: Some(MergeMode::LastNonNull), sst_format: None, primary_key_encoding: None, + write_buffer_size: None, }; assert_eq!(options, got); } diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index e9df654f834..3fdd15db4f1 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -38,9 +38,10 @@ use store_api::region_engine::{ }; use store_api::region_request::{ AffectedRows, ApplyStagingManifestRequest, EnterStagingRequest, RegionAlterRequest, - RegionBuildIndexRequest, RegionBulkInsertsRequest, RegionCatchupRequest, RegionCloseRequest, - RegionCompactRequest, RegionCreateRequest, RegionDropRequest, RegionFlushRequest, - RegionOpenRequest, RegionRequest, RegionTruncateRequest, StagingPartitionDirective, + RegionBuildIndexRequest, RegionBulkInsertsRequest, RegionCatchupRequest, RegionCleanUpRequest, + RegionCloseRequest, RegionCompactRequest, RegionCreateRequest, RegionDropRequest, + RegionFlushRequest, RegionOpenRequest, RegionRequest, RegionTruncateRequest, + StagingPartitionDirective, }; use store_api::storage::{FileId, RegionId}; use tokio::sync::oneshot::{self, Receiver, Sender}; @@ -711,6 +712,11 @@ impl WorkerRequest { sender: sender.into(), request: DdlRequest::Open((v, None)), }), + RegionRequest::CleanUp(v) => WorkerRequest::Ddl(SenderDdlRequest { + region_id, + sender: sender.into(), + request: DdlRequest::OfflineCleanup(v), + }), RegionRequest::Close(v) => WorkerRequest::Ddl(SenderDdlRequest { region_id, sender: sender.into(), @@ -863,6 +869,7 @@ pub(crate) enum DdlRequest { Create(RegionCreateRequest), Drop(RegionDropRequest), Open((RegionOpenRequest, Option)), + OfflineCleanup(RegionCleanUpRequest), Close(RegionCloseRequest), Alter(RegionAlterRequest), Flush(RegionFlushRequest), diff --git a/src/mito2/src/sst/parquet/file_range.rs b/src/mito2/src/sst/parquet/file_range.rs index db2ce371991..52c74ef70a6 100644 --- a/src/mito2/src/sst/parquet/file_range.rs +++ b/src/mito2/src/sst/parquet/file_range.rs @@ -332,7 +332,7 @@ impl FileRangeContext { /// Returns the estimated memory size of this context. /// Mainly accounts for the parquet metadata size. pub(crate) fn memory_size(&self) -> usize { - crate::cache::cache_size::parquet_meta_size(self.reader_builder.parquet_metadata()) + self.reader_builder.parquet_metadata_size() } } diff --git a/src/mito2/src/sst/parquet/metadata.rs b/src/mito2/src/sst/parquet/metadata.rs index 441be373b0c..9d38d6e2167 100644 --- a/src/mito2/src/sst/parquet/metadata.rs +++ b/src/mito2/src/sst/parquet/metadata.rs @@ -81,7 +81,10 @@ impl<'a> MetadataLoader<'a> { let file_size = self.get_file_size().await?; let reader = ParquetMetaDataReader::new() .with_prefetch_hint(Some(DEFAULT_PREFETCH_SIZE as usize)) - .with_page_index_policy(self.page_index_policy); + // Mito uses offset indexes to translate row selections into byte ranges. It does not + // consume Parquet column indexes, so decoding them only bloats the metadata cache. + .with_column_index_policy(PageIndexPolicy::Skip) + .with_offset_index_policy(self.page_index_policy); let num_reads = AtomicUsize::new(0); let bytes_read = AtomicU64::new(0); @@ -190,6 +193,7 @@ mod tests { }; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use datatypes::arrow::record_batch::RecordBatch; + use object_store::services::Memory; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use parquet::file::metadata::{KeyValue, ParquetMetaData}; @@ -199,12 +203,12 @@ mod tests { use crate::sst::parquet::PARQUET_METADATA_KEY; use crate::test_util::sst_util::sst_region_metadata; - fn build_test_metadata( + fn build_test_parquet_bytes( include_primary_key: bool, primary_keys: &[&[u8]], row_group_sizes: &[usize], stats_enabled: EnabledStatistics, - ) -> ParquetMetaData { + ) -> Vec { let total_rows = row_group_sizes.iter().sum::(); let mut fields = vec![Field::new("field", ArrowDataType::Int64, true)]; let mut columns: Vec = @@ -253,11 +257,44 @@ mod tests { } writer.close().unwrap(); - ParquetRecordBatchReaderBuilder::try_new(Bytes::from(parquet_bytes)) - .unwrap() - .metadata() - .as_ref() - .clone() + parquet_bytes + } + + fn build_test_metadata( + include_primary_key: bool, + primary_keys: &[&[u8]], + row_group_sizes: &[usize], + stats_enabled: EnabledStatistics, + ) -> ParquetMetaData { + ParquetRecordBatchReaderBuilder::try_new(Bytes::from(build_test_parquet_bytes( + include_primary_key, + primary_keys, + row_group_sizes, + stats_enabled, + ))) + .unwrap() + .metadata() + .as_ref() + .clone() + } + + #[tokio::test] + async fn test_metadata_loader_only_loads_offset_indexes() { + let parquet_bytes = build_test_parquet_bytes(false, &[], &[4], EnabledStatistics::Page); + let file_size = parquet_bytes.len() as u64; + let file_path = "test.parquet"; + let object_store = ObjectStore::new(Memory::default()).unwrap().finish(); + object_store.write(file_path, parquet_bytes).await.unwrap(); + + let mut loader = MetadataLoader::new(object_store, file_path, file_size); + loader.with_page_index_policy(PageIndexPolicy::Required); + let metadata = loader + .load(&mut MetadataCacheMetrics::default()) + .await + .unwrap(); + + assert!(metadata.column_index().is_none()); + assert!(metadata.offset_index().is_some()); } #[test] diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 6459a9877e8..87c521181c1 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -52,7 +52,7 @@ use table::predicate::Predicate; use self::stream::{NestedSchemaAligner, ProjectedRecordBatchStream}; use crate::cache::index::result_cache::PredicateKey; -use crate::cache::{CacheStrategy, CachedSstMeta}; +use crate::cache::{CacheStrategy, CachedSstMeta, SstMetaPreparation, prepare_sst_meta}; #[cfg(feature = "vector_index")] use crate::error::ApplyVectorIndexSnafu; use crate::error::{ @@ -393,6 +393,7 @@ impl ParquetReaderBuilder { ) .await?; let mut parquet_meta = sst_meta.parquet_metadata(); + let mut parquet_metadata_size = sst_meta.parquet_metadata_size(); let region_meta = sst_meta.region_metadata(); let region_partition_expr_str = self .expected_metadata @@ -504,6 +505,7 @@ impl ParquetReaderBuilder { ) .await?; parquet_meta = sst_meta.parquet_metadata(); + parquet_metadata_size = sst_meta.parquet_metadata_size(); cache_miss |= page_index_cache_miss; } @@ -551,6 +553,7 @@ impl ParquetReaderBuilder { file_handle: self.file_handle.clone(), file_path, parquet_meta, + parquet_metadata_size, arrow_metadata, output_schema, object_store: self.object_store.clone(), @@ -695,18 +698,34 @@ impl ParquetReaderBuilder { metadata_loader.with_page_index_policy(page_index_policy); let metadata = metadata_loader.load(cache_metrics).await?; - let metadata = Arc::new(CachedSstMeta::try_new_with_page_index_policy( - file_path, - metadata, - None, - page_index_policy, - )?); - // Cache the metadata. - self.cache_strategy - .put_sst_meta_data(file_id, metadata.clone()); + let decoded = if self.cache_strategy.sst_meta_cache_enabled() { + let metadata = prepare_sst_meta(file_path, metadata, None, page_index_policy).await?; + let decoded = metadata.decoded(); + match metadata { + SstMetaPreparation::Prepared(metadata) => { + self.cache_strategy + .put_prepared_sst_meta(file_id, metadata, true); + } + SstMetaPreparation::DecodedOnly { encoding_error, .. } => { + warn!( + encoding_error; + "Failed to encode SST metadata for cache, using decoded metadata for {}", + file_path + ); + } + } + decoded + } else { + Arc::new(CachedSstMeta::try_new_with_page_index_policy( + file_path, + metadata, + None, + page_index_policy, + )?) + }; cache_metrics.metadata_load_cost += start.elapsed(); - Ok((metadata, true)) + Ok((decoded, true)) } /// Computes row groups to read, along with their respective row selections. @@ -1735,6 +1754,8 @@ pub(crate) struct RowGroupReaderBuilder { file_path: String, /// Metadata of the parquet file. parquet_meta: Arc, + /// Immutable metadata size, computed once when the footer is decoded. + parquet_metadata_size: usize, /// Arrow reader metadata for building async stream. arrow_metadata: ArrowReaderMetadata, /// Projected output schema aligned with `projection.projected_root_presence`. @@ -1777,6 +1798,10 @@ impl RowGroupReaderBuilder { &self.parquet_meta } + pub(crate) fn parquet_metadata_size(&self) -> usize { + self.parquet_metadata_size + } + pub(crate) fn cache_strategy(&self) -> &CacheStrategy { &self.cache_strategy } diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index c90c93df297..5e027884985 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -1307,7 +1307,10 @@ pub async fn reopen_region( ) { // Close the region. engine - .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) .await .unwrap(); diff --git a/src/mito2/src/test_util/version_util.rs b/src/mito2/src/test_util/version_util.rs index b14eea09581..1c0d7bfebc8 100644 --- a/src/mito2/src/test_util/version_util.rs +++ b/src/mito2/src/test_util/version_util.rs @@ -89,6 +89,15 @@ impl VersionControlBuilder { } pub(crate) fn push_l0_file(&mut self, start_ms: i64, end_ms: i64) -> &mut Self { + self.push_l0_file_with_max_row_group_size(start_ms, end_ms, 0) + } + + pub(crate) fn push_l0_file_with_max_row_group_size( + &mut self, + start_ms: i64, + end_ms: i64, + max_row_group_uncompressed_size: u64, + ) -> &mut Self { let file_id = FileId::random(); self.files.insert( file_id, @@ -101,7 +110,7 @@ impl VersionControlBuilder { ), level: 0, file_size: 0, // We don't care file size. - max_row_group_uncompressed_size: 0, + max_row_group_uncompressed_size, available_indexes: Default::default(), indexes: Default::default(), index_file_size: 0, diff --git a/src/mito2/src/wal.rs b/src/mito2/src/wal.rs index ff152a4f603..316cbbd6126 100644 --- a/src/mito2/src/wal.rs +++ b/src/mito2/src/wal.rs @@ -159,6 +159,28 @@ impl Wal { .map_err(BoxedError::new) .context(DeleteWalSnafu { region_id }) } + + /// Deletes all WAL entries in the namespace represented by `provider`. + pub async fn delete_namespace(&self, region_id: RegionId, provider: &Provider) -> Result<()> { + if let Provider::Noop = provider { + return Ok(()); + } + self.store + .delete_namespace(provider) + .await + .map_err(BoxedError::new) + .context(DeleteWalSnafu { region_id }) + } + + /// Marks all WAL entries of a region as obsolete and removes its dedicated namespace when + /// supported by the backend. + pub async fn obsolete_all(&self, region_id: RegionId, provider: &Provider) -> Result<()> { + self.store + .obsolete_all(provider, region_id) + .await + .map_err(BoxedError::new) + .context(DeleteWalSnafu { region_id }) + } } /// WAL batch writer. diff --git a/src/mito2/src/wal/raw_entry_reader.rs b/src/mito2/src/wal/raw_entry_reader.rs index de63f92251b..8a79d6c3b16 100644 --- a/src/mito2/src/wal/raw_entry_reader.rs +++ b/src/mito2/src/wal/raw_entry_reader.rs @@ -187,6 +187,14 @@ mod tests { unreachable!() } + async fn obsolete_all( + &self, + _provider: &Provider, + _region_id: RegionId, + ) -> Result<(), Self::Error> { + unreachable!() + } + fn entry( &self, _data: Vec, diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 3e1d42caf2c..a3602d29ff6 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -771,6 +771,14 @@ pub(crate) struct StalledRequests { } impl StalledRequests { + /// Returns the estimated size of stalled requests for a region. + pub(crate) fn estimated_size(&self, region_id: &RegionId) -> usize { + self.requests + .get(region_id) + .map(|(size, _, _)| *size) + .unwrap_or_default() + } + /// Appends stalled requests. pub(crate) fn append( &mut self, @@ -1148,8 +1156,13 @@ impl RegionWorkerLoop { .await; continue; } - DdlRequest::Close(_) => { - self.handle_close_request(ddl.region_id, ddl.sender).await; + DdlRequest::OfflineCleanup(req) => { + self.handle_offline_cleanup_request(ddl.region_id, req) + .await + } + DdlRequest::Close(req) => { + self.handle_close_request(ddl.region_id, req, ddl.sender) + .await; continue; } DdlRequest::Alter(req) => { diff --git a/src/mito2/src/worker/handle_alter.rs b/src/mito2/src/worker/handle_alter.rs index e022d166315..87044d768de 100644 --- a/src/mito2/src/worker/handle_alter.rs +++ b/src/mito2/src/worker/handle_alter.rs @@ -188,6 +188,20 @@ impl RegionWorkerLoop { let mut current_options = version.options.clone(); for option in options { match option { + SetRegionOption::WriteBufferSize(new_write_buffer_size) => { + info!( + "Update region write_buffer_size: {}, previous: {:?} new: {:?}", + region.region_id, current_options.write_buffer_size, new_write_buffer_size + ); + current_options.write_buffer_size = new_write_buffer_size; + current_options.validate().map_err(|e| { + store_api::metadata::InvalidRegionRequestSnafu { + region_id: region.region_id, + err: e.to_string(), + } + .build() + })?; + } SetRegionOption::Ttl(new_ttl) => { info!( "Update region ttl: {}, previous: {:?} new: {:?}", @@ -273,7 +287,8 @@ fn new_region_options_on_empty_memtable( let mut current_options = current_options.clone(); for option in options { match option { - SetRegionOption::Ttl(_) + SetRegionOption::WriteBufferSize(_) + | SetRegionOption::Ttl(_) | SetRegionOption::Twsc(_, _) | SetRegionOption::AutoFlushInterval(_) => (), SetRegionOption::Format(format_str) => { diff --git a/src/mito2/src/worker/handle_close.rs b/src/mito2/src/worker/handle_close.rs index a616a0e32dc..ef5702c8367 100644 --- a/src/mito2/src/worker/handle_close.rs +++ b/src/mito2/src/worker/handle_close.rs @@ -27,6 +27,7 @@ impl RegionWorkerLoop { pub(crate) async fn handle_close_request( &mut self, region_id: RegionId, + request: RegionCloseRequest, sender: OptionOutputTx, ) { let Some(region) = self.regions.get_region(region_id) else { @@ -36,9 +37,10 @@ impl RegionWorkerLoop { info!("Try to close region {}, worker: {}", region_id, self.id); - // If the region is using Noop WAL and has data in memtable and region is flushable (like, - // not in follower state), we should flush it before closing to ensure durability. - if region.provider == Provider::Noop + // If the close request asks for a flush, or the region is using Noop WAL, + // and has data in memtable and region is flushable (like, not in follower state), + // we should flush it before closing to ensure durability. + if (request.flush_on_close || region.provider == Provider::Noop) && !region .version_control .current() @@ -53,7 +55,7 @@ impl RegionWorkerLoop { .add_ddl_request_to_pending(SenderDdlRequest { region_id, sender, - request: DdlRequest::Close(RegionCloseRequest {}), + request: DdlRequest::Close(request), }); return; } @@ -89,5 +91,13 @@ impl RegionWorkerLoop { // clean index build status. self.index_build_scheduler.on_region_closed(region_id).await; self.region_count.dec(); + + // Notify the region hook that the region has been closed. The region is + // fully stopped and unregistered, but its files/manifest are preserved. + // Runs inline; the hook contract requires it to be fast. + if let Some(hook) = region.manifest_ctx.hook() { + let metadata = region.metadata(); + hook.on_region_closed(region_id, &metadata).await; + } } } diff --git a/src/mito2/src/worker/handle_create.rs b/src/mito2/src/worker/handle_create.rs index 6e43e74a0e4..bf254dc1e6b 100644 --- a/src/mito2/src/worker/handle_create.rs +++ b/src/mito2/src/worker/handle_create.rs @@ -86,6 +86,12 @@ impl RegionWorkerLoop { self.region_count.inc(); + // Notify the region hook that the region has been opened (created). + // Fires before registration; allocates nothing when no hook is registered. + if let Some(hook) = region.manifest_ctx.hook() { + hook.on_region_opened(region_id, ®ion.metadata()).await; + } + // Insert the MitoRegion into the RegionMap. self.regions.insert_region(region); diff --git a/src/mito2/src/worker/handle_drop.rs b/src/mito2/src/worker/handle_drop.rs index 00a96b0d166..1efa7cb99ed 100644 --- a/src/mito2/src/worker/handle_drop.rs +++ b/src/mito2/src/worker/handle_drop.rs @@ -23,12 +23,16 @@ use object_store::util::join_path; use object_store::{EntryMode, ObjectStore}; use snafu::ResultExt; use store_api::logstore::LogStore; +use store_api::metadata::RegionMetadataRef; use store_api::region_request::{AffectedRows, PathType}; use store_api::storage::RegionId; use tokio::time::sleep; +use crate::cache::CacheManagerRef; +use crate::engine::region_hook::RegionHookRef; use crate::error::{OpenDalSnafu, Result}; use crate::region::{RegionLeaderState, RegionMapRef}; +use crate::sst::index::intermediate::IntermediateManager; use crate::worker::{DROPPING_MARKER_FILE, RegionWorkerLoop}; const GC_TASK_INTERVAL_SEC: u64 = 5 * 60; // 5 minutes @@ -88,14 +92,7 @@ where ®ion.provider, ) .await?; - // Notifies flush scheduler. - self.flush_scheduler.on_region_dropped(region_id); - // Notifies compaction scheduler. - self.compaction_scheduler.on_region_dropped(region_id); - // notifies index build scheduler. - self.index_build_scheduler - .on_region_dropped(region_id) - .await; + self.cleanup_dropped_region_runtime_state(region_id).await; // Marks region version as dropped region.version_control.mark_dropped(); @@ -106,6 +103,22 @@ where self.region_count.dec(); + // Notify registered hooks that the region has been logically dropped, and + // prepare a payload for the background GC task to fire the terminal + // `on_region_files_removed`. When no hook is registered this allocates + // nothing — no metadata snapshot is taken and no payload is built. + let hook_payload = match region.manifest_ctx.hook() { + Some(hook) => { + let region_metadata = region.metadata(); + hook.on_region_dropped(region_id, ®ion_metadata).await; + Some(DropHookPayload { + hook, + metadata: region_metadata, + }) + } + None => None, + }; + let object_store = region.access_layer.object_store().clone(); let dropping_regions = self.dropping_regions.clone(); let listener = self.listener.clone(); @@ -122,6 +135,7 @@ where object_store, dropping_regions, partial_drop, + hook_payload.as_ref(), ) .await } else { @@ -135,25 +149,75 @@ where object_store, dropping_regions, gc_duration, + hook_payload.as_ref(), ) .await }; - if let Err(err) = intm_manager.prune_region_dir(®ion_id).await { - warn!(err; "Failed to prune intermediate region directory, region_id: {}", region_id); - } - - if let Some(write_cache) = cache_manager.write_cache() - && let Some(manifest_cache) = write_cache.manifest_cache() - { - manifest_cache.clean_manifests(&table_dir).await; - } + cleanup_region_file_artifacts(region_id, &table_dir, &intm_manager, &cache_manager) + .await; listener.on_later_drop_end(region_id, removed); }); Ok(0) } + + /// Cleans runtime state for a region that is no longer available to serve requests. + pub(crate) async fn cleanup_dropped_region_runtime_state(&mut self, region_id: RegionId) { + // Notifies flush scheduler. + self.flush_scheduler.on_region_dropped(region_id); + // Notifies compaction scheduler. + self.compaction_scheduler.on_region_dropped(region_id); + // Notifies index build scheduler. + self.index_build_scheduler + .on_region_dropped(region_id) + .await; + } +} + +/// Cleans files and caches that are produced at runtime but are not part of the +/// primary region directory deletion. +pub(crate) async fn cleanup_region_file_artifacts( + region_id: RegionId, + table_dir: &str, + intermediate_manager: &IntermediateManager, + cache_manager: &CacheManagerRef, +) { + if let Err(err) = intermediate_manager.prune_region_dir(®ion_id).await { + warn!(err; "Failed to prune intermediate region directory, region_id: {}", region_id); + } + + if let Some(write_cache) = cache_manager.write_cache() + && let Some(manifest_cache) = write_cache.manifest_cache() + { + manifest_cache.clean_manifests(table_dir).await; + } +} + +/// Removes a region directory for full-drop style cleanup. +/// +/// Full drop and purge/offline cleanup force physical deletion. Only partial +/// drop may leave data files for global GC. +pub(crate) async fn remove_region_dir_for_full_drop( + region_path: &str, + object_store: &ObjectStore, +) -> Result<()> { + remove_region_dir_once(region_path, object_store, true).await?; + Ok(()) +} + +/// Carries the region hook and region metadata into the background GC task so +/// it can fire [`RegionHook::on_region_files_removed`] once the dropped region's +/// directory is physically deleted. +/// +/// Only constructed when a hook is registered; the task receives it as +/// `Option<&DropHookPayload>` so the no-hook path allocates nothing. +/// +/// [`RegionHook::on_region_files_removed`]: crate::engine::region_hook::RegionHook::on_region_files_removed +struct DropHookPayload { + hook: RegionHookRef, + metadata: RegionMetadataRef, } /// Background GC task to remove the entire region path once one of the following @@ -172,6 +236,7 @@ async fn later_drop_task_without_global_gc( object_store: ObjectStore, dropping_regions: RegionMapRef, gc_duration: Duration, + hook_payload: Option<&DropHookPayload>, ) -> bool { remove_region_with_retry( region_id, @@ -180,6 +245,7 @@ async fn later_drop_task_without_global_gc( dropping_regions, Some(gc_duration), false, + hook_payload, ) .await } @@ -191,6 +257,7 @@ async fn remove_region_with_retry( dropping_regions: std::sync::Arc, gc_duration: Option, mut force: bool, + hook_payload: Option<&DropHookPayload>, ) -> bool { for _ in 0..MAX_RETRY_TIMES { let result = remove_region_dir_once(®ion_path, &object_store, force).await; @@ -204,6 +271,15 @@ async fn remove_region_with_retry( Ok(true) => { dropping_regions.remove_region(region_id); info!("Region {} is dropped, force: {}", region_path, force); + // The region directory has been physically deleted; fire the + // terminal file-lifecycle event. The partial-drop/global-GC path + // never reaches here (it does not delete the directory itself). + if let Some(hook_payload) = hook_payload { + hook_payload + .hook + .on_region_files_removed(region_id, &hook_payload.metadata) + .await; + } return true; } Ok(false) => (), @@ -230,12 +306,13 @@ async fn later_drop_task_with_global_gc( object_store: ObjectStore, dropping_regions: RegionMapRef, partial_drop: bool, + hook_payload: Option<&DropHookPayload>, ) -> bool { // For metadata regions or regions marked for full deletion (such as when dropping a table) // the region directory is forcefully removed immediately. // // TODO(discord9): Evaluate removing files instantly rather than waiting for the GC period. - if path_type == PathType::Metadata || !partial_drop { + if should_force_remove_region_dir(path_type, partial_drop) { remove_region_with_retry( region_id, region_path, @@ -243,6 +320,7 @@ async fn later_drop_task_with_global_gc( dropping_regions, None, true, + hook_payload, ) .await } else { @@ -252,6 +330,10 @@ async fn later_drop_task_with_global_gc( } } +fn should_force_remove_region_dir(path_type: PathType, partial_drop: bool) -> bool { + path_type == PathType::Metadata || !partial_drop +} + // TODO(ruihang): place the marker in a separate dir /// Removes region dir if there is no parquet files, returns whether the directory is removed. /// If `force = true`, always removes the dir. diff --git a/src/mito2/src/worker/handle_flush.rs b/src/mito2/src/worker/handle_flush.rs index c8e86699ea1..ccbe31c3a32 100644 --- a/src/mito2/src/worker/handle_flush.rs +++ b/src/mito2/src/worker/handle_flush.rs @@ -14,9 +14,11 @@ //! Handling flush related requests. +use std::collections::HashSet; use std::sync::Arc; use std::sync::atomic::Ordering; +use common_base::readable_size::ReadableSize; use common_telemetry::{debug, error, info}; use store_api::logstore::LogStore; use store_api::region_request::{RegionFlushReason, RegionFlushRequest}; @@ -26,10 +28,24 @@ use crate::config::{IndexBuildMode, MitoConfig}; use crate::error::{RegionNotFoundSnafu, Result}; use crate::flush::{FlushReason, RegionFlushTask}; use crate::region::MitoRegionRef; +use crate::region::version::VersionRef; use crate::request::{BuildIndexRequest, FlushFailed, FlushFinished, OnFailure, OptionOutputTx}; use crate::sst::index::IndexBuildType; use crate::worker::RegionWorkerLoop; +#[derive(Debug, Default)] +pub(crate) struct RegionWriteBufferPressure { + pub(crate) stalled_region_ids: HashSet, + pub(crate) rejected_region_ids: HashSet, +} + +#[derive(Debug, Default, PartialEq, Eq)] +struct RegionWriteBufferStatus { + should_flush: bool, + should_stall: bool, + should_reject: bool, +} + fn resolve_flush_reason( request_reason: Option, is_downgrading: bool, @@ -85,8 +101,7 @@ impl RegionWorkerLoop { } let version = region.version(); - let region_memtable_size = - version.memtables.mutable_usage() + version.memtables.immutables_usage(); + let region_memtable_size = region_memtable_usage(&version); let auto_flush_interval = version .options @@ -152,6 +167,53 @@ impl RegionWorkerLoop { Ok(()) } + /// Flushes write regions that exceed their flush threshold and returns region pressure. + pub(crate) fn maybe_flush_write_regions( + &mut self, + region_ids: HashSet, + ) -> RegionWriteBufferPressure { + let mut pressure = RegionWriteBufferPressure::default(); + + for region_id in region_ids { + let Some(region) = self.regions.get_region(region_id) else { + continue; + }; + if !region.is_writable() { + continue; + } + + let status = region_write_buffer_status( + ®ion.version(), + self.config.default_region_write_buffer_size, + self.stalled_requests.estimated_size(®ion_id), + ); + + if status.should_flush && !self.flush_scheduler.is_flush_requested(region.region_id) { + let task = self.new_flush_task( + ®ion, + FlushReason::RegionFull, + None, + self.config.clone(), + ); + if let Err(e) = self.flush_scheduler.schedule_flush( + region.region_id, + ®ion.version_control, + task, + ) { + error!(e; "Failed to schedule flush task for region {}", region.region_id); + } + } + + if status.should_reject { + pressure.rejected_region_ids.insert(region_id); + } else if status.should_stall { + pressure.stalled_region_ids.insert(region_id); + } + } + + pressure + } + /// Creates a flush task with specific `reason` for the `region`. pub(crate) fn new_flush_task( &self, @@ -179,6 +241,61 @@ impl RegionWorkerLoop { } } +pub(crate) fn region_memtable_usage(version: &VersionRef) -> usize { + version.memtables.mutable_usage() + version.memtables.immutables_usage() +} + +pub(crate) fn region_write_buffer_size( + version: &VersionRef, + default_region_write_buffer_size: ReadableSize, +) -> Option { + let size = version + .options + .write_buffer_size + .unwrap_or(default_region_write_buffer_size); + (size.as_bytes() > 0).then_some(size.as_bytes() as usize) +} + +/// Returns the pressure state of a region's write buffer. +fn region_write_buffer_status( + version: &VersionRef, + default_region_write_buffer_size: ReadableSize, + stalled_request_size: usize, +) -> RegionWriteBufferStatus { + let Some(write_buffer_size) = + region_write_buffer_size(version, default_region_write_buffer_size) + else { + return RegionWriteBufferStatus::default(); + }; + + let mutable_usage = version.memtables.mutable_usage(); + let memory_usage = region_memtable_usage(version); + region_write_buffer_status_from_usage( + write_buffer_size, + mutable_usage, + memory_usage, + stalled_request_size, + ) +} + +fn region_write_buffer_status_from_usage( + write_buffer_size: usize, + mutable_usage: usize, + memory_usage: usize, + stalled_request_size: usize, +) -> RegionWriteBufferStatus { + let mutable_limit = std::cmp::max(1, write_buffer_size / 2); + let should_stall = memory_usage >= write_buffer_size; + let reject_limit = write_buffer_size.saturating_mul(2); + let should_reject = memory_usage.saturating_add(stalled_request_size) >= reject_limit; + + RegionWriteBufferStatus { + should_flush: mutable_usage >= mutable_limit || should_stall, + should_stall, + should_reject, + } +} + impl RegionWorkerLoop { /// Handles manual flush request. pub(crate) fn handle_flush_request( @@ -438,4 +555,48 @@ mod tests { assert_eq!(resolve_flush_reason(None, true), FlushReason::Downgrading); assert_eq!(resolve_flush_reason(None, false), FlushReason::Manual); } + + #[test] + fn test_region_write_buffer_status_boundaries() { + assert_eq!( + RegionWriteBufferStatus::default(), + region_write_buffer_status_from_usage(100, 49, 99, 0) + ); + assert_eq!( + RegionWriteBufferStatus { + should_flush: true, + should_stall: false, + should_reject: false, + }, + region_write_buffer_status_from_usage(100, 50, 99, 0) + ); + assert_eq!( + RegionWriteBufferStatus { + should_flush: true, + should_stall: true, + should_reject: false, + }, + region_write_buffer_status_from_usage(100, 50, 100, 99) + ); + assert_eq!( + RegionWriteBufferStatus { + should_flush: true, + should_stall: true, + should_reject: true, + }, + region_write_buffer_status_from_usage(100, 50, 100, 100) + ); + } + + #[test] + fn test_region_write_buffer_status_saturates_reject_limit() { + assert_eq!( + RegionWriteBufferStatus { + should_flush: true, + should_stall: true, + should_reject: true, + }, + region_write_buffer_status_from_usage(usize::MAX, usize::MAX, usize::MAX, usize::MAX,) + ); + } } diff --git a/src/mito2/src/worker/handle_open.rs b/src/mito2/src/worker/handle_open.rs index 3d3520b0cc8..802ab9c1a83 100644 --- a/src/mito2/src/worker/handle_open.rs +++ b/src/mito2/src/worker/handle_open.rs @@ -18,24 +18,68 @@ use std::sync::Arc; use std::time::Instant; use common_telemetry::info; -use object_store::util::join_path; +use object_store::util::{join_path, normalize_dir}; use snafu::{OptionExt, ResultExt}; use store_api::logstore::LogStore; -use store_api::region_request::RegionOpenRequest; +use store_api::region_request::{AffectedRows, RegionCleanUpRequest, RegionOpenRequest}; use store_api::storage::RegionId; use table::requests::STORAGE_KEY; use crate::error::{ - ObjectStoreNotFoundSnafu, OpenDalSnafu, OpenRegionSnafu, RegionNotFoundSnafu, Result, + ObjectStoreNotFoundSnafu, OpenDalSnafu, OpenRegionSnafu, RegionBusySnafu, RegionNotFoundSnafu, + Result, }; -use crate::region::opener::{RegionOpener, sanitize_open_request_options}; +use crate::region::opener::{ + RegionOpener, get_object_store, provider_from_wal_options, sanitize_open_request_options, +}; +use crate::region::options::RegionOptions; use crate::request::OptionOutputTx; use crate::sst::location::region_dir_from_table_dir; use crate::wal::entry_distributor::WalEntryReceiver; -use crate::worker::handle_drop::remove_region_dir_once; +use crate::worker::handle_drop::{ + cleanup_region_file_artifacts, remove_region_dir_for_full_drop, remove_region_dir_once, +}; use crate::worker::{DROPPING_MARKER_FILE, RegionWorkerLoop}; impl RegionWorkerLoop { + pub(crate) async fn handle_offline_cleanup_request( + &mut self, + region_id: RegionId, + mut request: RegionCleanUpRequest, + ) -> Result { + info!( + "Try to clean region {} offline, worker: {}", + region_id, self.id + ); + + if self.regions.is_region_exists(region_id) { + return RegionBusySnafu { region_id }.fail(); + } + + sanitize_open_request_options(&mut request.options); + + let options = RegionOptions::try_from_options(region_id, &request.options)?; + let object_store = get_object_store(&options.storage, &self.object_store_manager)?; + let provider = provider_from_wal_options::(region_id, &options.wal_options)?; + self.wal.obsolete_all(region_id, &provider).await?; + + let table_dir = normalize_dir(&request.table_dir); + let region_dir = region_dir_from_table_dir(&table_dir, region_id, request.path_type); + remove_region_dir_for_full_drop(®ion_dir, &object_store).await?; + + self.cleanup_dropped_region_runtime_state(region_id).await; + self.dropping_regions.remove_region(region_id); + cleanup_region_file_artifacts( + region_id, + &table_dir, + &self.intermediate_manager, + &self.cache_manager, + ) + .await; + + Ok(0) + } + async fn check_and_cleanup_region( &self, region_id: RegionId, @@ -149,6 +193,13 @@ impl RegionWorkerLoop { ); region_count.inc(); + // Notify the region hook that the region has been opened. + // Fires before registration; allocates nothing when no hook + // is registered. + if let Some(hook) = region.manifest_ctx.hook() { + hook.on_region_opened(region_id, ®ion.metadata()).await; + } + // Insert the Region into the RegionMap. regions.insert_region(region); diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index dd0acd63000..0eb4c7bc39e 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -14,7 +14,7 @@ //! Handling write requests. -use std::collections::{HashMap, hash_map}; +use std::collections::{HashMap, HashSet, hash_map}; use std::sync::Arc; use api::v1::OpType; @@ -50,8 +50,11 @@ impl RegionWorkerLoop { return; } - // Flush this worker if the engine needs to flush. + let write_region_ids = write_region_ids(write_requests, bulk_requests); + + // Check region pressure before writes to match the global write buffer behavior. self.maybe_flush_worker(); + let pressure = self.maybe_flush_write_regions(write_region_ids); if self.should_reject_write() { // The memory pressure is still too high, reject write requests. @@ -61,6 +64,20 @@ impl RegionWorkerLoop { return; } + if !pressure.rejected_region_ids.is_empty() { + reject_region_write_requests( + &pressure.rejected_region_ids, + write_requests, + bulk_requests, + ); + for region_id in &pressure.rejected_region_ids { + self.reject_region_stalled_requests(region_id); + } + if write_requests.is_empty() && bulk_requests.is_empty() { + return; + } + } + if self.write_buffer_manager.should_stall() && allow_stall { let stalled_count = (write_requests.len() + bulk_requests.len()) as i64; self.stalling_count.add(stalled_count); @@ -70,6 +87,17 @@ impl RegionWorkerLoop { return; } + if allow_stall { + self.stall_region_write_requests( + &pressure.stalled_region_ids, + write_requests, + bulk_requests, + ); + if write_requests.is_empty() && bulk_requests.is_empty() { + return; + } + } + // Prepare write context. let mut region_ctxs = { let _timer = WRITE_STAGE_ELAPSED @@ -163,15 +191,31 @@ impl RegionWorkerLoop { .inc_by(delete_rows as u64); } - /// Handles all stalled write requests. + /// Handles stalled write requests whose regions no longer need to stall. pub(crate) async fn handle_stalled_requests(&mut self) { - // Handle stalled requests. - let stalled = std::mem::take(&mut self.stalled_requests); - self.stalling_count.sub(stalled.stalled_count() as i64); - // We already stalled these requests, don't stall them again. - for (_, (_, mut requests, mut bulk)) in stalled.requests { - self.handle_write_requests(&mut requests, &mut bulk, false) - .await; + let region_ids = self + .stalled_requests + .requests + .keys() + .copied() + .collect::>(); + let pressure = self.maybe_flush_write_regions(region_ids); + for region_id in &pressure.rejected_region_ids { + self.reject_region_stalled_requests(region_id); + } + let ready_region_ids = self + .stalled_requests + .requests + .keys() + .filter(|region_id| !pressure.stalled_region_ids.contains(region_id)) + .copied() + .collect::>(); + + // These requests have already been stalled. Retry ready regions without stalling the + // same requests again. Regions that still exceed their limit remain in the queue until + // their own flush releases the pressure. + for region_id in ready_region_ids { + self.handle_region_stalled_requests(®ion_id, false).await; } } @@ -553,6 +597,34 @@ impl RegionWorkerLoop { self.write_buffer_manager.memory_usage() + self.stalled_requests.estimated_size >= self.config.global_write_buffer_reject_size.as_bytes() as usize } + + fn stall_region_write_requests( + &mut self, + stalled_region_ids: &HashSet, + write_requests: &mut Vec, + bulk_requests: &mut Vec, + ) { + let mut stalled_count = 0; + let mut stalled_write_requests = write_requests + .extract_if(.., |req| { + stalled_region_ids.contains(&req.request.region_id) + }) + .collect::>(); + let mut stalled_bulk_requests = bulk_requests + .extract_if(.., |req| stalled_region_ids.contains(&req.region_id)) + .collect::>(); + + stalled_count += stalled_write_requests.len() + stalled_bulk_requests.len(); + self.stalled_requests + .append(&mut stalled_write_requests, &mut stalled_bulk_requests); + + if stalled_count > 0 { + let stalled_count = stalled_count as i64; + self.stalling_count.add(stalled_count); + WRITE_STALL_TOTAL.inc_by(stalled_count as u64); + self.listener.on_write_stall(); + } + } } /// Send rejected error to all `write_requests`. @@ -576,6 +648,33 @@ fn reject_write_requests( } } +fn reject_region_write_requests( + rejected_region_ids: &HashSet, + write_requests: &mut Vec, + bulk_requests: &mut Vec, +) { + let mut rejected_write_requests = write_requests + .extract_if(.., |req| { + rejected_region_ids.contains(&req.request.region_id) + }) + .collect::>(); + let mut rejected_bulk_requests = bulk_requests + .extract_if(.., |req| rejected_region_ids.contains(&req.region_id)) + .collect::>(); + reject_write_requests(&mut rejected_write_requests, &mut rejected_bulk_requests); +} + +fn write_region_ids( + write_requests: &[SenderWriteRequest], + bulk_requests: &[SenderBulkRequest], +) -> HashSet { + write_requests + .iter() + .map(|req| req.request.region_id) + .chain(bulk_requests.iter().map(|req| req.region_id)) + .collect() +} + /// Rejects delete request under append mode. fn check_op_type(append_mode: bool, request: &WriteRequest) -> Result<()> { if append_mode { diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index 13b86140aa1..66fb66d7527 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -589,6 +589,22 @@ pub enum Error { location: Location, }, + #[snafu(display( + "CSV header mismatch in path: {}, unknown columns: {:?}, missing columns: {:?}, duplicate columns: {:?}", + path, + unknown_columns, + missing_columns, + duplicate_columns + ))] + CsvHeaderMismatch { + path: String, + unknown_columns: Vec, + missing_columns: Vec, + duplicate_columns: Vec, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to project schema"))] ProjectSchema { #[snafu(source)] @@ -936,6 +952,7 @@ impl ErrorExt for Error { | Error::ColumnNotFound { .. } | Error::BuildRegex { .. } | Error::InvalidSchema { .. } + | Error::CsvHeaderMismatch { .. } | Error::ProjectSchema { .. } | Error::UnsupportedFormat { .. } | Error::ColumnNoneDefaultValue { .. } diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index a433602a36a..d2fbe20d034 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -107,7 +107,7 @@ pub enum AutoCreateTableType { /// A table that merges rows by `last_non_null` strategy. LastNonNull, /// Create table that build index and default partition rules on trace_id - Trace, + Trace { alter_existing: bool }, } impl AutoCreateTableType { @@ -117,9 +117,18 @@ impl AutoCreateTableType { AutoCreateTableType::Physical => "physical", AutoCreateTableType::Log => "log", AutoCreateTableType::LastNonNull => "last_non_null", - AutoCreateTableType::Trace => "trace", + AutoCreateTableType::Trace { .. } => "trace", } } + + fn alter_existing(&self) -> bool { + !matches!( + self, + Self::Trace { + alter_existing: false + } + ) + } } /// Split insert requests into normal and instant requests. @@ -214,7 +223,9 @@ impl Inserter { requests, ctx, statement_executor, - AutoCreateTableType::Trace, + AutoCreateTableType::Trace { + alter_existing: true, + }, false, false, ) @@ -503,6 +514,39 @@ impl Inserter { }) } + /// Ensures a trace table has the request-global schema without requiring a + /// padded data row to drive on-demand creation or alteration. When + /// `alter_existing` is false, a table created after planning is left for the + /// caller to re-plan. + pub async fn ensure_trace_table_on_demand( + &self, + table_name: &str, + request_schema: Vec, + alter_existing: bool, + ctx: &QueryContextRef, + statement_executor: &StatementExecutor, + ) -> Result<()> { + let mut requests = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: table_name.to_string(), + rows: Some(api::v1::Rows { + schema: request_schema, + rows: Vec::new(), + }), + }], + }; + self.create_or_alter_tables_on_demand( + &mut requests, + ctx, + AutoCreateTableType::Trace { alter_existing }, + statement_executor, + false, + false, + ) + .await?; + Ok(()) + } + /// Creates or alter tables on demand: /// - if table does not exist, create table by inferred CreateExpr /// - if table exist, check if schema matches. If any new column found, alter table by inferred `AlterExpr` @@ -575,6 +619,7 @@ impl Inserter { ctx, accommodate_existing_schema, is_single_value, + auto_create_table_type.alter_existing(), )? { alter_tables.push(alter_expr); need_refresh_table_infos.insert(( @@ -639,7 +684,7 @@ impl Inserter { } } - AutoCreateTableType::Trace => { + AutoCreateTableType::Trace { .. } => { let trace_table_name = ctx .extension(TRACE_TABLE_NAME_SESSION_KEY) .unwrap_or(TRACE_TABLE_NAME); @@ -904,7 +949,12 @@ impl Inserter { ctx: &QueryContextRef, accommodate_existing_schema: bool, is_single_value: bool, + alter_existing: bool, ) -> Result> { + if !alter_existing { + return Ok(None); + } + let catalog_name = ctx.current_catalog(); let schema_name = ctx.current_schema(); let table_name = table.table_info().name.clone(); @@ -1187,7 +1237,7 @@ pub fn fill_table_options_for_create( table_options.insert(MERGE_MODE_KEY.to_string(), "last_non_null".to_string()); } } - AutoCreateTableType::Trace => { + AutoCreateTableType::Trace { .. } => { table_options.insert(APPEND_MODE_KEY.to_string(), "true".to_string()); } } @@ -1473,8 +1523,15 @@ mod tests { )), true, ); + // Do not apply an absent-table plan to a table that appeared concurrently. + assert!( + inserter + .get_alter_table_expr_on_demand(&mut req, &table, &ctx, false, false, false) + .unwrap() + .is_none() + ); let alter_expr = inserter - .get_alter_table_expr_on_demand(&mut req, &table, &ctx, true, true) + .get_alter_table_expr_on_demand(&mut req, &table, &ctx, true, true, true) .unwrap(); assert!(alter_expr.is_none()); diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 34b2550acf2..5ec7b4bc3cc 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -78,8 +78,7 @@ use table::table_name::TableName; use table::table_reference::TableReference; use self::set::{ - set_bytea_output, set_datestyle, set_intervalstyle, set_search_path, set_timezone, - validate_client_encoding, + set_bytea_output, set_datestyle, set_intervalstyle, set_timezone, validate_client_encoding, }; use crate::error::{ self, CatalogSnafu, ExecLogicalPlanSnafu, ExternalSnafu, InvalidSqlSnafu, NotSupportedSnafu, @@ -528,7 +527,10 @@ impl StatementExecutor { }, "SEARCH_PATH" => { if query_ctx.channel() == Channel::Postgres { - set_search_path(set_var.value, query_ctx)? + let search_path = set_var.search_path().context(NotSupportedSnafu { + feat: "Unsupported search path in set variable statement", + })?; + query_ctx.set_current_schema(search_path); } else { return NotSupportedSnafu { feat: format!("Unsupported set variable {}", var_name), diff --git a/src/operator/src/statement/copy_table_from.rs b/src/operator/src/statement/copy_table_from.rs index b64c654923c..014a2d32e9b 100644 --- a/src/operator/src/statement/copy_table_from.rs +++ b/src/operator/src/statement/copy_table_from.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; +use std::collections::{BTreeSet, HashMap, HashSet}; use std::future::Future; use std::path::Path; use std::pin::Pin; @@ -392,6 +392,7 @@ impl StatementExecutor { .collect_metadata(&object_store, format.clone(), path.to_string()) .await?; + validate_csv_headers_if_required(&file_metadata, &table_schema)?; let schema_mapping = copy_from_schema_mapping(&file_metadata, &table_schema); let projected_file_schema = Arc::new( file_metadata @@ -642,6 +643,62 @@ fn ensure_schema_compatible(from: &SchemaRef, to: &SchemaRef) -> Result<()> { } } +fn validate_csv_headers_if_required(file_metadata: &FileMetadata, table: &SchemaRef) -> Result<()> { + let FileMetadata::Csv { + schema, + format, + path, + } = file_metadata + else { + return Ok(()); + }; + + if !format.strict_headers { + return Ok(()); + } + + let mut seen_file_columns = HashSet::with_capacity(schema.fields().len()); + let duplicate_columns = schema + .fields() + .iter() + .filter_map(|field| { + if seen_file_columns.insert(field.name().clone()) { + None + } else { + Some(field.name().clone()) + } + }) + .collect::>() + .into_iter() + .collect::>(); + let file_columns = seen_file_columns.into_iter().collect::>(); + let table_columns = table + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + let unknown_columns = file_columns + .difference(&table_columns) + .cloned() + .collect::>(); + let missing_columns = table_columns + .difference(&file_columns) + .cloned() + .collect::>(); + + ensure!( + unknown_columns.is_empty() && missing_columns.is_empty() && duplicate_columns.is_empty(), + error::CsvHeaderMismatchSnafu { + path, + unknown_columns, + missing_columns, + duplicate_columns, + } + ); + + Ok(()) +} + /// Generates a maybe compatible schema of the file schema. /// /// If there is a field is found in table schema, @@ -998,12 +1055,131 @@ mod tests { } } + fn make_strict_csv_metadata(schema: Arc) -> FileMetadata { + FileMetadata::Csv { + schema, + format: CsvFormat { + strict_headers: true, + ..CsvFormat::default() + }, + path: "test.csv".to_string(), + } + } + fn assert_field(schema: &Schema, idx: usize, name: &str, data_type: &DataType) { let field = schema.field(idx); assert_eq!(field.name(), name); assert_eq!(field.data_type(), data_type); } + #[test] + fn test_strict_csv_headers_allows_reordered_columns() { + let file_schema = make_test_schema(&[ + Field::new("ts", DataType::Utf8, true), + Field::new("host_id", DataType::UInt8, true), + Field::new("reading_value", DataType::Float64, true), + ]); + let table_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt32, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + ]); + + validate_csv_headers_if_required(&make_strict_csv_metadata(file_schema), &table_schema) + .unwrap(); + } + + #[test] + fn test_strict_csv_headers_rejects_unknown_columns() { + let file_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt8, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + Field::new("extra", DataType::Utf8, true), + ]); + let table_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt32, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + ]); + + let err = + validate_csv_headers_if_required(&make_strict_csv_metadata(file_schema), &table_schema) + .unwrap_err(); + + assert!(matches!( + err, + error::Error::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + } if unknown_columns == vec!["extra".to_string()] + && missing_columns.is_empty() + && duplicate_columns.is_empty() + )); + } + + #[test] + fn test_strict_csv_headers_rejects_missing_columns() { + let file_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt8, true), + Field::new("ts", DataType::Utf8, true), + ]); + let table_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt32, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + ]); + + let err = + validate_csv_headers_if_required(&make_strict_csv_metadata(file_schema), &table_schema) + .unwrap_err(); + + assert!(matches!( + err, + error::Error::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + } if unknown_columns.is_empty() + && missing_columns == vec!["reading_value".to_string()] + && duplicate_columns.is_empty() + )); + } + + #[test] + fn test_strict_csv_headers_rejects_duplicate_columns() { + let file_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt8, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + Field::new("host_id", DataType::UInt16, true), + ]); + let table_schema = make_test_schema(&[ + Field::new("host_id", DataType::UInt32, true), + Field::new("reading_value", DataType::Float64, true), + Field::new("ts", DataType::Utf8, true), + ]); + + let err = + validate_csv_headers_if_required(&make_strict_csv_metadata(file_schema), &table_schema) + .unwrap_err(); + + assert!(matches!( + err, + error::Error::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + } if unknown_columns.is_empty() + && missing_columns.is_empty() + && duplicate_columns == vec!["host_id".to_string()] + )); + } + #[test] fn test_headerless_csv_schema_projection_is_positional() { let file_schema = make_test_schema(&[ diff --git a/src/operator/src/statement/set.rs b/src/operator/src/statement/set.rs index 38a16e1b051..77515e62e09 100644 --- a/src/operator/src/statement/set.rs +++ b/src/operator/src/statement/set.rs @@ -129,36 +129,6 @@ pub fn set_bytea_output(exprs: Vec, ctx: QueryContextRef) -> Result<()> { Ok(()) } -pub fn set_search_path(exprs: Vec, ctx: QueryContextRef) -> Result<()> { - let search_expr = exprs.first().context(NotSupportedSnafu { - feat: "No search path find in set variable statement", - })?; - match search_expr { - Expr::Value(ValueWithSpan { - value: Value::SingleQuotedString(search_path), - .. - }) - | Expr::Value(ValueWithSpan { - value: Value::DoubleQuotedString(search_path), - .. - }) => { - ctx.set_current_schema(search_path); - Ok(()) - } - Expr::Identifier(Ident { value, .. }) => { - ctx.set_current_schema(value); - Ok(()) - } - expr => NotSupportedSnafu { - feat: format!( - "Unsupported search path expr {} in set variable statement", - expr - ), - } - .fail(), - } -} - pub fn validate_client_encoding(set: SetVariables) -> Result<()> { let Some((encoding, [])) = set.value.split_first() else { return InvalidSqlSnafu { diff --git a/src/promql/Cargo.toml b/src/promql/Cargo.toml index 460be8ddd95..36ac213fef4 100644 --- a/src/promql/Cargo.toml +++ b/src/promql/Cargo.toml @@ -13,6 +13,7 @@ async-trait.workspace = true bytemuck.workspace = true common-error.workspace = true common-macro.workspace = true +common-query.workspace = true common-telemetry.workspace = true datafusion.workspace = true datafusion-common.workspace = true diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index 7318b9cba5d..1897b453f38 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -18,6 +18,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use common_query::prometheus::is_prometheus_stale_nan; use datafusion::arrow::array::{Array, Float64Array, TimestampMillisecondArray, UInt64Array}; use datafusion::arrow::datatypes::{DataType, SchemaRef}; use datafusion::arrow::record_batch::RecordBatch; @@ -579,9 +580,10 @@ impl InstantManipulateStream { match curr.cmp(&expected_ts) { Ordering::Equal => { if let Some(field_column) = &field_column - && field_column.value(cursor).is_nan() + && field_column.is_valid(cursor) + && is_prometheus_stale_nan(field_column.value(cursor)) { - // ignore the NaN value + // Ignore the stale marker. } else { take_indices.push(cursor as u64); aligned_ts.push(expected_ts); @@ -613,9 +615,10 @@ impl InstantManipulateStream { if prev_ts + self.lookback_delta > expected_ts { // only use the point in the time range if let Some(field_column) = &field_column - && field_column.value(prev_cursor).is_nan() + && field_column.is_valid(prev_cursor) + && is_prometheus_stale_nan(field_column.value(prev_cursor)) { - // if the newest value is NaN, it means the value is stale, so we should not use it + // Do not use a stale marker as the newest value. continue; } // use this point @@ -624,9 +627,10 @@ impl InstantManipulateStream { } } } else if let Some(field_column) = &field_column - && field_column.value(cursor).is_nan() + && field_column.is_valid(cursor) + && is_prometheus_stale_nan(field_column.value(cursor)) { - // if the newest value is NaN, it means the value is stale, so we should not use it + // Do not use a stale marker as the newest value. } else { // use this point take_indices.push(cursor as u64); @@ -688,6 +692,7 @@ fn reuse_constant_column(array: &Arc, len: usize) -> DataFusionResult #[cfg(test)] mod test { + use datafusion::arrow::buffer::NullBuffer; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::common::ToDFSchema; use datafusion::datasource::memory::MemorySourceConfig; @@ -697,7 +702,7 @@ mod test { use super::*; use crate::extension_plan::test_util::{ - TIME_INDEX_COLUMN, prepare_test_data, prepare_test_data_with_nan, + TIME_INDEX_COLUMN, prepare_test_data, prepare_test_data_with_stale_marker, }; async fn do_normalize_test( @@ -706,10 +711,10 @@ mod test { lookback_delta: Millisecond, interval: Millisecond, expected: String, - contains_nan: bool, + contains_stale_marker: bool, ) { - let memory_exec = if contains_nan { - Arc::new(prepare_test_data_with_nan()) + let memory_exec = if contains_stale_marker { + Arc::new(prepare_test_data_with_stale_marker()) } else { Arc::new(prepare_test_data()) }; @@ -1255,7 +1260,7 @@ mod test { } #[tokio::test] - async fn lookback_10s_interval_10s_with_nan() { + async fn lookback_10s_interval_10s_with_stale_marker() { let expected = String::from( "+---------------------+-------+\ \n| timestamp | value |\ @@ -1269,7 +1274,7 @@ mod test { } #[tokio::test] - async fn lookback_10s_interval_10s_with_nan_unaligned() { + async fn lookback_10s_interval_10s_with_stale_marker_unaligned() { let expected = String::from( "+-------------------------+-------+\ \n| timestamp | value |\ @@ -1303,4 +1308,190 @@ mod test { ) .await; } + + #[tokio::test] + async fn ordinary_nan_is_selected_for_exact_and_lookback() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000])), + Arc::new(Float64Array::from(vec![f64::NAN])), + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let exec = Arc::new(InstantManipulateExec { + start: 1_000, + end: 1_500, + lookback_delta: 1_000, + interval: 500, + time_index_column: TIME_INDEX_COLUMN.to_string(), + field_column: Some("value".to_string()), + reuse_tsid_column: false, + input, + metric: ExecutionPlanMetricsSet::new(), + }); + + let context = SessionContext::default(); + let batches = datafusion::physical_plan::collect(exec, context.task_ctx()) + .await + .unwrap(); + let values = batches + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + let timestamps = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + + assert_eq!(values.len(), 2); + assert_eq!(timestamps, vec![1_000, 1_500]); + assert!(values.iter().all(|value| value.is_nan())); + assert_eq!( + values + .iter() + .map(|value| value.to_bits()) + .collect::>(), + vec![f64::NAN.to_bits(); 2] + ); + } + + #[tokio::test] + async fn prometheus_stale_nan_suppresses_exact_and_lookback() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![500, 1_000])), + Arc::new(Float64Array::from(vec![ + 42.0, + f64::from_bits(0x7ff0_0000_0000_0002), + ])), + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let exec = Arc::new(InstantManipulateExec { + start: 1_000, + end: 1_500, + lookback_delta: 1_001, + interval: 500, + time_index_column: TIME_INDEX_COLUMN.to_string(), + field_column: Some("value".to_string()), + reuse_tsid_column: false, + input, + metric: ExecutionPlanMetricsSet::new(), + }); + + let context = SessionContext::default(); + let batches = datafusion::physical_plan::collect(exec, context.task_ctx()) + .await + .unwrap(); + + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 0, + "the stale marker must suppress both the exact and lookback selections rather than \ + falling back to 42.0" + ); + } + + #[tokio::test] + async fn null_value_backed_by_stale_bits_is_selected_for_exact_and_lookback() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let field_column = Float64Array::new( + vec![f64::from_bits(0x7ff0_0000_0000_0002)].into(), + Some(NullBuffer::from(vec![false])), + ); + assert!(!field_column.is_valid(0)); + assert_eq!(field_column.value(0).to_bits(), 0x7ff0_0000_0000_0002); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000])), + Arc::new(field_column), + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let exec = Arc::new(InstantManipulateExec { + start: 1_000, + end: 1_500, + lookback_delta: 1_000, + interval: 500, + time_index_column: TIME_INDEX_COLUMN.to_string(), + field_column: Some("value".to_string()), + reuse_tsid_column: false, + input, + metric: ExecutionPlanMetricsSet::new(), + }); + + let context = SessionContext::default(); + let batches = datafusion::physical_plan::collect(exec, context.task_ctx()) + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + let batch = batches.iter().find(|batch| batch.num_rows() == 2).unwrap(); + let timestamps = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(timestamps.values(), &[1_000, 1_500]); + assert!(!values.is_valid(0)); + assert!(!values.is_valid(1)); + } } diff --git a/src/promql/src/extension_plan/normalize.rs b/src/promql/src/extension_plan/normalize.rs index 6e24f3869b0..a84d5fc5c01 100644 --- a/src/promql/src/extension_plan/normalize.rs +++ b/src/promql/src/extension_plan/normalize.rs @@ -17,7 +17,8 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use datafusion::arrow::array::{BooleanArray, Float64Array}; +use common_query::prometheus::is_prometheus_stale_nan; +use datafusion::arrow::array::{Array, BooleanArray, Float64Array}; use datafusion::arrow::compute; use datafusion::common::{DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics}; use datafusion::error::DataFusionError; @@ -52,12 +53,12 @@ use crate::metrics::PROMQL_SERIES_COUNT; /// Roughly speaking, this method does these things: /// - bias sample's timestamp by offset /// - sort the record batch based on timestamp column -/// - remove NaN values (optional) +/// - remove Prometheus stale markers (optional) #[derive(Debug, PartialEq, Eq, Hash, PartialOrd)] pub struct SeriesNormalize { offset: Millisecond, time_index_column_name: String, - need_filter_out_nan: bool, + filter_stale_markers: bool, tag_columns: Vec, input: LogicalPlan, @@ -122,7 +123,7 @@ impl UserDefinedLogicalNodeCore for SeriesNormalize { write!( f, "PromSeriesNormalize: offset=[{}], time index=[{}], filter NaN: [{}]", - self.offset, self.time_index_column_name, self.need_filter_out_nan + self.offset, self.time_index_column_name, self.filter_stale_markers ) } @@ -158,7 +159,7 @@ impl UserDefinedLogicalNodeCore for SeriesNormalize { Ok(Self { offset: self.offset, time_index_column_name, - need_filter_out_nan: self.need_filter_out_nan, + filter_stale_markers: self.filter_stale_markers, tag_columns, input, unfix: None, @@ -167,7 +168,7 @@ impl UserDefinedLogicalNodeCore for SeriesNormalize { Ok(Self { offset: self.offset, time_index_column_name: self.time_index_column_name.clone(), - need_filter_out_nan: self.need_filter_out_nan, + filter_stale_markers: self.filter_stale_markers, tag_columns: self.tag_columns.clone(), input, unfix: None, @@ -180,14 +181,14 @@ impl SeriesNormalize { pub fn new>( offset: Millisecond, time_index_column_name: N, - need_filter_out_nan: bool, + filter_stale_markers: bool, tag_columns: Vec, input: LogicalPlan, ) -> Self { Self { offset, time_index_column_name: time_index_column_name.as_ref().to_string(), - need_filter_out_nan, + filter_stale_markers, tag_columns, input, unfix: None, @@ -202,7 +203,7 @@ impl SeriesNormalize { Arc::new(SeriesNormalizeExec { offset: self.offset, time_index_column_name: self.time_index_column_name.clone(), - need_filter_out_nan: self.need_filter_out_nan, + filter_stale_markers: self.filter_stale_markers, input: exec_input, tag_columns: self.tag_columns.clone(), metric: ExecutionPlanMetricsSet::new(), @@ -222,7 +223,7 @@ impl SeriesNormalize { pb::SeriesNormalize { offset: self.offset, time_index_idx, - filter_nan: self.need_filter_out_nan, + filter_nan: self.filter_stale_markers, tag_column_indices, ..Default::default() } @@ -244,7 +245,7 @@ impl SeriesNormalize { Ok(Self { offset: pb_normalize.offset, time_index_column_name: String::new(), - need_filter_out_nan: pb_normalize.filter_nan, + filter_stale_markers: pb_normalize.filter_nan, tag_columns: Vec::new(), input: placeholder_plan, unfix: Some(unfix), @@ -256,7 +257,7 @@ impl SeriesNormalize { pub struct SeriesNormalizeExec { offset: Millisecond, time_index_column_name: String, - need_filter_out_nan: bool, + filter_stale_markers: bool, tag_columns: Vec, input: Arc, @@ -303,7 +304,7 @@ impl ExecutionPlan for SeriesNormalizeExec { Ok(Arc::new(Self { offset: self.offset, time_index_column_name: self.time_index_column_name.clone(), - need_filter_out_nan: self.need_filter_out_nan, + filter_stale_markers: self.filter_stale_markers, input: children[0].clone(), tag_columns: self.tag_columns.clone(), metric: self.metric.clone(), @@ -334,7 +335,7 @@ impl ExecutionPlan for SeriesNormalizeExec { Ok(Box::pin(SeriesNormalizeStream { offset: self.offset, time_index, - need_filter_out_nan: self.need_filter_out_nan, + filter_stale_markers: self.filter_stale_markers, schema, input, metric: baseline_metric, @@ -364,7 +365,7 @@ impl DisplayAs for SeriesNormalizeExec { write!( f, "PromSeriesNormalizeExec: offset=[{}], time index=[{}], filter NaN: [{}]", - self.offset, self.time_index_column_name, self.need_filter_out_nan + self.offset, self.time_index_column_name, self.filter_stale_markers ) } } @@ -375,7 +376,7 @@ pub struct SeriesNormalizeStream { offset: Millisecond, // Column index of TIME INDEX column's position in schema time_index: usize, - need_filter_out_nan: bool, + filter_stale_markers: bool, schema: SchemaRef, input: SendableRecordBatchStream, @@ -408,25 +409,25 @@ impl SeriesNormalizeStream { columns[self.time_index] = ts_column_biased; let result_batch = RecordBatch::try_new(input.schema(), columns)?; - if !self.need_filter_out_nan { + if !self.filter_stale_markers { return Ok(result_batch); } - // TODO(ruihang): consider the "special NaN" - // filter out NaN - let mut filter = vec![true; input.num_rows()]; + // Filter out Prometheus stale markers. + let mut stale_marker_filter = vec![true; input.num_rows()]; for column in result_batch.columns() { if let Some(float_column) = column.as_any().downcast_ref::() { - for (i, flag) in filter.iter_mut().enumerate() { - if float_column.value(i).is_nan() { + for (i, flag) in stale_marker_filter.iter_mut().enumerate() { + if float_column.is_valid(i) && is_prometheus_stale_nan(float_column.value(i)) { *flag = false; } } } } - let result = compute::filter_record_batch(&result_batch, &BooleanArray::from(filter)) - .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; + let result = + compute::filter_record_batch(&result_batch, &BooleanArray::from(stale_marker_filter)) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; Ok(result) } } @@ -462,6 +463,7 @@ impl Stream for SeriesNormalizeStream { #[cfg(test)] mod test { use datafusion::arrow::array::Float64Array; + use datafusion::arrow::buffer::NullBuffer; use datafusion::arrow::datatypes::{ ArrowPrimitiveType, DataType, Field, Schema, TimestampMillisecondType, }; @@ -522,7 +524,7 @@ mod test { let normalize_exec = Arc::new(SeriesNormalizeExec { offset: 0, time_index_column_name: TIME_INDEX_COLUMN.to_string(), - need_filter_out_nan: true, + filter_stale_markers: true, input: memory_exec, tag_columns: vec!["path".to_string()], metric: ExecutionPlanMetricsSet::new(), @@ -556,7 +558,7 @@ mod test { let normalize_exec = Arc::new(SeriesNormalizeExec { offset: 1_000, time_index_column_name: TIME_INDEX_COLUMN.to_string(), - need_filter_out_nan: true, + filter_stale_markers: true, input: memory_exec, metric: ExecutionPlanMetricsSet::new(), tag_columns: vec!["path".to_string()], @@ -583,4 +585,77 @@ mod test { assert_eq!(result_literal, expected); } + + #[tokio::test] + async fn filters_stale_markers_and_preserves_ordinary_nan() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + TimestampMillisecondType::DATA_TYPE, + false, + ), + Field::new("value", DataType::Float64, true), + Field::new("auxiliary", DataType::Float64, true), + ])); + let value_column = Float64Array::new( + vec![ + 42.0, + f64::from_bits(0x7ff0_0000_0000_0002), + 24.0, + f64::from_bits(0x7ff0_0000_0000_0002), + ] + .into(), + Some(NullBuffer::from(vec![true, true, true, false])), + ); + assert!(!value_column.is_valid(3)); + assert_eq!(value_column.value(3).to_bits(), 0x7ff0_0000_0000_0002); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![ + 1_000, 2_000, 3_000, 4_000, + ])), + Arc::new(value_column), + Arc::new(Float64Array::from(vec![ + f64::from_bits(0x7ff8_0000_0000_0000), + 1.0, + f64::from_bits(0x7ff0_0000_0000_0002), + 2.0, + ])), + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let exec = Arc::new(SeriesNormalizeExec { + offset: 0, + time_index_column_name: TIME_INDEX_COLUMN.to_string(), + filter_stale_markers: true, + tag_columns: Vec::new(), + input, + metric: ExecutionPlanMetricsSet::new(), + }); + + let context = SessionContext::default(); + let batches = datafusion::physical_plan::collect(exec, context.task_ctx()) + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + let batch = batches.iter().find(|batch| batch.num_rows() == 2).unwrap(); + let value = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let auxiliary = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(value.value(0), 42.0); + assert_eq!(auxiliary.value(0).to_bits(), 0x7ff8_0000_0000_0000); + assert!(!value.is_valid(1)); + } } diff --git a/src/promql/src/extension_plan/test_util.rs b/src/promql/src/extension_plan/test_util.rs index 5521d151c21..43d0d686038 100644 --- a/src/promql/src/extension_plan/test_util.rs +++ b/src/promql/src/extension_plan/test_util.rs @@ -16,6 +16,7 @@ use std::sync::Arc; +use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use common_recordbatch::DfRecordBatch as RecordBatch; use datafusion::arrow::array::Float64Array; use datafusion::arrow::datatypes::{ @@ -52,7 +53,7 @@ pub(crate) fn prepare_test_data() -> DataSourceExec { )) } -pub(crate) fn prepare_test_data_with_nan() -> DataSourceExec { +pub(crate) fn prepare_test_data_with_stale_marker() -> DataSourceExec { let schema = Arc::new(Schema::new(vec![ Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true), Field::new("value", DataType::Float64, true), @@ -60,7 +61,13 @@ pub(crate) fn prepare_test_data_with_nan() -> DataSourceExec { let timestamp_column = Arc::new(TimestampMillisecondArray::from(vec![ 0, 30_000, 60_000, 90_000, 120_000, // every 30s ])) as _; - let field_column = Arc::new(Float64Array::from(vec![0.0, f64::NAN, 6.0, f64::NAN, 12.0])) as _; + let field_column = Arc::new(Float64Array::from(vec![ + 0.0, + f64::from_bits(PROMETHEUS_STALE_NAN_BITS), + 6.0, + f64::from_bits(PROMETHEUS_STALE_NAN_BITS), + 12.0, + ])) as _; let data = RecordBatch::try_new(schema.clone(), vec![timestamp_column, field_column]).unwrap(); DataSourceExec::new(Arc::new( diff --git a/src/promql/src/extension_plan/union_distinct_on.rs b/src/promql/src/extension_plan/union_distinct_on.rs index 149475bfe49..a7d8fd48594 100644 --- a/src/promql/src/extension_plan/union_distinct_on.rs +++ b/src/promql/src/extension_plan/union_distinct_on.rs @@ -17,7 +17,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use ahash::{HashMap, RandomState}; +use ahash::{HashSet, RandomState}; use datafusion::arrow::array::UInt64Array; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; @@ -34,14 +34,12 @@ use datafusion::physical_plan::{ }; use datafusion_expr::col; use datatypes::arrow::compute; -use futures::future::BoxFuture; -use futures::{Stream, StreamExt, TryStreamExt, ready}; +use futures::{Stream, StreamExt, ready}; use greptime_proto::substrait_extension as pb; use prost::Message; use snafu::ResultExt; -use crate::error::{DeserializeSnafu, Result}; -use crate::extension_plan::{resolve_column_name, serialize_column_index}; +use crate::error::{DataFusionPlanningSnafu, DeserializeSnafu, Result}; /// A special kind of `UNION`(`OR` in PromQL) operator, for PromQL specific use case. /// @@ -49,35 +47,23 @@ use crate::extension_plan::{resolve_column_name, serialize_column_index}; /// most different part is that it treat left child and right child differently: /// - All columns from left child will be outputted. /// - Only check collisions (when not distinct) on the columns specified by `compare_keys`. -/// - When there is a collision: -/// - If the collision is from right child itself, only the first observed row will be -/// preserved. All others are discarded. -/// - If the collision is from left child, the row in right child will be discarded. -/// - The output order is not maintained. This plan will output left child first, then right child. -/// - The output schema contains all columns from left or right child plans. +/// - Rows from the right child with the same comparison signature are all preserved. +/// - If a signature occurs in the left child, all matching rows from the right child are discarded. +/// - Output preserves each input's batch and row order, with all left rows before right rows. +/// - The output schema is based on the left child schema, with nullability widened from both inputs. /// -/// From the implementation perspective, this operator is similar to `HashJoin`, but the -/// probe side is the right child, and the build side is the left child. Another difference -/// is that the probe is opting-out. -/// -/// This plan will exhaust the right child first to build probe hash table, then streaming -/// on left side, and use the left side to "mask" the hash table. +/// Execution streams the left child first while retaining its comparison signatures. The +/// right child is polled only after the left child completes, then rows whose signatures were +/// observed on the left are omitted. #[derive(Debug, PartialEq, Eq, Hash)] pub struct UnionDistinctOn { left: LogicalPlan, right: LogicalPlan, /// The columns to compare for equality. /// TIME INDEX is included. - compare_keys: Vec, - ts_col: String, + compare_key_indices: Vec, + ts_col_idx: usize, output_schema: DFSchemaRef, - unfix: Option, -} - -#[derive(Debug, PartialEq, Eq, Hash, PartialOrd)] -struct UnfixIndices { - pub compare_key_indices: Vec, - pub ts_col_idx: u64, } impl UnionDistinctOn { @@ -85,21 +71,92 @@ impl UnionDistinctOn { "UnionDistinctOn" } - pub fn new( + pub fn try_new( left: LogicalPlan, right: LogicalPlan, - compare_keys: Vec, - ts_col: String, - output_schema: DFSchemaRef, - ) -> Self { - Self { + compare_key_indices: Vec, + ts_col_idx: usize, + ) -> DataFusionResult { + let output_schema = + Self::validate_children(&left, &right, &compare_key_indices, ts_col_idx)?; + Ok(Self { left, right, - compare_keys, - ts_col, + compare_key_indices, + ts_col_idx, output_schema, - unfix: None, + }) + } + + fn validate_children( + left: &LogicalPlan, + right: &LogicalPlan, + compare_key_indices: &[usize], + ts_col_idx: usize, + ) -> DataFusionResult { + let left_schema = left.schema(); + let right_schema = right.schema(); + let left_fields = left_schema.fields(); + let right_fields = right_schema.fields(); + + if left_fields.len() != right_fields.len() { + return Err(DataFusionError::Plan(format!( + "UnionDistinctOn inputs have different field counts: left={}, right={}", + left_fields.len(), + right_fields.len() + ))); } + + for (column_type, index) in compare_key_indices + .iter() + .map(|index| ("compare key", *index)) + .chain(std::iter::once(("timestamp", ts_col_idx))) + { + if index >= left_fields.len() || index >= right_fields.len() { + return Err(DataFusionError::Plan(format!( + "UnionDistinctOn {column_type} index {index} is out of bounds for inputs with {} fields", + left_fields.len() + ))); + } + } + + for (index, (left_field, right_field)) in left_fields.iter().zip(right_fields).enumerate() { + if left_field.data_type() != right_field.data_type() { + return Err(DataFusionError::Plan(format!( + "UnionDistinctOn input field at index {index} has incompatible data types: left={:?}, right={:?}", + left_field.data_type(), + right_field.data_type() + ))); + } + } + + let output_fields = left_fields + .iter() + .zip(right_fields) + .enumerate() + .map(|(index, (left_field, right_field))| { + let (qualifier, _) = left_schema.qualified_field(index); + ( + qualifier.cloned(), + Arc::new( + left_field + .as_ref() + .clone() + .with_nullable(left_field.is_nullable() || right_field.is_nullable()), + ), + ) + }) + .collect(); + let output_schema = + DFSchema::new_with_metadata(output_fields, left_schema.metadata().clone()).map_err( + |error| { + DataFusionError::Plan(format!( + "Failed to construct UnionDistinctOn output schema: {error}" + )) + }, + )?; + + Ok(Arc::new(output_schema)) } pub fn to_execution_plan( @@ -117,8 +174,8 @@ impl UnionDistinctOn { Arc::new(UnionDistinctOnExec { left: left_exec, right: right_exec, - compare_keys: self.compare_keys.clone(), - ts_col: self.ts_col.clone(), + compare_key_indices: self.compare_key_indices.clone(), + ts_col_idx: self.ts_col_idx, output_schema, metric: ExecutionPlanMetricsSet::new(), properties, @@ -128,12 +185,11 @@ impl UnionDistinctOn { pub fn serialize(&self) -> Vec { let compare_key_indices = self - .compare_keys + .compare_key_indices .iter() - .map(|name| serialize_column_index(&self.output_schema, name)) - .collect::>(); - - let ts_col_idx = serialize_column_index(&self.output_schema, &self.ts_col); + .map(|index| u64::try_from(*index).expect("usize always fits in u64")) + .collect(); + let ts_col_idx = u64::try_from(self.ts_col_idx).expect("usize always fits in u64"); pb::UnionDistinctOn { compare_key_indices, @@ -149,18 +205,33 @@ impl UnionDistinctOn { schema: Arc::new(DFSchema::empty()), }); - let unfix = UnfixIndices { - compare_key_indices: pb_union.compare_key_indices.clone(), - ts_col_idx: pb_union.ts_col_idx, - }; + let compare_key_indices = pb_union + .compare_key_indices + .into_iter() + .map(|index| { + usize::try_from(index).map_err(|_| { + DataFusionError::Plan(format!( + "UnionDistinctOn compare key index {index} does not fit in usize" + )) + }) + }) + .collect::>>() + .context(DataFusionPlanningSnafu)?; + let ts_col_idx = usize::try_from(pb_union.ts_col_idx) + .map_err(|_| { + DataFusionError::Plan(format!( + "UnionDistinctOn timestamp index {} does not fit in usize", + pb_union.ts_col_idx + )) + }) + .context(DataFusionPlanningSnafu)?; Ok(Self { left: placeholder_plan.clone(), right: placeholder_plan, - compare_keys: Vec::new(), - ts_col: String::new(), + compare_key_indices, + ts_col_idx, output_schema: Arc::new(DFSchema::empty()), - unfix: Some(unfix), }) } } @@ -176,11 +247,14 @@ impl PartialOrd for UnionDistinctOn { Some(core::cmp::Ordering::Equal) => {} ord => return ord, } - match self.compare_keys.partial_cmp(&other.compare_keys) { + match self + .compare_key_indices + .partial_cmp(&other.compare_key_indices) + { Some(core::cmp::Ordering::Equal) => {} ord => return ord, } - self.ts_col.partial_cmp(&other.ts_col) + self.ts_col_idx.partial_cmp(&other.ts_col_idx) } } @@ -198,22 +272,21 @@ impl UserDefinedLogicalNodeCore for UnionDistinctOn { } fn expressions(&self) -> Vec { - if self.unfix.is_some() { - return vec![]; - } - - let mut exprs: Vec = self.compare_keys.iter().map(col).collect(); - if !self.compare_keys.iter().any(|key| key == &self.ts_col) { - exprs.push(col(&self.ts_col)); + let fields = self.left.schema().fields(); + let mut exprs = self + .compare_key_indices + .iter() + .filter_map(|index| fields.get(*index).map(|field| col(field.name()))) + .collect::>(); + if !self.compare_key_indices.contains(&self.ts_col_idx) + && let Some(field) = fields.get(self.ts_col_idx) + { + exprs.push(col(field.name())); } exprs } fn necessary_children_exprs(&self, _output_columns: &[usize]) -> Option>> { - if self.unfix.is_some() { - return None; - } - let left_len = self.left.schema().fields().len(); let right_len = self.right.schema().fields().len(); Some(vec![ @@ -223,10 +296,20 @@ impl UserDefinedLogicalNodeCore for UnionDistinctOn { } fn fmt_for_explain(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let fields = self.left.schema().fields(); + let display_column = |index: usize| match fields.get(index) { + Some(field) => format!("{}@{index}", field.name()), + None => format!("@{index}"), + }; + let compare_keys = self + .compare_key_indices + .iter() + .map(|index| display_column(*index)) + .collect::>(); write!( f, - "UnionDistinctOn: on col=[{:?}], ts_col=[{}]", - self.compare_keys, self.ts_col + "UnionDistinctOn: on col={compare_keys:?}, ts_col={}", + display_column(self.ts_col_idx) ) } @@ -245,38 +328,12 @@ impl UserDefinedLogicalNodeCore for UnionDistinctOn { let left = inputs.next().unwrap(); let right = inputs.next().unwrap(); - if let Some(unfix) = &self.unfix { - let output_schema = left.schema().clone(); - - let compare_keys = unfix - .compare_key_indices - .iter() - .map(|idx| { - resolve_column_name(*idx, &output_schema, "UnionDistinctOn", "compare key") - }) - .collect::>>()?; - - let ts_col = - resolve_column_name(unfix.ts_col_idx, &output_schema, "UnionDistinctOn", "ts")?; - - Ok(Self { - left, - right, - compare_keys, - ts_col, - output_schema, - unfix: None, - }) - } else { - Ok(Self { - left, - right, - compare_keys: self.compare_keys.clone(), - ts_col: self.ts_col.clone(), - output_schema: self.output_schema.clone(), - unfix: None, - }) - } + Self::try_new( + left, + right, + self.compare_key_indices.clone(), + self.ts_col_idx, + ) } } @@ -284,8 +341,8 @@ impl UserDefinedLogicalNodeCore for UnionDistinctOn { pub struct UnionDistinctOnExec { left: Arc, right: Arc, - compare_keys: Vec, - ts_col: String, + compare_key_indices: Vec, + ts_col_idx: usize, output_schema: SchemaRef, metric: ExecutionPlanMetricsSet, properties: Arc, @@ -326,8 +383,8 @@ impl ExecutionPlan for UnionDistinctOnExec { Ok(Arc::new(UnionDistinctOnExec { left, right, - compare_keys: self.compare_keys.clone(), - ts_col: self.ts_col.clone(), + compare_key_indices: self.compare_key_indices.clone(), + ts_col_idx: self.ts_col_idx, output_schema: self.output_schema.clone(), metric: self.metric.clone(), properties: self.properties.clone(), @@ -341,41 +398,23 @@ impl ExecutionPlan for UnionDistinctOnExec { context: Arc, ) -> DataFusionResult { let left_stream = self.left.execute(partition, context.clone())?; - let right_stream = self.right.execute(partition, context.clone())?; - // Convert column name to column index. Add one for the time column. - let mut key_indices = Vec::with_capacity(self.compare_keys.len() + 1); - for key in &self.compare_keys { - let index = self - .output_schema - .column_with_name(key) - .map(|(i, _)| i) - .ok_or_else(|| DataFusionError::Internal(format!("Column {} not found", key)))?; - key_indices.push(index); - } - let ts_index = self - .output_schema - .column_with_name(&self.ts_col) - .map(|(i, _)| i) - .ok_or_else(|| { - DataFusionError::Internal(format!("Column {} not found", self.ts_col)) - })?; - key_indices.push(ts_index); + let mut key_indices = self.compare_key_indices.clone(); + key_indices.push(self.ts_col_idx); - // Build right hash table future. - let hashed_data_future = HashedDataFut::Pending(Box::pin(HashedData::new( - right_stream, - self.random_state.clone(), - key_indices.clone(), - ))); - - let baseline_metric = BaselineMetrics::new(&self.metric, partition); Ok(Box::pin(UnionDistinctOnStream { left: left_stream, - right: hashed_data_future, + right_plan: self.right.clone(), + right_partition: partition, + right_context: context, + right: None, compare_keys: key_indices, output_schema: self.output_schema.clone(), - metric: baseline_metric, + random_state: self.random_state.clone(), + lhs_signatures: HashSet::default(), + hashes: Vec::new(), + phase: StreamPhase::Left, + metric: BaselineMetrics::new(&self.metric, partition), })) } @@ -396,60 +435,139 @@ impl DisplayAs for UnionDistinctOnExec { | DisplayFormatType::TreeRender => { write!( f, - "UnionDistinctOnExec: on col=[{:?}], ts_col=[{}]", - self.compare_keys, self.ts_col + "UnionDistinctOnExec: on col={:?}, ts_col={}", + self.compare_key_indices, self.ts_col_idx ) } } } } -// TODO(ruihang): some unused fields are for metrics, which will be implemented later. -#[allow(dead_code)] pub struct UnionDistinctOnStream { left: SendableRecordBatchStream, - right: HashedDataFut, + right_plan: Arc, + right_partition: usize, + right_context: Arc, + right: Option, /// Include time index compare_keys: Vec, output_schema: SchemaRef, + random_state: RandomState, + lhs_signatures: HashSet, + hashes: Vec, + phase: StreamPhase, metric: BaselineMetrics, } -impl UnionDistinctOnStream { - fn poll_impl(&mut self, cx: &mut Context<'_>) -> Poll::Item>> { - // resolve the right stream - let right = match self.right { - HashedDataFut::Pending(ref mut fut) => { - let right = ready!(fut.as_mut().poll(cx))?; - self.right = HashedDataFut::Ready(right); - let HashedDataFut::Ready(right_ref) = &mut self.right else { - unreachable!() - }; - right_ref - } - HashedDataFut::Ready(ref mut right) => right, - HashedDataFut::Empty => return Poll::Ready(None), - }; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum StreamPhase { + Left, + Right, + Done, +} - // poll left and probe with right - let next_left = ready!(self.left.poll_next_unpin(cx)); - match next_left { - Some(Ok(left)) => { - // observe left batch and return it - right.update_map(&left)?; - Poll::Ready(Some(Ok(left))) +impl UnionDistinctOnStream { + fn hash_batch(&mut self, batch: &RecordBatch) -> DataFusionResult<()> { + let arrays = self + .compare_keys + .iter() + .map(|index| batch.column(*index).clone()) + .collect::>(); + self.hashes.clear(); + self.hashes.resize(batch.num_rows(), 0); + hash_utils::create_hashes(&arrays, &self.random_state, &mut self.hashes)?; + Ok(()) + } + + fn filter_rhs_batch(&mut self, batch: RecordBatch) -> DataFusionResult> { + self.hash_batch(&batch)?; + let mut survivor_indices: Option> = None; + for (index, hash) in self.hashes.iter().enumerate() { + if self.lhs_signatures.contains(hash) { + survivor_indices.get_or_insert_with(|| (0..index).collect()); + } else if let Some(indices) = &mut survivor_indices { + indices.push(index); } - Some(Err(e)) => Poll::Ready(Some(Err(e))), + } + + match survivor_indices { + None => Ok(Some(with_schema(batch, self.output_schema.clone())?)), + Some(indices) if indices.is_empty() => Ok(None), + Some(indices) => Ok(Some(with_schema( + take_batch(&batch, &indices)?, + self.output_schema.clone(), + )?)), + } + } + + fn terminal_error(&mut self, error: DataFusionError) -> Poll::Item>> { + self.phase = StreamPhase::Done; + Poll::Ready(Some(Err(error))) + } + + fn poll_right(&mut self, cx: &mut Context<'_>) -> Poll::Item>> { + if self.right.is_none() { + let right = match self + .right_plan + .execute(self.right_partition, self.right_context.clone()) + { + Ok(right) => right, + Err(error) => return self.terminal_error(error), + }; + self.right = Some(right); + } + + let right = self.right.as_mut().expect("right stream is initialized"); + match ready!(right.poll_next_unpin(cx)) { + Some(Ok(batch)) if self.lhs_signatures.is_empty() => { + match with_schema(batch, self.output_schema.clone()) { + Ok(batch) => Poll::Ready(Some(Ok(batch))), + Err(error) => self.terminal_error(error), + } + } + Some(Ok(batch)) => match self.filter_rhs_batch(batch) { + Ok(Some(batch)) => Poll::Ready(Some(Ok(batch))), + Ok(None) => { + // One fully filtered input batch has been consumed. Yield so a sequence + // of filtered batches cannot monopolize a single downstream poll. + cx.waker().wake_by_ref(); + Poll::Pending + } + Err(error) => self.terminal_error(error), + }, + Some(Err(error)) => self.terminal_error(error), None => { - // left stream is exhausted, so we can send the right part - let right = std::mem::replace(&mut self.right, HashedDataFut::Empty); - let HashedDataFut::Ready(data) = right else { - unreachable!() - }; - Poll::Ready(Some(data.finish())) + self.phase = StreamPhase::Done; + Poll::Ready(None) } } } + + fn poll_impl(&mut self, cx: &mut Context<'_>) -> Poll::Item>> { + match self.phase { + StreamPhase::Left => match ready!(self.left.poll_next_unpin(cx)) { + Some(Ok(batch)) => { + if batch.num_rows() > 0 { + if let Err(error) = self.hash_batch(&batch) { + return self.terminal_error(error); + } + self.lhs_signatures.extend(self.hashes.iter().copied()); + } + match with_schema(batch, self.output_schema.clone()) { + Ok(batch) => Poll::Ready(Some(Ok(batch))), + Err(error) => self.terminal_error(error), + } + } + Some(Err(error)) => self.terminal_error(error), + None => { + self.phase = StreamPhase::Right; + self.poll_right(cx) + } + }, + StreamPhase::Right => self.poll_right(cx), + StreamPhase::Done => Poll::Ready(None), + } + } } impl RecordBatchStream for UnionDistinctOnStream { @@ -462,185 +580,49 @@ impl Stream for UnionDistinctOnStream { type Item = DataFusionResult; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - self.poll_impl(cx) + let poll = self.poll_impl(cx); + self.metric.record_poll(poll) } } -/// Simple future state for [HashedData] -enum HashedDataFut { - /// The result is not ready - Pending(BoxFuture<'static, DataFusionResult>), - /// The result is ready - Ready(HashedData), - /// The result is taken - Empty, -} - -/// ALL input batches and its hash table -struct HashedData { - // TODO(ruihang): use `JoinHashMap` instead after upgrading to DF 34.0 - /// Hash table for all input batches. The key is hash value, and the value - /// is the index of `bathc`. - hash_map: HashMap, - /// Output batch. - batch: RecordBatch, - /// The indices of the columns to be hashed. - hash_key_indices: Vec, - random_state: RandomState, -} - -impl HashedData { - pub async fn new( - input: SendableRecordBatchStream, - random_state: RandomState, - hash_key_indices: Vec, - ) -> DataFusionResult { - // Collect all batches from the input stream - let initial = (Vec::new(), 0); - let schema = input.schema(); - let (batches, _num_rows) = input - .try_fold(initial, |mut acc, batch| async { - // Update rowcount - acc.1 += batch.num_rows(); - // Push batch to output - acc.0.push(batch); - Ok(acc) - }) - .await?; - - // Create hash for each batch - let mut hash_map = HashMap::default(); - let mut hashes_buffer = Vec::new(); - let mut interleave_indices = Vec::new(); - for (batch_number, batch) in batches.iter().enumerate() { - hashes_buffer.resize(batch.num_rows(), 0); - // get columns for hashing - let arrays = hash_key_indices - .iter() - .map(|i| batch.column(*i).clone()) - .collect::>(); - - // compute hash - let hash_values = - hash_utils::create_hashes(&arrays, &random_state, &mut hashes_buffer)?; - for (row_number, hash_value) in hash_values.iter().enumerate() { - // Only keeps the first observed row for each hash value - if hash_map - .try_insert(*hash_value, interleave_indices.len()) - .is_ok() - { - interleave_indices.push((batch_number, row_number)); - } - } - } - - // Finalize the hash map - let batch = interleave_batches(schema, batches, interleave_indices)?; - - Ok(Self { - hash_map, - batch, - hash_key_indices, - random_state, - }) - } - - /// Remove rows that hash value present in the input - /// record batch from the hash map. - pub fn update_map(&mut self, input: &RecordBatch) -> DataFusionResult<()> { - // get columns for hashing - let mut hashes_buffer = Vec::new(); - let arrays = self - .hash_key_indices - .iter() - .map(|i| input.column(*i).clone()) - .collect::>(); - - // compute hash - hashes_buffer.resize(input.num_rows(), 0); - let hash_values = - hash_utils::create_hashes(&arrays, &self.random_state, &mut hashes_buffer)?; - - // remove those hashes - for hash in hash_values { - self.hash_map.remove(hash); - } - - Ok(()) - } - - pub fn finish(self) -> DataFusionResult { - let valid_indices = self.hash_map.values().copied().collect::>(); - let result = take_batch(&self.batch, &valid_indices)?; - Ok(result) - } -} - -/// Utility function to interleave batches. Based on [interleave](datafusion::arrow::compute::interleave) -fn interleave_batches( - schema: SchemaRef, - batches: Vec, - indices: Vec<(usize, usize)>, -) -> DataFusionResult { - if batches.is_empty() { - if indices.is_empty() { - return Ok(RecordBatch::new_empty(schema)); - } else { - return Err(DataFusionError::Internal( - "Cannot interleave empty batches with non-empty indices".to_string(), - )); - } - } - - // transform batches into arrays - let mut arrays = vec![vec![]; schema.fields().len()]; - for batch in &batches { - for (i, array) in batch.columns().iter().enumerate() { - arrays[i].push(array.as_ref()); - } - } - - // interleave arrays - let interleaved_arrays: Vec<_> = arrays - .into_iter() - .map(|array| compute::interleave(&array, &indices)) - .collect::>()?; - - // assemble new record batch - RecordBatch::try_new(schema, interleaved_arrays) +fn with_schema(batch: RecordBatch, schema: SchemaRef) -> DataFusionResult { + RecordBatch::try_new(schema, batch.columns().to_vec()) .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) } -/// Utility function to take rows from a record batch. Based on [take](datafusion::arrow::compute::take) +/// Utility function to take ordered rows from a record batch. fn take_batch(batch: &RecordBatch, indices: &[usize]) -> DataFusionResult { - // fast path if batch.num_rows() == indices.len() { return Ok(batch.clone()); } - let schema = batch.schema(); - - let indices_array = UInt64Array::from_iter(indices.iter().map(|i| *i as u64)); + let indices_array = UInt64Array::from_iter(indices.iter().map(|index| *index as u64)); let arrays = batch .columns() .iter() .map(|array| compute::take(array, &indices_array, None)) .collect::, _>>() - .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; + .map_err(|error| DataFusionError::ArrowError(Box::new(error), None))?; - let result = RecordBatch::try_new(schema, arrays) - .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; - Ok(result) + RecordBatch::try_new(batch.schema(), arrays) + .map_err(|error| DataFusionError::ArrowError(Box::new(error), None)) } #[cfg(test)] mod test { + use std::collections::VecDeque; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::task::{Wake, Waker}; - use datafusion::arrow::array::Int32Array; - use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::array::{Array, Float64Array, Int32Array, Int64Array, StringArray}; + use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::common::ToDFSchema; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; use datafusion::logical_expr::{EmptyRelation, LogicalPlan}; + use datafusion::prelude::SessionContext; + use futures::StreamExt; use super::*; @@ -660,13 +642,7 @@ mod test { produce_one_row: false, schema: df_schema.clone(), }); - let plan = UnionDistinctOn::new( - left, - right, - vec!["k".to_string()], - "ts".to_string(), - df_schema, - ); + let plan = UnionDistinctOn::try_new(left, right, vec![1], 0).unwrap(); // Simulate a parent projection requesting only one output column. let output_columns = [2usize]; @@ -676,56 +652,6 @@ mod test { assert_eq!(required[1].as_slice(), &[0, 1, 2]); } - #[test] - fn test_interleave_batches() { - let schema = Schema::new(vec![ - Field::new("a", DataType::Int32, false), - Field::new("b", DataType::Int32, false), - ]); - - let batch1 = RecordBatch::try_new( - Arc::new(schema.clone()), - vec![ - Arc::new(Int32Array::from(vec![1, 2, 3])), - Arc::new(Int32Array::from(vec![4, 5, 6])), - ], - ) - .unwrap(); - - let batch2 = RecordBatch::try_new( - Arc::new(schema.clone()), - vec![ - Arc::new(Int32Array::from(vec![7, 8, 9])), - Arc::new(Int32Array::from(vec![10, 11, 12])), - ], - ) - .unwrap(); - - let batch3 = RecordBatch::try_new( - Arc::new(schema.clone()), - vec![ - Arc::new(Int32Array::from(vec![13, 14, 15])), - Arc::new(Int32Array::from(vec![16, 17, 18])), - ], - ) - .unwrap(); - - let batches = vec![batch1, batch2, batch3]; - let indices = vec![(0, 0), (1, 0), (2, 0), (0, 1), (1, 1), (2, 1)]; - let result = interleave_batches(Arc::new(schema.clone()), batches, indices).unwrap(); - - let expected = RecordBatch::try_new( - Arc::new(schema), - vec![ - Arc::new(Int32Array::from(vec![1, 7, 13, 2, 8, 14])), - Arc::new(Int32Array::from(vec![4, 10, 16, 5, 11, 17])), - ], - ) - .unwrap(); - - assert_eq!(result, expected); - } - #[test] fn test_take_batch() { let schema = Schema::new(vec![ @@ -757,39 +683,958 @@ mod test { assert_eq!(result, expected); } + fn empty_plan(schema: datafusion::common::DFSchemaRef) -> LogicalPlan { + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema, + }) + } + + fn schemas(left: SchemaRef, right: SchemaRef) -> (LogicalPlan, LogicalPlan) { + ( + empty_plan(left.to_dfschema_ref().unwrap()), + empty_plan(right.to_dfschema_ref().unwrap()), + ) + } + + fn schema(prefix: &str, key: &str, nullable: bool) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(format!("{prefix}_ts"), DataType::Int64, false), + Field::new(format!("{prefix}_{key}"), DataType::Utf8, nullable), + Field::new(format!("{prefix}_value"), DataType::Float64, nullable), + ])) + } + + fn batch(schema: SchemaRef, ts: i64, key: Option<&str>, value: Option) -> RecordBatch { + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![ts])), + Arc::new(StringArray::from(vec![key])), + Arc::new(Float64Array::from(vec![value])), + ], + ) + .unwrap() + } + + fn source_exec(batch: RecordBatch) -> Arc { + source_exec_batches(batch.schema(), vec![batch]) + } + + async fn execute( + plan: &UnionDistinctOn, + left: RecordBatch, + right: RecordBatch, + ) -> Vec { + datafusion::physical_plan::collect( + plan.to_execution_plan(source_exec(left), source_exec(right)), + SessionContext::default().task_ctx(), + ) + .await + .unwrap() + } + + fn source_exec_batches(schema: SchemaRef, batches: Vec) -> Arc { + Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[batches], schema, None).unwrap(), + ))) + } + + async fn execute_batches( + plan: &UnionDistinctOn, + left_schema: SchemaRef, + left: Vec, + right_schema: SchemaRef, + right: Vec, + ) -> Vec { + datafusion::physical_plan::collect( + plan.to_execution_plan( + source_exec_batches(left_schema, left), + source_exec_batches(right_schema, right), + ), + SessionContext::default().task_ctx(), + ) + .await + .unwrap() + } + + fn simple_schema(prefix: &str, nullable: bool) -> SchemaRef { + schema(prefix, "label", nullable) + } + + fn simple_batch(schema: SchemaRef, rows: &[(i64, &str, f64)]) -> RecordBatch { + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from_iter_values(rows.iter().map(|row| row.0))), + Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.1))), + Arc::new(Float64Array::from_iter_values(rows.iter().map(|row| row.2))), + ], + ) + .unwrap() + } + + fn simple_rows(batches: &[RecordBatch]) -> Vec<(i64, String, f64)> { + batches + .iter() + .flat_map(|batch| { + let timestamps = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let labels = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(move |row| { + ( + timestamps.value(row), + labels.value(row).to_string(), + values.value(row), + ) + }) + }) + .collect() + } + + fn simple_plan(left_schema: SchemaRef, right_schema: SchemaRef) -> UnionDistinctOn { + let (left, right) = schemas(left_schema, right_schema); + UnionDistinctOn::try_new(left, right, vec![1], 0).unwrap() + } + + #[derive(Clone, Debug)] + enum TestEvent { + Batch(RecordBatch), + Error(&'static str), + } + + struct TestStream { + schema: SchemaRef, + events: VecDeque, + polls: Arc, + } + + impl Stream for TestStream { + type Item = DataFusionResult; + + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + self.polls.fetch_add(1, Ordering::SeqCst); + match self.events.pop_front() { + Some(TestEvent::Batch(batch)) => Poll::Ready(Some(Ok(batch))), + Some(TestEvent::Error(message)) => { + Poll::Ready(Some(Err(DataFusionError::Internal(message.to_string())))) + } + None => Poll::Ready(None), + } + } + } + + impl RecordBatchStream for TestStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + } + + struct CountingWake(AtomicUsize); + + impl Wake for CountingWake { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::SeqCst); + } + + fn wake_by_ref(self: &Arc) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + fn test_stream( + schema: SchemaRef, + events: Vec, + polls: Arc, + ) -> SendableRecordBatchStream { + Box::pin(TestStream { + schema, + events: events.into(), + polls, + }) + } + + #[derive(Debug)] + struct TestExec { + schema: SchemaRef, + events: Vec, + polls: Arc, + executions: Arc, + execute_error: Option<&'static str>, + properties: Arc, + } + + impl TestExec { + fn new( + schema: SchemaRef, + events: Vec, + polls: Arc, + executions: Arc, + ) -> Self { + Self { + properties: Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )), + schema, + events, + polls, + executions, + execute_error: None, + } + } + + fn with_execute_error(mut self, error: &'static str) -> Self { + self.execute_error = Some(error); + self + } + } + + impl DisplayAs for TestExec { + fn fmt_as(&self, _t: DisplayFormatType, _f: &mut std::fmt::Formatter) -> std::fmt::Result { + Ok(()) + } + } + + impl ExecutionPlan for TestExec { + fn name(&self) -> &str { + "TestExec" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DataFusionResult> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> DataFusionResult { + self.executions.fetch_add(1, Ordering::SeqCst); + if let Some(error) = self.execute_error { + return Err(DataFusionError::Internal(error.to_string())); + } + Ok(test_stream( + self.schema.clone(), + self.events.clone(), + self.polls.clone(), + )) + } + } + + fn test_union_stream( + left: SendableRecordBatchStream, + right: Arc, + output_schema: SchemaRef, + ) -> UnionDistinctOnStream { + let metrics = ExecutionPlanMetricsSet::new(); + UnionDistinctOnStream { + left, + right_plan: right, + right_partition: 0, + right_context: SessionContext::default().task_ctx(), + right: None, + compare_keys: vec![1, 0], + output_schema, + random_state: RandomState::new(), + lhs_signatures: HashSet::default(), + hashes: Vec::new(), + phase: StreamPhase::Left, + metric: BaselineMetrics::new(&metrics, 0), + } + } + + fn series_rows(batches: &[RecordBatch]) -> Vec<(i64, &str, f64, &str)> { + batches + .iter() + .flat_map(|batch| { + let timestamps = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let series = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let keys = batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(move |row| { + ( + timestamps.value(row), + series.value(row), + values.value(row), + keys.value(row), + ) + }) + }) + .collect() + } + #[tokio::test] - async fn encode_decode_union_distinct_on() { + async fn serialize_deserialize_and_execute_with_different_input_names() { + let (left_schema, right_schema) = + (schema("left", "job", false), schema("right", "job", false)); + let (left_plan, right_plan) = schemas(left_schema.clone(), right_schema.clone()); + let decoded = UnionDistinctOn::deserialize( + &UnionDistinctOn::try_new(left_plan.clone(), right_plan.clone(), vec![1], 0) + .unwrap() + .serialize(), + ) + .unwrap() + .with_exprs_and_inputs(vec![], vec![left_plan, right_plan]) + .unwrap(); + assert_eq!( + (decoded.compare_key_indices.as_slice(), decoded.ts_col_idx), + (&[1usize][..], 0) + ); + assert_eq!( + decoded.output_schema, + left_schema.clone().to_dfschema_ref().unwrap() + ); + let result = execute( + &decoded, + batch(left_schema.clone(), 1, Some("left"), Some(10.0)), + batch(right_schema, 2, Some("right"), Some(20.0)), + ) + .await; + assert!(result.iter().all(|batch| batch.schema() == left_schema)); + } + + #[tokio::test] + async fn execute_widens_nullable_fields_and_emits_rhs_nulls() { + let (left_schema, right_schema) = ( + schema("left", "label", false), + schema("right", "label", true), + ); + let (left_plan, right_plan) = schemas(left_schema.clone(), right_schema.clone()); + let plan = UnionDistinctOn::try_new(left_plan, right_plan, vec![1, 2], 0).unwrap(); + let declared_schema = plan.output_schema.inner().clone(); + assert!( + declared_schema.fields()[1..] + .iter() + .all(|field| field.is_nullable()) + ); + let result = execute( + &plan, + batch(left_schema, 1, Some("present"), Some(10.0)), + batch(right_schema, 2, None, None), + ) + .await; + assert_eq!(result.len(), 2); + assert!(result.iter().all(|batch| batch.schema() == declared_schema)); + assert!(result[1].column(1).is_null(0)); + assert!(result[1].column(2).is_null(0)); + } + + #[tokio::test] + async fn empty_lhs_preserves_distinct_rhs_series_with_same_normalized_key() { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("ts", DataType::Int64, false), + Field::new("series", DataType::Utf8, false), + Field::new("value", DataType::Float64, false), + Field::new("__normalized_absent_label", DataType::Utf8, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("rhs_ts", DataType::Int64, false), + Field::new("rhs_series", DataType::Utf8, false), + Field::new("rhs_value", DataType::Float64, false), + Field::new("rhs_normalized_absent_label", DataType::Utf8, false), + ])); + let (left, right) = schemas(left_schema.clone(), right_schema.clone()); + let plan = UnionDistinctOn::try_new(left, right, vec![3], 0).unwrap(); + let declared_schema = plan.output_schema.inner().clone(); + let rhs = RecordBatch::try_new( + right_schema, + vec![ + Arc::new(Int64Array::from(vec![1_000, 1_000])), + Arc::new(StringArray::from(vec!["series_a", "series_b"])), + Arc::new(Float64Array::from(vec![10.0, 20.0])), + Arc::new(StringArray::from(vec!["", ""])), + ], + ) + .unwrap(); + + let result = execute(&plan, RecordBatch::new_empty(left_schema), rhs).await; + assert!(result.iter().all(|batch| batch.schema() == declared_schema)); + assert_eq!(result.iter().map(RecordBatch::num_rows).sum::(), 2); + + let mut rows = series_rows(&result); + rows.sort_by_key(|row| row.1); + assert_eq!( + rows, + vec![(1_000, "series_a", 10.0, ""), (1_000, "series_b", 20.0, "")] + ); + } + + #[tokio::test] + async fn lhs_signature_suppresses_all_matching_rhs_rows_and_preserves_lhs() { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("ts", DataType::Int64, false), + Field::new("series", DataType::Utf8, false), + Field::new("value", DataType::Float64, false), + Field::new("__normalized_absent_label", DataType::Utf8, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("rhs_ts", DataType::Int64, false), + Field::new("rhs_series", DataType::Utf8, false), + Field::new("rhs_value", DataType::Float64, false), + Field::new("rhs_normalized_absent_label", DataType::Utf8, false), + ])); + let (left, right) = schemas(left_schema.clone(), right_schema.clone()); + let plan = UnionDistinctOn::try_new(left, right, vec![3], 0).unwrap(); + let declared_schema = plan.output_schema.inner().clone(); + let lhs = RecordBatch::try_new( + left_schema, + vec![ + Arc::new(Int64Array::from(vec![1_000, 2_000])), + Arc::new(StringArray::from(vec!["lhs_match", "lhs_only"])), + Arc::new(Float64Array::from(vec![1.0, 2.0])), + Arc::new(StringArray::from(vec!["", "left-only"])), + ], + ) + .unwrap(); + let rhs = RecordBatch::try_new( + right_schema, + vec![ + Arc::new(Int64Array::from(vec![1_000, 1_000, 3_000])), + Arc::new(StringArray::from(vec![ + "rhs_match_a", + "rhs_match_b", + "rhs_unmatched", + ])), + Arc::new(Float64Array::from(vec![10.0, 20.0, 30.0])), + Arc::new(StringArray::from(vec!["", "", ""])), + ], + ) + .unwrap(); + + let result = execute(&plan, lhs, rhs).await; + assert_eq!(result.iter().map(RecordBatch::num_rows).sum::(), 3); + assert!(result.iter().all(|batch| batch.schema() == declared_schema)); + let mut rows = series_rows(&result); + rows.sort_by_key(|row| row.1); + assert_eq!( + rows, + vec![ + (1_000, "lhs_match", 1.0, ""), + (2_000, "lhs_only", 2.0, "left-only"), + (3_000, "rhs_unmatched", 30.0, ""), + ] + ); + } + + #[tokio::test] + async fn empty_lhs_preserves_duplicate_rhs_rows_across_batches() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + let result = execute_batches( + &plan, + left_schema, + vec![], + right_schema.clone(), + vec![ + simple_batch( + right_schema.clone(), + &[(1, "same", 10.0), (1, "same", 11.0)], + ), + simple_batch(right_schema, &[(1, "same", 12.0)]), + ], + ) + .await; + assert_eq!( + simple_rows(&result), + vec![ + (1, "same".to_string(), 10.0), + (1, "same".to_string(), 11.0), + (1, "same".to_string(), 12.0), + ] + ); + } + + #[tokio::test] + async fn lhs_signature_suppresses_rhs_duplicates_across_batches() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + let result = execute_batches( + &plan, + left_schema.clone(), + vec![simple_batch(left_schema, &[(1, "match", 1.0)])], + right_schema.clone(), + vec![ + simple_batch(right_schema.clone(), &[(1, "match", 10.0)]), + simple_batch(right_schema, &[(1, "match", 11.0)]), + ], + ) + .await; + assert_eq!(simple_rows(&result), vec![(1, "match".to_string(), 1.0)]); + } + + #[tokio::test] + async fn hash_scratch_reset_suppresses_null_label_after_same_sized_lhs_batches() { + let left_schema = simple_schema("left", true); + let right_schema = simple_schema("right", true); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + let lhs_poison = RecordBatch::try_new( + left_schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1])), + Arc::new(StringArray::from(vec![Some("poison")])), + Arc::new(Float64Array::from(vec![1.0])), + ], + ) + .unwrap(); + let lhs_null = RecordBatch::try_new( + left_schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![42])), + Arc::new(StringArray::from(vec![None::<&str>])), + Arc::new(Float64Array::from(vec![2.0])), + ], + ) + .unwrap(); + let rhs_null = RecordBatch::try_new( + right_schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![42])), + Arc::new(StringArray::from(vec![None::<&str>])), + Arc::new(Float64Array::from(vec![999.0])), + ], + ) + .unwrap(); + + let result = execute_batches( + &plan, + left_schema, + vec![lhs_poison, lhs_null], + right_schema, + vec![rhs_null], + ) + .await; + assert_eq!(result.iter().map(RecordBatch::num_rows).sum::(), 2); + let values = result + .iter() + .flat_map(|batch| { + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + assert_eq!(values, vec![1.0, 2.0]); + } + + #[tokio::test] + async fn mixed_rhs_duplicates_retain_unmatched_rows_in_order() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + let result = execute_batches( + &plan, + left_schema.clone(), + vec![simple_batch(left_schema, &[(1, "match", 1.0)])], + right_schema.clone(), + vec![simple_batch( + right_schema, + &[ + (1, "match", 10.0), + (1, "keep", 11.0), + (1, "keep", 12.0), + (1, "match", 13.0), + (1, "later", 14.0), + ], + )], + ) + .await; + assert_eq!( + simple_rows(&result), + vec![ + (1, "match".to_string(), 1.0), + (1, "keep".to_string(), 11.0), + (1, "keep".to_string(), 12.0), + (1, "later".to_string(), 14.0), + ] + ); + } + + #[tokio::test] + async fn same_label_at_different_timestamps_is_unmatched() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + let result = execute_batches( + &plan, + left_schema.clone(), + vec![simple_batch(left_schema, &[(1, "label", 1.0)])], + right_schema.clone(), + vec![simple_batch(right_schema, &[(2, "label", 2.0)])], + ) + .await; + assert_eq!( + simple_rows(&result), + vec![(1, "label".to_string(), 1.0), (2, "label".to_string(), 2.0)] + ); + } + + #[tokio::test] + async fn empty_input_combinations_complete_cleanly() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + + let empty_left = execute_batches( + &plan, + left_schema.clone(), + vec![], + right_schema.clone(), + vec![simple_batch(right_schema.clone(), &[(1, "rhs", 1.0)])], + ) + .await; + assert_eq!(simple_rows(&empty_left), vec![(1, "rhs".to_string(), 1.0)]); + + let empty_right = execute_batches( + &plan, + left_schema.clone(), + vec![simple_batch(left_schema.clone(), &[(1, "lhs", 1.0)])], + right_schema.clone(), + vec![], + ) + .await; + assert_eq!(simple_rows(&empty_right), vec![(1, "lhs".to_string(), 1.0)]); + + let both_empty = execute_batches(&plan, left_schema, vec![], right_schema, vec![]).await; + assert!(both_empty.is_empty()); + } + + #[tokio::test] + async fn zero_row_batches_on_either_side_are_handled() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + + let zero_row_left = execute_batches( + &plan, + left_schema.clone(), + vec![simple_batch(left_schema.clone(), &[])], + right_schema.clone(), + vec![simple_batch(right_schema.clone(), &[(1, "rhs", 1.0)])], + ) + .await; + assert_eq!( + simple_rows(&zero_row_left), + vec![(1, "rhs".to_string(), 1.0)] + ); + + let zero_row_right = execute_batches( + &plan, + left_schema.clone(), + vec![simple_batch(left_schema, &[(1, "lhs", 1.0)])], + right_schema.clone(), + vec![simple_batch(right_schema, &[])], + ) + .await; + assert_eq!( + simple_rows(&zero_row_right), + vec![(1, "lhs".to_string(), 1.0)] + ); + } + + #[tokio::test] + async fn metrics_record_output_rows_and_completion() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + let exec = plan.to_execution_plan( + source_exec(simple_batch(left_schema, &[(1, "lhs", 1.0)])), + source_exec(simple_batch(right_schema, &[(2, "rhs", 2.0)])), + ); + + let output = + datafusion::physical_plan::collect(exec.clone(), SessionContext::default().task_ctx()) + .await + .unwrap(); + assert_eq!(simple_rows(&output).len(), 2); + + let metrics = exec.metrics().unwrap(); + assert_eq!(metrics.output_rows(), Some(2)); + assert!(metrics.iter().any(|metric| { + metric.value().name() == "end_timestamp" && metric.value().as_usize() > 0 + })); + } + + #[tokio::test] + async fn output_keeps_lhs_before_rhs_and_input_order() { + let left_schema = simple_schema("left", false); + let right_schema = simple_schema("right", false); + let plan = simple_plan(left_schema.clone(), right_schema.clone()); + let result = execute_batches( + &plan, + left_schema.clone(), + vec![ + simple_batch(left_schema.clone(), &[(1, "lhs-a", 1.0)]), + simple_batch(left_schema, &[(2, "lhs-b", 2.0)]), + ], + right_schema.clone(), + vec![ + simple_batch(right_schema.clone(), &[(3, "rhs-a", 3.0)]), + simple_batch(right_schema, &[(4, "rhs-b", 4.0)]), + ], + ) + .await; + assert_eq!( + simple_rows(&result), + vec![ + (1, "lhs-a".to_string(), 1.0), + (2, "lhs-b".to_string(), 2.0), + (3, "rhs-a".to_string(), 3.0), + (4, "rhs-b".to_string(), 4.0), + ] + ); + } + + #[test] + fn fully_filtered_rhs_batch_wakes_before_later_unmatched_batch() { + let schema = simple_schema("stream", false); + let left_polls = Arc::new(AtomicUsize::new(0)); + let right_polls = Arc::new(AtomicUsize::new(0)); + let right_executions = Arc::new(AtomicUsize::new(0)); + let left = test_stream( + schema.clone(), + vec![TestEvent::Batch(simple_batch( + schema.clone(), + &[(1, "match", 1.0)], + ))], + left_polls, + ); + let right = Arc::new(TestExec::new( + schema.clone(), + vec![ + TestEvent::Batch(simple_batch(schema.clone(), &[(1, "match", 2.0)])), + TestEvent::Batch(simple_batch(schema.clone(), &[(2, "keep", 3.0)])), + ], + right_polls.clone(), + right_executions.clone(), + )); + let mut stream = Box::pin(test_union_stream(left, right, schema)); + let wake = Arc::new(CountingWake(AtomicUsize::new(0))); + let waker = Waker::from(wake.clone()); + let mut context = Context::from_waker(&waker); + + assert!(matches!( + stream.as_mut().poll_next(&mut context), + Poll::Ready(Some(Ok(_))) + )); + assert!(matches!( + stream.as_mut().poll_next(&mut context), + Poll::Pending + )); + assert_eq!(right_executions.load(Ordering::SeqCst), 1); + assert_eq!(right_polls.load(Ordering::SeqCst), 1); + assert_eq!(wake.0.load(Ordering::SeqCst), 1); + assert!(matches!( + stream.as_mut().poll_next(&mut context), + Poll::Ready(Some(Ok(batch))) + if simple_rows(std::slice::from_ref(&batch)) + == vec![(2, "keep".to_string(), 3.0)] + )); + } + + #[tokio::test] + async fn rhs_is_not_polled_before_lhs_eof_and_drop_preserves_backpressure() { + let schema = simple_schema("stream", false); + let left_polls = Arc::new(AtomicUsize::new(0)); + let right_polls = Arc::new(AtomicUsize::new(0)); + let left = test_stream( + schema.clone(), + vec![TestEvent::Batch(simple_batch( + schema.clone(), + &[(1, "lhs", 1.0)], + ))], + left_polls.clone(), + ); + let right_executions = Arc::new(AtomicUsize::new(0)); + let right = Arc::new(TestExec::new( + schema.clone(), + vec![TestEvent::Batch(simple_batch( + schema.clone(), + &[(2, "rhs", 2.0)], + ))], + right_polls.clone(), + right_executions.clone(), + )); + let mut stream = Box::pin(test_union_stream(left, right, schema)); + + let first = stream.next().await.unwrap().unwrap(); + assert_eq!(simple_rows(&[first]), vec![(1, "lhs".to_string(), 1.0)]); + assert_eq!(left_polls.load(Ordering::SeqCst), 1); + assert_eq!(right_polls.load(Ordering::SeqCst), 0); + assert_eq!(right_executions.load(Ordering::SeqCst), 0); + drop(stream); + assert_eq!(right_polls.load(Ordering::SeqCst), 0); + assert_eq!(right_executions.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn lhs_and_delayed_rhs_errors_propagate() { + let schema = simple_schema("stream", false); + let polls = Arc::new(AtomicUsize::new(0)); + let mut left_error = Box::pin(test_union_stream( + test_stream( + schema.clone(), + vec![TestEvent::Error("left failed")], + polls.clone(), + ), + Arc::new(TestExec::new( + schema.clone(), + vec![], + polls.clone(), + Arc::new(AtomicUsize::new(0)), + )), + schema.clone(), + )); + assert!(left_error.next().await.unwrap().is_err()); + assert!(left_error.next().await.is_none()); + + let mut right_error = Box::pin(test_union_stream( + test_stream( + schema.clone(), + vec![TestEvent::Batch(simple_batch( + schema.clone(), + &[(1, "lhs", 1.0)], + ))], + polls.clone(), + ), + Arc::new(TestExec::new( + schema.clone(), + vec![TestEvent::Error("right failed")], + polls, + Arc::new(AtomicUsize::new(0)), + )), + schema.clone(), + )); + assert!(right_error.next().await.unwrap().is_ok()); + assert!(right_error.next().await.unwrap().is_err()); + assert!(right_error.next().await.is_none()); + + let right_execute_error = Arc::new( + TestExec::new( + schema.clone(), + vec![], + Arc::new(AtomicUsize::new(0)), + Arc::new(AtomicUsize::new(0)), + ) + .with_execute_error("right execute failed"), + ); + let mut right_execute_error = Box::pin(test_union_stream( + test_stream(schema.clone(), vec![], Arc::new(AtomicUsize::new(0))), + right_execute_error, + schema, + )); + assert!(right_execute_error.next().await.unwrap().is_err()); + assert!(right_execute_error.next().await.is_none()); + } + + #[tokio::test] + async fn rhs_partial_and_full_batches_rebind_to_declared_left_schema() { + let left_schema = schema("left", "label", false); + let right_schema = schema("right", "label", true); + let (left, right) = schemas(left_schema.clone(), right_schema.clone()); + let plan = UnionDistinctOn::try_new(left, right, vec![1], 0).unwrap(); + let declared_schema = plan.output_schema.inner().clone(); + let partial_rhs = RecordBatch::try_new( + right_schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(StringArray::from(vec![Some("match"), None])), + Arc::new(Float64Array::from(vec![Some(10.0), None])), + ], + ) + .unwrap(); + let full_rhs = batch(right_schema.clone(), 3, Some("full"), Some(30.0)); + let result = execute_batches( + &plan, + left_schema.clone(), + vec![batch(left_schema, 1, Some("match"), Some(1.0))], + right_schema, + vec![partial_rhs, full_rhs], + ) + .await; + assert!(result.iter().all(|batch| batch.schema() == declared_schema)); + assert_eq!(result.iter().map(RecordBatch::num_rows).sum::(), 3); + assert!(result[1].column(1).is_null(0)); + } + + #[test] + fn malformed_indices_and_incompatible_inputs_fail_before_execution() { let schema = Arc::new(Schema::new(vec![ Field::new("ts", DataType::Int64, false), Field::new("job", DataType::Utf8, false), - Field::new("value", DataType::Float64, false), ])); - let df_schema = schema.clone().to_dfschema_ref().unwrap(); - let left_plan = LogicalPlan::EmptyRelation(EmptyRelation { - produce_one_row: false, - schema: df_schema.clone(), - }); - let right_plan = LogicalPlan::EmptyRelation(EmptyRelation { - produce_one_row: false, - schema: df_schema.clone(), - }); - let plan_node = UnionDistinctOn::new( - left_plan.clone(), - right_plan.clone(), - vec!["job".to_string()], - "ts".to_string(), - df_schema.clone(), - ); - - let bytes = plan_node.serialize(); - - let union_distinct_on = UnionDistinctOn::deserialize(&bytes).unwrap(); - let union_distinct_on = union_distinct_on - .with_exprs_and_inputs(vec![], vec![left_plan, right_plan]) + let invalid = |compare_key_indices, ts_col_idx| { + let decoded = UnionDistinctOn::deserialize( + &pb::UnionDistinctOn { + compare_key_indices, + ts_col_idx, + } + .encode_to_vec(), + ) .unwrap(); - - assert_eq!(union_distinct_on.compare_keys, vec!["job".to_string()]); - assert_eq!(union_distinct_on.ts_col, "ts"); - assert_eq!(union_distinct_on.output_schema, df_schema); + let (left, right) = schemas(schema.clone(), schema.clone()); + decoded + .with_exprs_and_inputs(vec![], vec![left, right]) + .is_err() + }; + assert!(invalid(vec![2], 0)); + assert!(invalid(vec![1], 2)); + let incompatible_schema = Arc::new(Schema::new(vec![ + Field::new("other_ts", DataType::Int64, false), + Field::new("other_job", DataType::Int64, false), + ])); + let (left, right) = schemas(schema, incompatible_schema); + assert!(UnionDistinctOn::try_new(left, right, vec![1], 0).is_err()); } } diff --git a/src/promql/src/functions/aggr_over_time.rs b/src/promql/src/functions/aggr_over_time.rs index 518c8689eff..05b2b1b2e2d 100644 --- a/src/promql/src/functions/aggr_over_time.rs +++ b/src/promql/src/functions/aggr_over_time.rs @@ -43,7 +43,14 @@ pub fn avg_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Op display_name = prom_min_over_time )] pub fn min_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Option { - compute::min(values) + let mut valid_values = values.iter().flatten(); + let mut min = valid_values.next()?; + for value in valid_values { + if value < min || min.is_nan() { + min = value; + } + } + Some(min) } /// The maximum value of all points in the specified interval. @@ -53,7 +60,14 @@ pub fn min_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Op display_name = prom_max_over_time )] pub fn max_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Option { - compute::max(values) + let mut valid_values = values.iter().flatten(); + let mut max = valid_values.next()?; + for value in valid_values { + if value > max || max.is_nan() { + max = value; + } + } + Some(max) } /// The sum of all values in the specified interval. @@ -183,6 +197,60 @@ mod test { use super::*; use crate::functions::test_util::simple_range_udf_runner; + fn assert_over_time_value(actual: Option, expected: Option) { + match (actual, expected) { + (Some(actual), Some(expected)) if expected.is_nan() => assert!(actual.is_nan()), + (Some(actual), Some(expected)) => assert_eq!(actual, expected), + (None, None) => {} + (actual, expected) => panic!("expected {expected:?}, got {actual:?}"), + } + } + + fn assert_min_max( + values: Vec>, + expected_min: Option, + expected_max: Option, + ) { + let timestamps = TimestampMillisecondArray::from(vec![0; values.len()]); + let values = Float64Array::from(values); + + assert_over_time_value(min_over_time(×tamps, &values), expected_min); + assert_over_time_value(max_over_time(×tamps, &values), expected_max); + } + + #[test] + fn min_max_over_time_ignore_ordinary_nan_when_finite_values_exist() { + let ordinary_nan = f64::from_bits(0x7ff8_0000_0000_0000); + + assert_min_max( + vec![Some(ordinary_nan), Some(3.0), Some(-2.0)], + Some(-2.0), + Some(3.0), + ); + assert_min_max( + vec![Some(3.0), Some(ordinary_nan), Some(-2.0)], + Some(-2.0), + Some(3.0), + ); + assert_min_max( + vec![Some(-2.0), Some(3.0), Some(ordinary_nan)], + Some(-2.0), + Some(3.0), + ); + assert_min_max( + vec![Some(ordinary_nan), Some(ordinary_nan)], + Some(ordinary_nan), + Some(ordinary_nan), + ); + assert_min_max( + vec![Some(3.0), Some(-2.0), Some(1.0)], + Some(-2.0), + Some(3.0), + ); + assert_min_max(vec![], None, None); + assert_min_max(vec![None, None], None, None); + } + // build timestamp range and value range arrays for test fn build_test_range_arrays() -> (RangeArray, RangeArray) { let ts_array = Arc::new(TimestampMillisecondArray::from_iter( diff --git a/src/promql/src/lib.rs b/src/promql/src/lib.rs index 2de1c6508e6..78bd012c4d1 100644 --- a/src/promql/src/lib.rs +++ b/src/promql/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(map_try_insert)] - pub mod error; pub mod extension_plan; pub mod functions; diff --git a/src/query/src/analyze.rs b/src/query/src/analyze.rs index 6f6864ef2b9..e71a10e2f20 100644 --- a/src/query/src/analyze.rs +++ b/src/query/src/analyze.rs @@ -33,7 +33,7 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, accept, }; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; -use datafusion_common::{DataFusionError, internal_err}; +use datafusion_common::{DataFusionError, assert_eq_or_internal_err, internal_err}; use datafusion_physical_expr::{Distribution, EquivalenceProperties, Partitioning}; use futures::StreamExt; use serde::Serialize; @@ -172,8 +172,13 @@ impl ExecutionPlan for DistAnalyzeExec { self: Arc, mut children: Vec>, ) -> DfResult> { + assert_eq_or_internal_err!( + children.len(), + 1, + "DistAnalyzeExec requires exactly one child" + ); Ok(Arc::new(Self::new( - children.pop().unwrap(), + children.swap_remove(0), self.verbose, self.format, ))) @@ -381,3 +386,71 @@ impl Display for JsonMetrics { write!(f, "{}", serde_json::to_string(self).unwrap()) } } + +#[cfg(test)] +mod tests { + use datafusion::physical_plan::empty::EmptyExec; + + use super::*; + + fn empty_plan(name: &str) -> Arc { + Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new( + name, + DataType::Utf8, + true, + )])))) + } + + #[test] + fn qbs_dist_analyze_rejects_zero_children() { + let analyze = Arc::new(DistAnalyzeExec::new( + empty_plan("original"), + false, + AnalyzeFormat::TEXT, + )); + + assert!(ExecutionPlan::with_new_children(analyze, vec![]).is_err()); + } + + #[test] + fn qbs_dist_analyze_rejects_multiple_children() { + let analyze = Arc::new(DistAnalyzeExec::new( + empty_plan("original"), + false, + AnalyzeFormat::TEXT, + )); + + let result = ExecutionPlan::with_new_children( + analyze, + vec![empty_plan("first"), empty_plan("second")], + ); + + if let Ok(plan) = result { + let retained = plan + .as_any() + .downcast_ref::() + .unwrap() + .input() + .schema() + .field(0) + .name() + .clone(); + panic!("expected an arity error for multiple children, but retained `{retained}`"); + } + } + + #[test] + fn qbs_dist_analyze_accepts_exactly_one_child() { + let analyze = Arc::new(DistAnalyzeExec::new( + empty_plan("original"), + false, + AnalyzeFormat::TEXT, + )); + let replacement = empty_plan("replacement"); + + let rebuilt = ExecutionPlan::with_new_children(analyze, vec![replacement]).unwrap(); + let rebuilt = rebuilt.as_any().downcast_ref::().unwrap(); + + assert_eq!(rebuilt.input().schema().field(0).name(), "replacement"); + } +} diff --git a/src/query/src/dist_plan.rs b/src/query/src/dist_plan.rs index 1983e7824c7..969f00e7c9c 100644 --- a/src/query/src/dist_plan.rs +++ b/src/query/src/dist_plan.rs @@ -27,6 +27,7 @@ mod remote_dyn_filter_registry; pub use analyzer::{DistPlannerAnalyzer, DistPlannerOptions}; pub use filter_id::{FilterFingerprint, FilterId, ParseFilterIdError, RemoteDynFilterProducerId}; pub use merge_scan::{MergeScanExec, MergeScanLogicalPlan}; +pub(crate) use merge_sort::MergeSortExec; pub use planner::{DistExtensionPlanner, MergeSortExtensionPlanner}; pub use predicate_extractor::PredicateExtractor; pub use region_pruner::ConstraintPruner; diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 27a0d9196bb..126ee9cb068 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -554,6 +554,24 @@ impl MergeScanExec { return None; } + let hash_expr_col_names: HashSet<_> = hash_exprs + .iter() + .filter_map(|expr| { + expr.as_any() + .downcast_ref::() + .map(|col_expr| col_expr.name()) + }) + .collect(); + + let covers_all_partition_cols = self.partition_cols.values().all(|aliases| { + aliases + .iter() + .any(|col| hash_expr_col_names.contains(col.name())) + }); + if !covers_all_partition_cols { + return None; + } + let all_partition_col_aliases: HashSet<_> = self .partition_cols .values() @@ -1013,7 +1031,7 @@ mod tests { use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::filter_pushdown::ChildFilterPushdownResult; use datafusion_common::TableReference; - use datafusion_expr::{LogicalPlanBuilder, lit}; + use datafusion_expr::{LogicalPlanBuilder, col, lit}; use datafusion_physical_expr::Distribution; use datafusion_physical_expr::expressions::{ Column, DynamicFilterPhysicalExpr, lit as physical_lit, @@ -1033,6 +1051,65 @@ mod tests { QueryId::from(Uuid::from_u128(value)) } + fn merge_scan_exec_with_sorted_input( + region_count: u64, + target_partition: usize, + ) -> MergeScanExec { + let session_state = SessionStateBuilder::new().build(); + let plan = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("ts")]) + .unwrap() + .sort(vec![col("ts").sort(false, true)]) + .unwrap() + .build() + .unwrap(); + let schema = plan.schema().as_arrow().clone(); + let regions = (0..region_count) + .map(|region_number| RegionId::new(1024, region_number as u32)) + .collect(); + + MergeScanExec::new( + &session_state, + // The table name is not relevant to these ordering metadata tests; + // `MergeScanExec::new` requires one to model the production plan. + TableName::new("catalog", "schema", "table"), + regions, + plan, + &schema, + Arc::new(TestRegionQueryHandler), + QueryContext::arc(), + target_partition, + AliasMapping::new(), + None, + false, + ) + .unwrap() + } + + #[test] + fn merge_scan_does_not_advertise_ordering_when_partition_may_merge_regions() { + let exec = merge_scan_exec_with_sorted_input(3, 2); + + assert!( + exec.properties().output_ordering().is_none(), + "target_partition < region_count means one output partition may concatenate multiple sorted region streams" + ); + } + + #[test] + fn merge_scan_advertises_ordering_when_each_partition_reads_at_most_one_region() { + let exec = merge_scan_exec_with_sorted_input(3, 3); + + assert!(exec.properties().output_ordering().is_some()); + } + + #[test] + fn merge_scan_advertises_ordering_when_partitions_exceed_regions() { + let exec = merge_scan_exec_with_sorted_input(3, 4); + + assert!(exec.properties().output_ordering().is_some()); + } + #[test] fn remote_dyn_filter_region_query_context_registers_before_do_get() { let registry_manager = Arc::new(DynFilterRegistryManager::default()); diff --git a/src/query/src/dist_plan/merge_sort.rs b/src/query/src/dist_plan/merge_sort.rs index d9bc6119c08..2c8f4d9fd8a 100644 --- a/src/query/src/dist_plan/merge_sort.rs +++ b/src/query/src/dist_plan/merge_sort.rs @@ -16,11 +16,23 @@ //! `SortPreservingMergeExec` operator in datafusion //! +use std::any::Any; use std::fmt; use std::sync::Arc; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::metrics::MetricsSet; +use datafusion::physical_plan::projection::{ProjectionExec, make_with_child, update_ordering}; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, + Statistics, +}; use datafusion_common::{DataFusionError, Result}; use datafusion_expr::{Extension, LogicalPlan, SortExpr, UserDefinedLogicalNodeCore}; +use datafusion_physical_expr::{Distribution, LexOrdering, OrderingRequirements}; /// MergeSort Logical Plan, have same field as `Sort`, but indicate it is a merge sort, /// which assume each input partition is a sorted stream, and will use `SortPreserveingMergeExec` @@ -53,14 +65,249 @@ impl MergeSortLogicalPlan { node: Arc::new(self), }) } +} - /// Convert self to a [`Sort`] logical plan with same input and expressions - pub fn into_sort(self) -> LogicalPlan { - LogicalPlan::Sort(datafusion::logical_expr::Sort { - input: self.input.clone(), - expr: self.expr, - fetch: self.fetch, - }) +/// An opaque physical execution node for [`MergeSortLogicalPlan`]. +/// +/// It delegates execution and physical properties to DataFusion's +/// [`SortPreservingMergeExec`], but intentionally does not expose itself as a +/// `SortPreservingMergeExec`. `EnforceSorting` is allowed to replace a bare +/// `SortPreservingMergeExec` with `CoalescePartitionsExec` when the parent does +/// not require ordering. `MergeSortExec` represents the distributed TopK merge +/// stage itself, so later physical optimizer rules must not rewrite it into an +/// unordered fetch. +#[derive(Debug, Clone)] +pub(crate) struct MergeSortExec { + inner: SortPreservingMergeExec, +} + +impl MergeSortExec { + pub(crate) fn new( + ordering: LexOrdering, + input: Arc, + fetch: Option, + ) -> Self { + Self { + inner: SortPreservingMergeExec::new(ordering, input).with_fetch(fetch), + } + } + + fn input_with_fetch(&self, fetch: Option) -> Arc { + let input = Arc::clone(self.inner.input()); + if let Some(sort) = input.as_any().downcast_ref::() + && sort.preserve_partitioning() + && sort.expr() == self.inner.expr() + { + // Mirror DataFusion's bare SPM plan quality for distributed TopK: + // keep the parent `MergeSortExec(fetch)` as the global merge, and + // bound the partition-preserving child sort to the same local TopK. + // Local top-K is safe because every global top-K row must be within + // the top-K rows of its own input partition. + Arc::new(sort.with_fetch(fetch)) + } else { + input + } + } +} + +impl DisplayAs for MergeSortExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "MergeSortExec: [{}]", self.inner.expr())?; + if let Some(fetch) = self.inner.fetch() { + write!(f, ", fetch={fetch}")?; + } + Ok(()) + } + DisplayFormatType::TreeRender => { + if let Some(fetch) = self.inner.fetch() { + writeln!(f, "limit={fetch}")?; + } + + for (i, expr) in self.inner.expr().iter().enumerate() { + expr.fmt_sql(f)?; + if i != self.inner.expr().len() - 1 { + write!(f, ", ")?; + } + } + + Ok(()) + } + } + } +} + +impl ExecutionPlan for MergeSortExec { + fn name(&self) -> &str { + "MergeSortExec" + } + + /// Keeps this node intentionally opaque to DataFusion's type-specialized + /// optimizer rewrites. + /// + /// `MergeSortExec` delegates most behavior to DataFusion's + /// `SortPreservingMergeExec`, but it must not expose itself as that type. + /// DataFusion's `EnforceSorting` optimizer recognizes a bare + /// `SortPreservingMergeExec` via `as_any().downcast_ref::<...>()` and may + /// replace it with an unordered `CoalescePartitionsExec(fetch)` when the + /// parent does not require sorted output. + /// + /// That rewrite is valid for an ordinary SPM used only to satisfy parent + /// ordering, but not for GreptimeDB's distributed TopK merge stage. In a + /// scalar-subquery shape like `ORDER BY ts DESC LIMIT 1`, this node is the + /// operator that merges region-local TopK streams into the global TopK. + /// Replacing it with unordered coalescing can return a partial/latest row + /// from one region instead of the global latest row. + /// + /// `required_input_ordering()` separately tells DataFusion what ordering this + /// node needs from its child, so `EnforceSorting` can insert a `SortExec` + /// below `MergeSortExec` when `MergeScanExec` cannot preserve per-partition + /// ordering. This opacity is specifically about protecting the merge stage + /// itself from the `EnforceSorting` rewrite above. + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &Arc { + self.inner.properties() + } + + /// Forwards DataFusion's order-preserving scan hint through this wrapper. + /// + /// This mirrors `SortPreservingMergeExec::with_preserve_order()`: if the + /// child can produce an order-preserving variant, rebuild the same merge + /// stage on top of that child. The returned plan must stay a + /// `MergeSortExec`, not a bare SPM, so the distributed TopK merge remains + /// opaque to `EnforceSorting`'s SPM-specific rewrite. + fn with_preserve_order(&self, preserve_order: bool) -> Option> { + self.inner + .input() + .with_preserve_order(preserve_order) + .map(|new_input| { + Arc::new(Self::new( + self.inner.expr().clone(), + new_input, + self.inner.fetch(), + )) as Arc + }) + } + + fn required_input_distribution(&self) -> Vec { + self.inner.required_input_distribution() + } + + fn benefits_from_input_partitioning(&self) -> Vec { + self.inner.benefits_from_input_partitioning() + } + + /// Tells DataFusion that `MergeSortExec` requires each input partition to be + /// ordered. This is the contract that makes `EnforceSorting` insert a + /// `SortExec` below `MergeSortExec` when the input cannot preserve ordering. + /// + /// The opacity of `MergeSortExec::as_any`, not this requirement, is what + /// prevents DataFusion from rewriting the merge stage itself as a bare + /// `SortPreservingMergeExec`. + fn required_input_ordering(&self) -> Vec> { + vec![Some(OrderingRequirements::from(self.inner.expr().clone()))] + } + + fn maintains_input_order(&self) -> Vec { + self.inner.maintains_input_order() + } + + fn children(&self) -> Vec<&Arc> { + self.inner.children() + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + if children.len() != 1 { + return Err(DataFusionError::Internal(format!( + "MergeSortExec expects exactly one child, got {}", + children.len() + ))); + } + + Ok(Arc::new(Self::new( + self.inner.expr().clone(), + children.swap_remove(0), + self.inner.fetch(), + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.inner.execute(partition, context) + } + + fn metrics(&self) -> Option { + self.inner.metrics() + } + + fn partition_statistics(&self, partition: Option) -> Result { + self.inner.partition_statistics(partition) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + self.inner.cardinality_effect() + } + + /// Intentionally keeps DataFusion's generic limit pushdown disabled. + /// + /// `MergeSortExec` still supports its own global fetch through + /// `with_fetch()`. What we must not allow is pushing an external limit below + /// this required distributed TopK merge. DataFusion's limit pushdown rules + /// know how to treat a bare `SortPreservingMergeExec` as a + /// partition-combining node, but `MergeSortExec` is intentionally opaque to + /// those SPM-specific downcasts. Enabling generic limit pushdown without also + /// teaching the optimizer about this wrapper could return partition-local + /// rows instead of the global TopK. + fn supports_limit_pushdown(&self) -> bool { + false + } + + fn fetch(&self) -> Option { + self.inner.fetch() + } + + fn with_fetch(&self, limit: Option) -> Option> { + Some(Arc::new(Self::new( + self.inner.expr().clone(), + self.input_with_fetch(limit), + limit, + ))) + } + + /// Lets DataFusion push a projection below this merge when it can rewrite + /// the ordering expressions safely. + /// + /// This mirrors `SortPreservingMergeExec::try_swapping_with_projection()` + /// for plan quality, but re-wraps the result as `MergeSortExec` so the + /// distributed merge stage keeps its type identity and opacity. + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result>> { + if projection.expr().len() >= projection.input().schema().fields().len() { + return Ok(None); + } + + let Some(updated_exprs) = update_ordering(self.inner.expr().clone(), projection.expr())? + else { + return Ok(None); + }; + + Ok(Some(Arc::new(Self::new( + updated_exprs, + make_with_child(projection, self.inner.input())?, + self.inner.fetch(), + )))) } } @@ -127,3 +374,420 @@ pub fn merge_sort_transformer(plan: &LogicalPlan) -> Option { None } } + +#[cfg(test)] +mod tests { + use arrow_schema::{DataType, Field, Schema, SortOptions}; + use datafusion::physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{ + OrderPreservationContext, plan_with_order_breaking_variants, + }; + use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; + use datafusion::physical_plan::displayable; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_physical_expr::expressions::col as physical_col; + + use super::*; + + /// Test double that records DataFusion's preserve-order signal while + /// otherwise behaving like a transparent wrapper around its child. + #[derive(Debug, Clone)] + struct PreserveOrderProbeExec { + inner: Arc, + preserve_order: bool, + } + + impl PreserveOrderProbeExec { + fn new(inner: Arc) -> Self { + Self { + inner, + preserve_order: false, + } + } + } + + impl DisplayAs for PreserveOrderProbeExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "PreserveOrderProbeExec: preserve_order={}", + self.preserve_order + ) + } + } + + impl ExecutionPlan for PreserveOrderProbeExec { + fn name(&self) -> &str { + "PreserveOrderProbeExec" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &Arc { + self.inner.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.inner] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + if children.len() != 1 { + return Err(DataFusionError::Internal(format!( + "PreserveOrderProbeExec expects exactly one child, got {}", + children.len() + ))); + } + + Ok(Arc::new(Self { + inner: children.swap_remove(0), + preserve_order: self.preserve_order, + })) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.inner.execute(partition, context) + } + + fn with_preserve_order(&self, preserve_order: bool) -> Option> { + Some(Arc::new(Self { + inner: Arc::clone(&self.inner), + preserve_order, + })) + } + } + + fn test_ordering(schema: &Schema) -> LexOrdering { + LexOrdering::new([PhysicalSortExpr::new( + physical_col("ts", schema).unwrap(), + SortOptions { + descending: true, + nulls_first: false, + }, + )]) + .unwrap() + } + + #[test] + fn merge_sort_exec_is_opaque_and_preserves_topk_requirements() { + let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)])); + let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _; + let ordering = test_ordering(schema.as_ref()); + + let merge_sort = + Arc::new(MergeSortExec::new(ordering, input, Some(1))) as Arc; + + assert_eq!(merge_sort.name(), "MergeSortExec"); + assert!( + merge_sort + .as_any() + .downcast_ref::() + .is_none(), + "MergeSortExec must stay opaque to EnforceSorting's bare SortPreservingMerge rewrite" + ); + assert_eq!(merge_sort.fetch(), Some(1)); + assert!(!merge_sort.supports_limit_pushdown()); + assert!(merge_sort.required_input_ordering()[0].is_some()); + + let tree = displayable(merge_sort.as_ref()).tree_render().to_string(); + assert!(tree.contains("MergeSortExec")); + assert!(!tree.contains("SortPreservingMergeExec")); + + let fetched = merge_sort.with_fetch(Some(2)).unwrap(); + assert!(fetched.as_any().downcast_ref::().is_some()); + assert_eq!(fetched.fetch(), Some(2)); + } + + #[test] + fn merge_sort_exec_required_input_ordering_matches_spm() { + let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)])); + let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _; + let ordering = test_ordering(schema.as_ref()); + + let merge_sort = MergeSortExec::new(ordering.clone(), Arc::clone(&input), Some(1)); + let bare_spm = + SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1)); + + assert_eq!( + merge_sort.required_input_ordering(), + vec![Some(OrderingRequirements::from(ordering))], + "MergeSortExec must require locally sorted input partitions for the merge key" + ); + assert_eq!( + merge_sort.required_input_ordering(), + bare_spm.required_input_ordering(), + "MergeSortExec's child ordering contract should mirror SortPreservingMergeExec" + ); + assert_eq!( + merge_sort.maintains_input_order(), + bare_spm.maintains_input_order() + ); + } + + #[test] + fn merge_sort_exec_with_fetch_pushes_fetch_to_child_sort() { + let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)])); + let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _; + let ordering = test_ordering(schema.as_ref()); + let child_sort = + Arc::new(SortExec::new(ordering.clone(), input).with_preserve_partitioning(true)) + as Arc; + let merge_sort = MergeSortExec::new(ordering, child_sort, None); + + let fetched = merge_sort.with_fetch(Some(2)).unwrap(); + + assert!(fetched.as_any().downcast_ref::().is_some()); + assert_eq!(fetched.fetch(), Some(2)); + let child_sort = fetched.children()[0] + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(child_sort.fetch(), Some(2)); + assert!(child_sort.preserve_partitioning()); + } + + #[test] + fn merge_sort_exec_with_preserve_order_matches_spm_but_keeps_wrapper() { + let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)])); + let input = Arc::new(PreserveOrderProbeExec::new(Arc::new( + EmptyExec::new(schema.clone()).with_partitions(2), + ))) as _; + let ordering = test_ordering(schema.as_ref()); + + let bare_spm = + SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1)); + let preserved_spm = bare_spm.with_preserve_order(true).unwrap(); + assert!( + preserved_spm + .as_any() + .downcast_ref::() + .is_some(), + "bare SPM should rebuild as bare SPM" + ); + assert!( + preserved_spm.children()[0] + .as_any() + .downcast_ref::() + .unwrap() + .preserve_order + ); + + let merge_sort = MergeSortExec::new(ordering, input, Some(1)); + let preserved_merge_sort = merge_sort.with_preserve_order(true).unwrap(); + assert!( + preserved_merge_sort + .as_any() + .downcast_ref::() + .is_some(), + "MergeSortExec must rewrap the preserve-order child as MergeSortExec" + ); + assert!( + preserved_merge_sort + .as_any() + .downcast_ref::() + .is_none(), + "MergeSortExec must not expose a bare SPM after with_preserve_order" + ); + assert_eq!(preserved_merge_sort.fetch(), Some(1)); + assert_eq!( + preserved_merge_sort.required_input_ordering(), + preserved_spm.required_input_ordering(), + "preserve-order rewrite should keep the same SPM child-ordering contract" + ); + assert!( + preserved_merge_sort.children()[0] + .as_any() + .downcast_ref::() + .unwrap() + .preserve_order + ); + } + + #[test] + fn merge_sort_exec_projection_swap_matches_spm_but_keeps_wrapper() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("value", DataType::Int64, false), + Field::new("ts", DataType::Int64, false), + Field::new("tag", DataType::Utf8, false), + ])); + let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _; + let ordering = test_ordering(schema.as_ref()); + + let bare_spm = Arc::new( + SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1)), + ) as Arc; + let spm_projection = ProjectionExec::try_new( + vec![ + (physical_col("ts", schema.as_ref())?, "ts".to_string()), + (physical_col("tag", schema.as_ref())?, "tag".to_string()), + ], + Arc::clone(&bare_spm), + )?; + let swapped_spm = bare_spm + .try_swapping_with_projection(&spm_projection)? + .expect("SPM should accept a narrowing projection that preserves the sort key"); + assert!( + swapped_spm + .as_any() + .downcast_ref::() + .is_some(), + "bare SPM should rebuild as bare SPM" + ); + + let merge_sort = + Arc::new(MergeSortExec::new(ordering, input, Some(1))) as Arc; + let merge_projection = ProjectionExec::try_new( + vec![ + (physical_col("ts", schema.as_ref())?, "ts".to_string()), + (physical_col("tag", schema.as_ref())?, "tag".to_string()), + ], + Arc::clone(&merge_sort), + )?; + let swapped_merge_sort = merge_sort + .try_swapping_with_projection(&merge_projection)? + .expect("MergeSortExec should accept the same projection swap as SPM"); + + assert!( + swapped_merge_sort + .as_any() + .downcast_ref::() + .is_some(), + "MergeSortExec must rewrap projection swaps as MergeSortExec" + ); + assert!( + swapped_merge_sort + .as_any() + .downcast_ref::() + .is_none(), + "MergeSortExec must not expose a bare SPM after projection swap" + ); + assert_eq!(swapped_merge_sort.fetch(), Some(1)); + assert!( + swapped_merge_sort.children()[0] + .as_any() + .downcast_ref::() + .is_some(), + "the projection should move below MergeSortExec" + ); + let swapped_schema = swapped_merge_sort.schema(); + assert_eq!( + swapped_schema + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + vec!["ts", "tag"], + "swapped MergeSortExec should expose the projected schema" + ); + + let projected_ordering = LexOrdering::new([PhysicalSortExpr::new( + physical_col("ts", swapped_merge_sort.children()[0].schema().as_ref())?, + SortOptions { + descending: true, + nulls_first: false, + }, + )]) + .unwrap(); + assert_eq!( + swapped_merge_sort.required_input_ordering(), + vec![Some(OrderingRequirements::from(projected_ordering))], + "projection swap must rewrite the ordering to the child projection's schema" + ); + assert_eq!( + swapped_merge_sort.required_input_ordering(), + swapped_spm.required_input_ordering(), + "MergeSortExec projection swap should mirror SPM's ordering rewrite" + ); + + let spm_projection_without_sort_key = ProjectionExec::try_new( + vec![(physical_col("tag", schema.as_ref())?, "tag".to_string())], + Arc::clone(&bare_spm), + )?; + let merge_projection_without_sort_key = ProjectionExec::try_new( + vec![(physical_col("tag", schema.as_ref())?, "tag".to_string())], + Arc::clone(&merge_sort), + )?; + assert!( + bare_spm + .try_swapping_with_projection(&spm_projection_without_sort_key)? + .is_none(), + "SPM must reject projection swaps that drop the sort key" + ); + assert!( + merge_sort + .try_swapping_with_projection(&merge_projection_without_sort_key)? + .is_none(), + "MergeSortExec should reject the same projection swap as SPM" + ); + + Ok(()) + } + + #[test] + fn enforce_sorting_rewrite_keeps_merge_sort_exec_opaque() { + let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)])); + let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _; + let ordering = test_ordering(schema.as_ref()); + + let bare_spm = Arc::new( + SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1)), + ) as Arc; + let optimized_spm = plan_with_order_breaking_variants(OrderPreservationContext::new( + bare_spm, + false, + vec![OrderPreservationContext::new( + Arc::clone(&input), + false, + vec![], + )], + )) + .unwrap() + .plan; + assert!( + optimized_spm + .as_any() + .downcast_ref::() + .is_some(), + "this regression test must exercise EnforceSorting's bare SPM -> CoalescePartitionsExec rewrite" + ); + + let merge_sort = + Arc::new(MergeSortExec::new(ordering, input, Some(1))) as Arc; + let optimized_merge_sort = + plan_with_order_breaking_variants(OrderPreservationContext::new( + Arc::clone(&merge_sort), + false, + vec![OrderPreservationContext::new( + Arc::clone(merge_sort.children()[0]), + false, + vec![], + )], + )) + .unwrap() + .plan; + assert!( + optimized_merge_sort + .as_any() + .downcast_ref::() + .is_some(), + "MergeSortExec must stay opaque to the bare SPM rewrite" + ); + assert!( + optimized_merge_sort + .as_any() + .downcast_ref::() + .is_none(), + "MergeSortExec(fetch) is the required distributed TopK merge stage, not an unordered coalesce" + ); + assert_eq!(optimized_merge_sort.fetch(), Some(1)); + } +} diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index f613807b291..f3b7851fde3 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -17,6 +17,7 @@ use std::sync::Arc; use ahash::HashMap; +use arrow_schema::SortOptions; use async_trait::async_trait; use catalog::CatalogManagerRef; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; @@ -29,6 +30,7 @@ use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner}; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}; use datafusion_common::{DataFusionError, TableReference}; use datafusion_expr::{LogicalPlan, UserDefinedLogicalNode}; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; use partition::manager::{PartitionRuleManagerRef, create_partitions_from_region_routes}; use session::context::QueryContext; use snafu::{OptionExt, ResultExt}; @@ -39,27 +41,54 @@ use table::table_name::TableName; use crate::dist_plan::PredicateExtractor; use crate::dist_plan::merge_scan::{MergeScanExec, MergeScanLogicalPlan}; -use crate::dist_plan::merge_sort::MergeSortLogicalPlan; +use crate::dist_plan::merge_sort::{MergeSortExec, MergeSortLogicalPlan}; use crate::dist_plan::region_pruner::ConstraintPruner; use crate::error::{CatalogSnafu, PartitionRuleManagerSnafu, TableNotFoundSnafu}; use crate::region_query::RegionQueryHandlerRef; -/// Planner for convert merge sort logical plan to physical plan +/// Planner for converting merge sort logical plan to physical plan. /// -/// it is currently a fallback to sort, and doesn't change the execution plan: -/// `MergeSort(MergeScan) -> Sort(MergeScan) - to physical plan -> ...` -/// It should be applied after `DistExtensionPlanner` -/// -/// (Later when actually impl this merge sort) -/// -/// We should ensure the number of partition is not smaller than the number of region at present. Otherwise this would result in incorrect output. +/// `MergeSortExec` always represents the distributed merge stage. It declares +/// the required input ordering to DataFusion, so `EnforceSorting` inserts a +/// `SortExec` below it when the input `MergeScanExec` cannot preserve per-region +/// ordering, for example when one output partition may merge multiple region +/// streams. pub struct MergeSortExtensionPlanner {} +impl MergeSortExtensionPlanner { + fn ordering( + session_state: &SessionState, + merge_sort: &MergeSortLogicalPlan, + ) -> Result { + let ordering = merge_sort + .expr + .iter() + .map(|sort_expr| { + let physical_expr = session_state + .create_physical_expr(sort_expr.expr.clone(), merge_sort.input.schema())?; + Ok(PhysicalSortExpr::new( + physical_expr, + SortOptions { + descending: !sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }, + )) + }) + .collect::>>()?; + + LexOrdering::new(ordering).ok_or_else(|| { + DataFusionError::Internal( + "Expect MergeSort to have non-empty sort expressions".to_string(), + ) + }) + } +} + #[async_trait] impl ExtensionPlanner for MergeSortExtensionPlanner { async fn plan_extension( &self, - planner: &dyn PhysicalPlanner, + _planner: &dyn PhysicalPlanner, node: &dyn UserDefinedLogicalNode, _logical_inputs: &[&LogicalPlan], physical_inputs: &[Arc], @@ -73,28 +102,24 @@ impl ExtensionPlanner for MergeSortExtensionPlanner { .downcast_ref::() .is_some() { - let merge_scan_exec = physical_inputs - .first() - .and_then(|p| p.as_any().downcast_ref::()) - .ok_or(DataFusionError::Internal(format!( + let input = physical_inputs.first().cloned().ok_or_else(|| { + DataFusionError::Internal( + "Expect MergeSort to have one physical input".to_string(), + ) + })?; + if input.as_any().downcast_ref::().is_none() { + return Err(DataFusionError::Internal(format!( "Expect MergeSort's input is a MergeScanExec, found {:?}", physical_inputs - )))?; - - let partition_cnt = merge_scan_exec.partition_count(); - let region_cnt = merge_scan_exec.region_count(); - // if partition >= region, we know that every partition stream of merge scan is ordered - // and we only need to do a merge sort, otherwise fallback to quick sort - let can_merge_sort = partition_cnt >= region_cnt; - if can_merge_sort { - // TODO(discord9): use `SortPreservingMergeExec here` + ))); } - // for now merge sort only exist in logical plan, and have the same effect as `Sort` - // doesn't change the execution plan, this will change in the future - let ret = planner - .create_physical_plan(&merge_sort.clone().into_sort(), session_state) - .await?; - Ok(Some(ret)) + + let ordering = Self::ordering(session_state, merge_sort)?; + Ok(Some(Arc::new(MergeSortExec::new( + ordering, + input, + merge_sort.fetch, + )))) } else { Ok(None) } diff --git a/src/query/src/log_query/planner.rs b/src/query/src/log_query/planner.rs index f9c641b23b5..05aeb0795d2 100644 --- a/src/query/src/log_query/planner.rs +++ b/src/query/src/log_query/planner.rs @@ -34,8 +34,6 @@ use crate::log_query::error::{ UnknownTableSnafu, }; -const DEFAULT_LIMIT: usize = 1000; - pub struct LogQueryPlanner { table_provider: DfTableSourceProvider, session_state: SessionState, @@ -99,19 +97,18 @@ impl LogQueryPlanner { .context(DataFusionPlanningSnafu)?; } - // Apply limit - plan_builder = plan_builder - .limit( - query.limit.skip.unwrap_or(0), - Some(query.limit.fetch.unwrap_or(DEFAULT_LIMIT)), - ) - .context(DataFusionPlanningSnafu)?; - // Apply log expressions for expr in &query.exprs { plan_builder = self.process_log_expr(plan_builder, expr)?; } + // Apply pagination to the final result after all log expressions. + if query.limit.skip.is_some() || query.limit.fetch.is_some() { + plan_builder = plan_builder + .limit(query.limit.skip.unwrap_or(0), query.limit.fetch) + .context(DataFusionPlanningSnafu)?; + } + // Build the final plan let plan = plan_builder.build().context(DataFusionPlanningSnafu)?; @@ -126,15 +123,14 @@ impl LogQueryPlanner { .clone(); let start_time = ScalarValue::Utf8(time_filter.start.clone()); - let end_time = ScalarValue::Utf8( - time_filter - .end - .clone() - .or(Some("9999-12-31T23:59:59Z".to_string())), - ); let expr = col(timestamp_col.clone()) .gt_eq(lit(start_time)) - .and(col(timestamp_col).lt_eq(lit(end_time))); + .and(match &time_filter.end { + Some(end) => col(timestamp_col).lt(lit(ScalarValue::Utf8(Some(end.clone())))), + None => col(timestamp_col).lt_eq(lit(ScalarValue::Utf8(Some( + "9999-12-31T23:59:59Z".to_string(), + )))), + }); Ok(expr) } @@ -730,7 +726,7 @@ mod tests { let plan = planner.query_to_plan(log_query).await.unwrap(); let expected = "Limit: skip=0, fetch=100 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ -\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp <= Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ +\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); @@ -757,7 +753,7 @@ mod tests { .gt_eq(lit(ScalarValue::Utf8(Some( "2021-01-01T00:00:00Z".to_string(), )))) - .and(col("timestamp").lt_eq(lit(ScalarValue::Utf8(Some( + .and(col("timestamp").lt(lit(ScalarValue::Utf8(Some( "2021-01-02T00:00:00Z".to_string(), ))))); @@ -850,8 +846,8 @@ mod tests { }; let plan = planner.query_to_plan(log_query).await.unwrap(); - let expected = "Limit: skip=10, fetch=1000 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ -\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp <= Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ + let expected = "Limit: skip=10, fetch=None [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ +\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); @@ -885,9 +881,8 @@ mod tests { }; let plan = planner.query_to_plan(log_query).await.unwrap(); - let expected = "Limit: skip=0, fetch=1000 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ -\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp <= Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ -\n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; + let expected = "Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") AND greptime.public.test_table.message LIKE Utf8(\"%error%\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ +\n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); } @@ -932,9 +927,9 @@ mod tests { }; let plan = planner.query_to_plan(log_query).await.unwrap(); - let expected = "Aggregate: groupBy=[[greptime.public.test_table.host]], aggr=[[count(greptime.public.test_table.message) AS count_result]] [host:Utf8;N, count_result:Int64]\ -\n Limit: skip=0, fetch=100 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ -\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp <= Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ + let expected = "Limit: skip=0, fetch=100 [host:Utf8;N, count_result:Int64]\ +\n Aggregate: groupBy=[[greptime.public.test_table.host]], aggr=[[count(greptime.public.test_table.message) AS count_result]] [host:Utf8;N, count_result:Int64]\ +\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); @@ -972,9 +967,9 @@ mod tests { }; let plan = planner.query_to_plan(log_query).await.unwrap(); - let expected = "Projection: date_trunc(Utf8(\"day\"), greptime.public.test_table.timestamp) AS time_bucket [time_bucket:Timestamp(ms)]\ - \n Limit: skip=0, fetch=100 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ - \n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp <= Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ + let expected = "Limit: skip=0, fetch=100 [time_bucket:Timestamp(ms)]\ + \n Projection: date_trunc(Utf8(\"day\"), greptime.public.test_table.timestamp) AS time_bucket [time_bucket:Timestamp(ms)]\ + \n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); @@ -1055,10 +1050,10 @@ mod tests { }; let plan = planner.query_to_plan(log_query).await.unwrap(); - let expected = "Aggregate: groupBy=[[2__date_histogram__time_bucket]], aggr=[[count(2__date_histogram__time_bucket) AS count_result]] [2__date_histogram__time_bucket:Timestamp(ns);N, count_result:Int64]\ -\n Projection: date_bin(Utf8(\"30 seconds\"), greptime.public.test_table.timestamp) AS 2__date_histogram__time_bucket [2__date_histogram__time_bucket:Timestamp(ns);N]\ -\n Limit: skip=0, fetch=1000 [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ -\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp <= Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ + let expected = "Limit: skip=0, fetch=None [2__date_histogram__time_bucket:Timestamp(ns);N, count_result:Int64]\ +\n Aggregate: groupBy=[[2__date_histogram__time_bucket]], aggr=[[count(2__date_histogram__time_bucket) AS count_result]] [2__date_histogram__time_bucket:Timestamp(ns);N, count_result:Int64]\ +\n Projection: date_bin(Utf8(\"30 seconds\"), greptime.public.test_table.timestamp) AS 2__date_histogram__time_bucket [2__date_histogram__time_bucket:Timestamp(ns);N]\ +\n Filter: greptime.public.test_table.timestamp >= Utf8(\"2021-01-01T00:00:00Z\") AND greptime.public.test_table.timestamp < Utf8(\"2021-01-02T00:00:00Z\") [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]\ \n TableScan: greptime.public.test_table [message:Utf8, timestamp:Timestamp(ms), host:Utf8;N, is_active:Boolean;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); diff --git a/src/query/src/optimizer/global_limit.rs b/src/query/src/optimizer/global_limit.rs index 2e5350f7797..4b2ff721f64 100644 --- a/src/query/src/optimizer/global_limit.rs +++ b/src/query/src/optimizer/global_limit.rs @@ -25,6 +25,8 @@ use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use datafusion_common::Result as DfResult; use datafusion_physical_expr::{Distribution, OrderingRequirements, Partitioning}; +use crate::dist_plan::MergeSortExec; + #[derive(Debug)] pub struct EnsureGlobalLimitForFetch; @@ -149,7 +151,8 @@ fn provided_global_fetch(plan: &Arc) -> Option { let fetch = plan.fetch()?; (plan.as_any().is::() || plan.as_any().is::() - || plan.as_any().is::()) + || plan.as_any().is::() + || plan.as_any().is::()) .then_some(fetch) } @@ -327,6 +330,36 @@ mod tests { assert_eq!(child.fetch(), Some(5)); } + #[test] + fn keeps_filter_under_parent_merge_sort_fetch() { + let (input, ordering) = ordered_input(); + let filter = filter_fetch(input, 1); + let merge = merge_sort_fetch(ordering, filter, 1); + + let optimized = + EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap(); + let child = optimized.children()[0]; + + assert!(optimized.as_any().is::()); + assert!(child.as_any().is::()); + } + + #[test] + fn adds_tighter_global_fetch_under_looser_merge_sort_fetch() { + let (input, ordering) = ordered_input(); + let filter = filter_fetch(input, 5); + let merge = merge_sort_fetch(ordering, filter, 10); + + let optimized = + EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap(); + let child = optimized.children()[0]; + + assert!(optimized.as_any().is::()); + assert!(child.as_any().is::()); + assert_eq!(child.fetch(), Some(5)); + assert!(child.children()[0].as_any().is::()); + } + #[test] fn preserves_parent_ordering_requirement() { let (input, ordering) = ordered_input(); @@ -552,6 +585,14 @@ mod tests { ) } + fn merge_sort_fetch( + ordering: LexOrdering, + input: Arc, + fetch: usize, + ) -> Arc { + Arc::new(MergeSortExec::new(ordering, input, Some(fetch))) + } + fn hash_repartition(input: Arc) -> Arc { let partitioning = Partitioning::Hash(vec![col("a", input.schema().as_ref()).unwrap()], 3); Arc::new(RepartitionExec::try_new(input, partitioning).unwrap()) diff --git a/src/query/src/optimizer/pass_distribution.rs b/src/query/src/optimizer/pass_distribution.rs index 22e2b903593..8ff21e046d4 100644 --- a/src/query/src/optimizer/pass_distribution.rs +++ b/src/query/src/optimizer/pass_distribution.rs @@ -16,8 +16,9 @@ use std::sync::Arc; use datafusion::config::ConfigOptions; use datafusion::physical_optimizer::PhysicalOptimizerRule; -use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::{ExecutionPlan, Partitioning}; use datafusion_common::Result as DfResult; use datafusion_physical_expr::Distribution; use datafusion_physical_expr::utils::map_columns_before_projection; @@ -68,11 +69,26 @@ impl PassDistribution { ) -> DfResult> { // If this is a MergeScanExec, try to apply the current requirement. if let Some(merge_scan) = plan.as_any().downcast_ref::() - && let Some(distribution) = current_req.as_ref() - && let Some(new_plan) = merge_scan.try_with_new_distribution(distribution.clone()) + && let Some(Distribution::HashPartitioned(hash_exprs)) = current_req.as_ref() { - // Leaf node; no children to process - return Ok(Arc::new(new_plan) as _); + if let Partitioning::Hash(current_hash_exprs, _) = &merge_scan.properties().partitioning + && *current_hash_exprs == *hash_exprs + { + return Ok(plan); + } + + if let Some(new_plan) = merge_scan + .try_with_new_distribution(Distribution::HashPartitioned(hash_exprs.clone())) + { + // Leaf node; no children to process + return Ok(Arc::new(new_plan) as _); + } + + let partitioning = Partitioning::Hash( + hash_exprs.clone(), + merge_scan.properties().partitioning.partition_count(), + ); + return Ok(Arc::new(RepartitionExec::try_new(plan, partitioning)?) as _); } // Compute per-child requirements from the current node. @@ -266,6 +282,20 @@ mod tests { ); } + #[test] + fn merge_scan_rejects_hash_requirement_on_partition_key_subset() { + let merge_scan = test_merge_scan_exec(test_schema()); + + let new_plan = merge_scan.try_with_new_distribution(Distribution::HashPartitioned(vec![ + partition_column(DATA_SCHEMA_TSID_COLUMN_NAME, 1), + ])); + + assert!( + new_plan.is_none(), + "partitioning by a subset of multi-column partition keys is not sufficient" + ); + } + fn test_merge_scan_exec(schema: SchemaRef) -> MergeScanExec { let session_state = SessionStateBuilder::new().with_default_features().build(); let partition_cols = BTreeMap::from([ diff --git a/src/query/src/parser.rs b/src/query/src/parser.rs index 53f24557c73..87b827e0cc1 100644 --- a/src/query/src/parser.rs +++ b/src/query/src/parser.rs @@ -52,45 +52,37 @@ pub enum QueryStatement { } impl QueryStatement { - pub fn post_process(&self, params: HashMap) -> Result { + pub fn post_process(self, params: HashMap) -> Result { match self { QueryStatement::Sql(_) => UnimplementedSnafu { operation: "sql post process", } .fail(), - QueryStatement::Promql(eval_stmt, alias) => { + QueryStatement::Promql(mut eval_stmt, alias) => { let node_name = match params.get("name") { Some(name) => name.as_str(), None => "", }; - let extension_node = Self::create_extension_node(node_name, &eval_stmt.expr); - Ok(QueryStatement::Promql( - EvalStmt { - expr: Extension(extension_node.unwrap()), - start: eval_stmt.start, - end: eval_stmt.end, - interval: eval_stmt.interval, - lookback_delta: eval_stmt.lookback_delta, - }, - alias.clone(), - )) + let extension_node = Self::create_extension_node(node_name, eval_stmt.expr); + eval_stmt.expr = Extension(extension_node.unwrap()); + Ok(QueryStatement::Promql(eval_stmt, alias)) } } } - fn create_extension_node(node_name: &str, expr: &Expr) -> Option { + fn create_extension_node(node_name: &str, expr: Expr) -> Option { match node_name { ANALYZE_NODE_NAME => Some(NodeExtension { - expr: Arc::new(AnalyzeExpr { expr: expr.clone() }), + expr: Arc::new(AnalyzeExpr { expr }), }), ANALYZE_VERBOSE_NODE_NAME => Some(NodeExtension { - expr: Arc::new(AnalyzeVerboseExpr { expr: expr.clone() }), + expr: Arc::new(AnalyzeVerboseExpr { expr }), }), EXPLAIN_NODE_NAME => Some(NodeExtension { - expr: Arc::new(ExplainExpr { expr: expr.clone() }), + expr: Arc::new(ExplainExpr { expr }), }), EXPLAIN_VERBOSE_NODE_NAME => Some(NodeExtension { - expr: Arc::new(ExplainVerboseExpr { expr: expr.clone() }), + expr: Arc::new(ExplainVerboseExpr { expr }), }), _ => None, } @@ -204,9 +196,10 @@ impl QueryLanguageParser { } pub(crate) fn apply_alias_extension(mut eval_stmt: EvalStmt, alias: &str) -> EvalStmt { + let expr = eval_stmt.expr; eval_stmt.expr = Extension(NodeExtension { expr: Arc::new(AliasExpr { - expr: eval_stmt.expr.clone(), + expr, alias: alias.to_string(), }), }); @@ -263,6 +256,14 @@ macro_rules! define_node_ast_extension { fn children(&self) -> &[Expr] { std::slice::from_ref(&self.expr) } + + fn with_new_children(&self, children: Vec) -> Arc { + let mut iter = children.into_iter(); + match (iter.next(), iter.next()) { + (Some(expr), None) => Arc::new($name_expr { expr }), + _ => Arc::new(self.clone()), + } + } } #[allow(rustdoc::broken_intra_doc_links)] @@ -313,6 +314,16 @@ impl ExtensionExpr for AliasExpr { fn children(&self) -> &[Expr] { std::slice::from_ref(&self.expr) } + fn with_new_children(&self, children: Vec) -> Arc { + let mut iter = children.into_iter(); + match (iter.next(), iter.next()) { + (Some(expr), None) => Arc::new(Self { + expr, + alias: self.alias.clone(), + }), + _ => Arc::new(self.clone()), + } + } } #[derive(Debug, Clone)] pub struct Alias { diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index c2cd6316ada..284701f3375 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -50,6 +50,7 @@ use datafusion_expr::utils::conjunction; use datafusion_expr::{ ExprSchemable, Literal, Projection, SortExpr, TableScan, TableSource, col, lit, }; +use datafusion_functions::core::coalesce; use datatypes::arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit}; use datatypes::data_type::ConcreteDataType; use itertools::Itertools; @@ -229,6 +230,8 @@ impl IslandExpr { !modifier.return_bool && modifier.matching.is_none() && matches!(modifier.card, VectorMatchCardinality::OneToOne) + && modifier.fill_values.lhs.is_none() + && modifier.fill_values.rhs.is_none() }) => { let lhs = Self::try_new(lhs, env)?; @@ -1093,6 +1096,18 @@ impl PromPlanner { query_engine_state: &QueryEngineState, binary_expr: &PromBinaryExpr, ) -> Result { + // promql-parser accepts fill modifiers, but Greptime does not implement the + // required outer joins and missing-value substitution. Reject them before the + // binary-island fast path so they cannot silently behave like normal inner joins. + if let Some(modifier) = &binary_expr.modifier { + ensure!( + modifier.fill_values.lhs.is_none() && modifier.fill_values.rhs.is_none(), + UnsupportedExprSnafu { + name: "PromQL fill modifiers" + } + ); + } + if let Some(plan) = self.try_plan_binary_island(binary_expr).await? { return Ok(plan); } @@ -4409,6 +4424,59 @@ impl PromPlanner { // Take the name of first field column. The length is checked above. let left_field_col = left_context.field_columns.first().unwrap(); let right_field_col = right_context.field_columns.first().unwrap(); + let left_field = left + .schema() + .iter() + .find(|(_, field)| field.name() == left_field_col) + .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { + col: left_field_col.clone(), + })?; + let right_field = right + .schema() + .iter() + .find(|(_, field)| field.name() == right_field_col) + .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { + col: right_field_col.clone(), + })?; + let target_field_type = if left_field.1 == right_field.1 { + left_field.1.clone() + } else if matches!( + left_field.1, + ArrowDataType::Int8 + | ArrowDataType::Int16 + | ArrowDataType::Int32 + | ArrowDataType::Int64 + | ArrowDataType::UInt8 + | ArrowDataType::UInt16 + | ArrowDataType::UInt32 + | ArrowDataType::UInt64 + | ArrowDataType::Float32 + | ArrowDataType::Float64 + ) && matches!( + right_field.1, + ArrowDataType::Int8 + | ArrowDataType::Int16 + | ArrowDataType::Int32 + | ArrowDataType::Int64 + | ArrowDataType::UInt8 + | ArrowDataType::UInt16 + | ArrowDataType::UInt32 + | ArrowDataType::UInt64 + | ArrowDataType::Float32 + | ArrowDataType::Float64 + ) { + ArrowDataType::Float64 + } else { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR value fields have incompatible types: {:?} and {:?}", + left_field.1, right_field.1 + ), + } + .fail(); + }; let left_has_tsid = left .schema() .fields() @@ -4445,10 +4513,26 @@ impl PromPlanner { all_columns.sort_unstable(); // use left time index column name as the result time index column name all_columns.insert(0, left_time_index_column.clone()); + let mut occupied_column_names = left + .schema() + .fields() + .iter() + .chain(right.schema().fields().iter()) + .map(|field| field.name().clone()) + .collect::>(); // step 1: align schema using project, fill non-exist columns with null let left_proj_exprs = all_columns.iter().map(|col| { - if tags_not_in_left.contains(col) { + if col == left_field_col && left_field.1 != target_field_type { + DfExpr::Cast(Cast { + expr: Box::new(DfExpr::Column(Column::new( + left_field.0.clone(), + left_field_col, + ))), + data_type: target_field_type.clone(), + }) + .alias(left_field_col.clone()) + } else if tags_not_in_left.contains(col) { DfExpr::Literal(ScalarValue::Utf8(None), None).alias(col.clone()) } else { DfExpr::Column(Column::new(None::, col)) @@ -4461,25 +4545,22 @@ impl PromPlanner { .alias(left_time_index_column.clone()); // The field column in right side may not have qualifier (it may be removed by join operation), // so we need to find it from the schema. - let right_qualifier_for_field = right - .schema() - .iter() - .find(|(_, f)| f.name() == right_field_col) - .map(|(q, _)| q) - .with_context(|| ColumnNotFoundSnafu { - col: right_field_col.clone(), - })? - .cloned(); - // `skip（1)` to skip the time index column let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| { // expr - if col == left_field_col && left_field_col != right_field_col { - // qualify field in right side if necessary to handle different field name - DfExpr::Column(Column::new( - right_qualifier_for_field.clone(), - right_field_col, - )) + if col == left_field_col { + let expr = DfExpr::Column(Column::new(right_field.0.clone(), right_field_col)); + if right_field.1 != target_field_type { + DfExpr::Cast(Cast { + expr: Box::new(expr), + data_type: target_field_type.clone(), + }) + .alias(left_field_col.clone()) + } else if left_field_col != right_field_col { + expr.alias(left_field_col.clone()) + } else { + expr + } } else if tags_not_in_right.contains(col) { DfExpr::Literal(ScalarValue::Utf8(None), None).alias(col.clone()) } else { @@ -4523,24 +4604,168 @@ impl PromPlanner { }; // sort to ensure the generated plan is not volatile match_columns.sort_unstable(); - // step 3: build `UnionDistinctOn` plan - let schema = left_projected.schema().clone(); - let union_distinct_on = UnionDistinctOn::new( - left_projected, - right_projected, - match_columns, - left_time_index_column.clone(), - schema, + match_columns.dedup(); + occupied_column_names.extend( + left_projected + .schema() + .fields() + .iter() + .chain(right_projected.schema().fields().iter()) + .map(|field| field.name().clone()), ); - let result = LogicalPlan::Extension(Extension { + + let visible_schema = left_projected.schema().clone(); + let visible_left_exprs = left_projected + .schema() + .iter() + .map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + }) + .collect::>(); + let visible_right_exprs = right_projected + .schema() + .iter() + .map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + }) + .collect::>(); + let mut left_match_exprs = Vec::with_capacity(match_columns.len()); + let mut right_match_exprs = Vec::with_capacity(match_columns.len()); + let mut next_internal_column = 0; + + for label in &match_columns { + let left_field = if left_tag_cols_set.contains(label) { + Some( + left_projected + .schema() + .iter() + .find(|(_, field)| field.name() == label) + .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, + ) + } else { + None + }; + let right_field = if right_tag_cols_set.contains(label) { + Some( + right_projected + .schema() + .iter() + .find(|(_, field)| field.name() == label) + .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, + ) + } else { + None + }; + let data_type = match (left_field.as_ref(), right_field.as_ref()) { + (Some((_, left_type)), Some((_, right_type))) if left_type == right_type => { + left_type.clone() + } + (Some((_, left_type)), Some((_, right_type))) => { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR match label {label} has incompatible types: {left_type:?} and {right_type:?}" + ), + } + .fail(); + } + (Some((_, data_type)), None) | (None, Some((_, data_type))) => data_type.clone(), + (None, None) => ArrowDataType::Utf8, + }; + let empty = match data_type { + ArrowDataType::Utf8 => ScalarValue::Utf8(Some(String::new())), + ArrowDataType::LargeUtf8 => ScalarValue::LargeUtf8(Some(String::new())), + _ => { + return UnexpectedPlanExprSnafu { + desc: format!("OR match label {label} must be a string"), + } + .fail(); + } + }; + let internal_name = loop { + let name = format!("__promql_or_match_{next_internal_column}"); + next_internal_column += 1; + if occupied_column_names.insert(name.clone()) { + break name; + } + }; + let normalize = |field: Option<(Option, ArrowDataType)>| { + let expr = if let Some((qualifier, _)) = field { + DfExpr::ScalarFunction(ScalarFunction { + func: coalesce(), + args: vec![ + DfExpr::Column(Column::new(qualifier, label.clone())), + DfExpr::Literal(empty.clone(), None), + ], + }) + } else { + DfExpr::Literal(empty.clone(), None) + }; + expr.alias(internal_name.clone()) + }; + left_match_exprs.push(normalize(left_field)); + right_match_exprs.push(normalize(right_field)); + } + + let left_augmented = LogicalPlanBuilder::from(left_projected) + .project(visible_left_exprs.into_iter().chain(left_match_exprs)) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + let right_augmented = LogicalPlanBuilder::from(right_projected) + .project(visible_right_exprs.into_iter().chain(right_match_exprs)) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + // step 3: build `UnionDistinctOn` with normalized internal match keys. + let visible_field_count = visible_schema.fields().len(); + let compare_key_indices = + (visible_field_count..visible_field_count + match_columns.len()).collect::>(); + let (time_qualifier, _) = visible_schema + .iter() + .find(|(_, field)| field.name() == &left_time_index_column) + .with_context(|| TimeIndexNotFoundSnafu { + table: left_qualifier_string.clone(), + })?; + let ts_col_idx = left_augmented + .schema() + .iter() + .position(|(qualifier, field)| { + qualifier == time_qualifier && field.name() == &left_time_index_column + }) + .with_context(|| TimeIndexNotFoundSnafu { + table: left_qualifier_string.clone(), + })?; + let union_distinct_on = UnionDistinctOn::try_new( + left_augmented, + right_augmented, + compare_key_indices, + ts_col_idx, + ) + .context(DataFusionPlanningSnafu)?; + let augmented_result = LogicalPlan::Extension(Extension { node: Arc::new(union_distinct_on), }); + let result = LogicalPlanBuilder::from(augmented_result) + .project(visible_schema.iter().map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + })) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; // step 4: update context - self.ctx.time_index_column = Some(left_time_index_column); - self.ctx.tag_columns = all_tags.into_iter().collect(); - self.ctx.field_columns = vec![left_field_col.clone()]; - self.ctx.use_tsid = left_has_tsid && right_has_tsid; + let output_field_col = left_field_col.clone(); + let mut output_context = left_context; + let mut visible_tags = all_tags.into_iter().collect::>(); + visible_tags.sort_unstable(); + output_context.time_index_column = Some(left_time_index_column); + output_context.tag_columns = visible_tags; + output_context.field_columns = vec![output_field_col]; + output_context.use_tsid = left_has_tsid && right_has_tsid; + self.ctx = output_context; Ok(result) } @@ -4747,15 +4972,23 @@ mod test { use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_query::prelude::greptime_timestamp; use common_query::test_util::DummyDecoder; - use datafusion::arrow::datatypes::Schema as ArrowSchema; + use datafusion::arrow::array::{ + Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, + }; + use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; + use datafusion::datasource::{MemTable, provider_as_source}; + use datafusion::execution::context::SessionContext; use datafusion::logical_expr::Extension; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use promql_parser::label::Labels; use promql_parser::parser; use session::context::QueryContext; + use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use table::metadata::{TableInfoBuilder, TableMetaBuilder}; use table::test_util::EmptyTable; @@ -4763,6 +4996,7 @@ mod test { use crate::QueryEngineContext; use crate::options::QueryOptions; use crate::parser::QueryLanguageParser; + use crate::query_engine::DefaultSerializer; fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> { if let LogicalPlan::Extension(Extension { node }) = plan @@ -4858,6 +5092,262 @@ mod test { } } + enum DirectOrValue { + Float64(f64), + Int64(i64), + Utf8(&'static str), + } + + impl DirectOrValue { + fn data_type(&self) -> ArrowDataType { + match self { + Self::Float64(_) => ArrowDataType::Float64, + Self::Int64(_) => ArrowDataType::Int64, + Self::Utf8(_) => ArrowDataType::Utf8, + } + } + fn array(&self) -> Arc { + match self { + Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])), + Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])), + Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])), + } + } + } + + struct DirectOrSource { + name: &'static str, + empty: bool, + timestamp: i64, + tags: Vec<(&'static str, Option<&'static str>)>, + value: DirectOrValue, + } + + fn source( + name: &'static str, + empty: bool, + timestamp: i64, + tags: Vec<(&'static str, Option<&'static str>)>, + value: DirectOrValue, + ) -> DirectOrSource { + DirectOrSource { + name, + empty, + timestamp, + tags, + value, + } + } + + fn tagged_source( + name: &'static str, + empty: bool, + tag: (&'static str, Option<&'static str>), + value: DirectOrValue, + ) -> DirectOrSource { + source(name, empty, 1, vec![("job", Some("job")), tag], value) + } + + fn job_source(name: &'static str, value: DirectOrValue) -> DirectOrSource { + source(name, true, 1, vec![("job", Some("job"))], value) + } + + fn table(source: &DirectOrSource) -> Arc { + let mut fields = vec![Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + )]; + fields.extend( + source + .tags + .iter() + .map(|(name, _)| Field::new(*name, ArrowDataType::Utf8, true)), + ); + fields.push(Field::new("v", source.value.data_type(), true)); + let schema = Arc::new(ArrowSchema::new(fields)); + let partitions = if source.empty { + vec![vec![]] + } else { + let mut columns: Vec> = + vec![Arc::new(TimestampMillisecondArray::from(vec![ + source.timestamp, + ]))]; + columns.extend( + source + .tags + .iter() + .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as Arc), + ); + columns.push(source.value.array()); + vec![vec![RecordBatch::try_new(schema.clone(), columns).unwrap()]] + }; + Arc::new(MemTable::try_new(schema, partitions).unwrap()) + } + + fn scan(source: &DirectOrSource) -> LogicalPlan { + LogicalPlanBuilder::scan(source.name, provider_as_source(table(source)), None) + .unwrap() + .build() + .unwrap() + } + + fn direct_or_context(qualifier: &str, tags: &[&str], field: &str) -> PromPlannerContext { + PromPlannerContext { + table_name: Some(qualifier.to_string()), + time_index_column: Some("ts".to_string()), + field_columns: vec![field.to_string()], + tag_columns: tags.iter().map(|tag| (*tag).to_string()).collect(), + ..Default::default() + } + } + + fn or_modifier(expr: &str) -> Option { + let PromExpr::Binary(expr) = parser::parse(expr).unwrap() else { + unreachable!() + }; + expr.modifier + } + + async fn plan_direct_or( + left: LogicalPlan, + right: LogicalPlan, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + modifier: &Option, + ) -> LogicalPlan { + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + }; + planner + .or_operator( + left, + right, + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + modifier, + ) + .unwrap() + } + + async fn execute( + plan: LogicalPlan, + state: &QueryEngineState, + ) -> (LogicalPlan, Vec) { + let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); + let optimized = state.optimize_by_extension_rules(plan, &context).unwrap(); + let physical = state + .session_state() + .create_physical_plan(&optimized) + .await + .unwrap(); + let batches = + datafusion::physical_plan::collect(physical, state.session_state().task_ctx()) + .await + .unwrap(); + (optimized, batches) + } + + async fn run( + left: &DirectOrSource, + right: &DirectOrSource, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + modifier: &Option, + ) -> (LogicalPlan, Vec) { + let plan = plan_direct_or( + scan(left), + scan(right), + left_context, + right_context, + modifier, + ) + .await; + execute(plan, &build_query_engine_state()).await + } + + fn assert_no_internal_or_keys(schema: &DFSchema) { + assert!( + schema + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")), + "{schema:?}" + ); + } + + fn values(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + batch + .column_by_name(column) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .collect() + } + + fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option)> { + let mut rows = batches + .iter() + .flat_map(|batch| { + let values = batch + .column_by_name("v") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let labels = batch + .column_by_name("k") + .map(|column| column.as_any().downcast_ref::().unwrap()); + (0..batch.num_rows()).map(move |i| { + ( + values.value(i), + labels.and_then(|labels| { + (!labels.is_null(i)).then(|| labels.value(i).to_string()) + }), + ) + }) + }) + .collect::>(); + rows.sort_by(|left, right| left.0.total_cmp(&right.0)); + rows + } + + fn matrix_source( + name: &'static str, + k: Option>, + timestamp: i64, + value: f64, + ) -> DirectOrSource { + let mut tags = vec![("job", Some("job"))]; + if let Some(k) = k { + tags.push(("k", k)); + } + source(name, false, timestamp, tags, DirectOrValue::Float64(value)) + } + + fn matrix_context(name: &str, k: Option>) -> PromPlannerContext { + direct_or_context( + name, + if k.is_some() { &["job", "k"] } else { &["job"] }, + "v", + ) + } + async fn build_test_table_provider( table_name_tuples: &[(String, String)], num_tag: usize, @@ -5542,6 +6032,33 @@ mod test { ); } + #[tokio::test] + async fn reject_binary_fill_modifiers() { + let state = build_query_engine_state(); + + for query in [ + "some_metric + fill(0) some_alt_metric", + "some_metric + fill_left(0) some_alt_metric", + "some_metric + fill_right(0) some_alt_metric", + "(some_metric + fill(0) some_alt_metric) + some_metric", + ] { + let eval_stmt = build_eval_stmt(query); + let table_provider = build_test_table_provider(&[], 0, 0).await; + let err = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &state) + .await + .unwrap_err(); + + assert!( + matches!( + &err, + crate::promql::error::Error::UnsupportedExpr { name, .. } + if name == "PromQL fill modifiers" + ), + "{err}" + ); + } + } + #[tokio::test] async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() { for query in [ @@ -8150,47 +8667,33 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom #[tokio::test] async fn test_or_not_exists_table_label() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = r#"sum by (job, tag0, tag2) (metric_exists) or sum by (job, tag0, tag2) (metric_not_exists)"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "metric_exists".to_string())], + let state = build_query_engine_state(); + let provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string())], &["job"], ) .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = r#"UnionDistinctOn: on col=[["job"]], ts_col=[greptime_timestamp] [greptime_timestamp:Timestamp(ms), job:Utf8, sum(metric_exists.greptime_value):Float64;N] - SubqueryAlias: metric_exists [greptime_timestamp:Timestamp(ms), job:Utf8, sum(metric_exists.greptime_value):Float64;N] - Projection: metric_exists.greptime_timestamp, metric_exists.job, sum(metric_exists.greptime_value) [greptime_timestamp:Timestamp(ms), job:Utf8, sum(metric_exists.greptime_value):Float64;N] - Sort: metric_exists.job ASC NULLS LAST, metric_exists.greptime_timestamp ASC NULLS LAST [job:Utf8, greptime_timestamp:Timestamp(ms), sum(metric_exists.greptime_value):Float64;N] - Aggregate: groupBy=[[metric_exists.job, metric_exists.greptime_timestamp]], aggr=[[sum(metric_exists.greptime_value)]] [job:Utf8, greptime_timestamp:Timestamp(ms), sum(metric_exists.greptime_value):Float64;N] - PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - PromSeriesDivide: tags=["job"] [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - Sort: metric_exists.job ASC NULLS FIRST, metric_exists.greptime_timestamp ASC NULLS FIRST [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - Filter: metric_exists.greptime_timestamp >= TimestampMillisecond(-999, None) AND metric_exists.greptime_timestamp <= TimestampMillisecond(100000000, None) [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - TableScan: metric_exists [job:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - SubqueryAlias: [greptime_timestamp:Timestamp(ms), job:Utf8;N, sum(.value):Float64;N] - Projection: .time AS greptime_timestamp, Utf8(NULL) AS job, sum(.value) [greptime_timestamp:Timestamp(ms), job:Utf8;N, sum(.value):Float64;N] - Sort: .time ASC NULLS LAST [time:Timestamp(ms), sum(.value):Float64;N] - Aggregate: groupBy=[[.time]], aggr=[[sum(.value)]] [time:Timestamp(ms), sum(.value):Float64;N] - EmptyMetric: range=[0..-1], interval=[5000] [time:Timestamp(ms), value:Float64;N] - TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#; - - assert_eq!(plan.display_indent_schema().to_string(), expected); + let raw = PromPlanner::stmt_to_plan( + provider, + &build_eval_stmt(r#"missing_metric or on(absent_label) normal_metric"#), + &state, + ) + .await + .unwrap(); + assert!( + raw.display_indent_schema() + .to_string() + .contains("__promql_or_match_0@") + ); + let (optimized, batches) = execute(raw, &state).await; + assert_no_internal_or_keys(optimized.schema()); + assert!(batches.iter().all(|batch| { + batch + .schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")) + })); } #[tokio::test] @@ -8242,4 +8745,250 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom _ => panic!("Expected EmptyRelation, but got: {:?}", plan), } } + + #[tokio::test] + async fn test_direct_or_normalizes_missing_match_labels() { + type Case<'a> = ( + Option>, + Option>, + i64, + i64, + &'a [(f64, Option<&'a str>)], + ); + + let modifier = or_modifier("lhs or on(k) rhs"); + #[rustfmt::skip] + let cases: &[Case<'_>] = &[ + (None, None, 1, 1, &[(1.0, None)]), + (None, Some(Some("")), 1, 1, &[(1.0, None)]), + (Some(Some("")), None, 1, 1, &[(1.0, Some(""))]), + (None, Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), + (Some(Some("l")), None, 1, 1, &[(1.0, Some("l")), (2.0, None)]), + (Some(None), Some(Some("")), 1, 1, &[(1.0, None)]), + (Some(None), Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), + (Some(Some("same")), Some(Some("same")), 1, 2, &[(1.0, Some("same")), (2.0, Some("same"))]), + ]; + for &(left, right, left_ts, right_ts, expected) in cases { + let (optimized, batches) = run( + &matrix_source("lhs", left, left_ts, 1.0), + &matrix_source("rhs", right, right_ts, 2.0), + matrix_context("lhs", left), + matrix_context("rhs", right), + &modifier, + ) + .await; + assert_no_internal_or_keys(optimized.schema()); + assert_eq!( + rows(&batches), + expected + .iter() + .map(|(value, label)| (*value, label.map(str::to_string))) + .collect::>() + ); + } + } + + #[tokio::test] + async fn test_direct_or_match_modifiers() { + for (modifier, left, right, expected) in [ + (None, "left", "right", 2), + (or_modifier("lhs or on(k) rhs"), "same", "same", 1), + (or_modifier("lhs or on() rhs"), "left", "right", 1), + (or_modifier("lhs or ignoring(k) rhs"), "left", "right", 1), + ] { + let (_, batches) = run( + &matrix_source("lhs", Some(Some(left)), 1, 1.0), + &matrix_source("rhs", Some(Some(right)), 1, 2.0), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &modifier, + ) + .await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected + ); + } + } + + #[tokio::test] + async fn test_direct_or_nested_projection_uses_left_context() { + let left = matrix_source("lhs", Some(Some("k")), 1, 1.0); + let right = matrix_source("rhs", Some(Some("k")), 1, 2.0); + let raw = plan_direct_or( + scan(&left), + scan(&right), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + assert!(raw.schema().iter().any(|(qualifier, field)| { + qualifier.as_ref().is_some_and(|q| q.to_string() == "lhs") && field.name() == "v" + })); + let nested = LogicalPlanBuilder::from(raw) + .project(vec![ + DfExpr::BinaryExpr(BinaryExpr { + left: Box::new(DfExpr::Column(Column::new( + Some(TableReference::bare("lhs")), + "v", + ))), + op: Operator::Plus, + right: Box::new(lit(1.0)), + }) + .alias("v_plus"), + ]) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(nested, &build_query_engine_state()).await; + assert_eq!(values(&batches, "v_plus"), vec![2.0]); + } + + #[tokio::test] + async fn test_direct_or_skips_user_internal_key_name() { + const USER_TAG: &str = "__promql_or_match_0"; + let left = tagged_source( + "lhs", + false, + (USER_TAG, Some("left")), + DirectOrValue::Float64(1.0), + ); + let right = tagged_source( + "rhs", + false, + (USER_TAG, Some("right")), + DirectOrValue::Float64(2.0), + ); + let raw = plan_direct_or( + scan(&left), + scan(&right), + direct_or_context("lhs", &["job", USER_TAG], "v"), + direct_or_context("rhs", &["job", USER_TAG], "v"), + &or_modifier("lhs or on(missing_label) rhs"), + ) + .await; + assert!( + raw.display_indent_schema() + .to_string() + .contains("__promql_or_match_1@") + ); + let (_, batches) = execute(raw, &build_query_engine_state()).await; + assert!( + batches + .iter() + .all(|batch| batch.column_by_name(USER_TAG).is_some()) + ); + } + + #[tokio::test] + async fn test_direct_or_substrait_round_trip_with_normalized_key() { + let state = build_query_engine_state(); + let ctx = SessionContext::new_with_state(state.session_state()); + let catalog = Arc::new(MemoryCatalogProvider::new()); + catalog + .register_schema("public", Arc::new(MemorySchemaProvider::new())) + .unwrap(); + ctx.register_catalog("datafusion", catalog); + let left = matrix_source("lhs", Some(Some("")), 1, 1.0); + let right = matrix_source("rhs", None, 1, 2.0); + ctx.register_table( + TableReference::full("datafusion", "public", "lhs"), + table(&left), + ) + .unwrap(); + ctx.register_table( + TableReference::full("datafusion", "public", "rhs"), + table(&right), + ) + .unwrap(); + let raw = plan_direct_or( + ctx.table("datafusion.public.lhs") + .await + .unwrap() + .into_unoptimized_plan(), + ctx.table("datafusion.public.rhs") + .await + .unwrap() + .into_unoptimized_plan(), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + let decoded = DFLogicalSubstraitConvertor + .decode( + DFLogicalSubstraitConvertor + .encode(&raw, DefaultSerializer) + .unwrap(), + ctx.state(), + ) + .await + .unwrap(); + let (optimized, batches) = execute(decoded, &state).await; + assert_no_internal_or_keys(optimized.schema()); + assert!(batches.iter().all(|batch| { + batch + .schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")) + })); + assert_eq!(values(&batches, "v"), vec![1.0]); + } + + #[tokio::test] + async fn test_direct_or_numeric_value_types() { + let left = tagged_source("lhs", true, ("k", Some("lhs")), DirectOrValue::Int64(0)); + let right = tagged_source( + "rhs", + false, + ("k", Some("rhs")), + DirectOrValue::Float64(0.5), + ); + let (optimized, batches) = run( + &left, + &right, + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + assert_eq!( + optimized + .schema() + .field_with_name(None, "v") + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!(values(&batches, "v"), vec![0.5]); + let provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider: provider, + ctx: PromPlannerContext::default(), + }; + let left_context = direct_or_context("lhs", &["job"], "v"); + let right_context = direct_or_context("rhs", &["job"], "v"); + let error = planner + .or_operator( + scan(&job_source("lhs", DirectOrValue::Utf8("x"))), + scan(&job_source("rhs", DirectOrValue::Float64(1.0))), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on() rhs"), + ) + .unwrap_err(); + assert!( + error + .to_string() + .contains("OR value fields have incompatible types") + ); + } } diff --git a/src/servers/src/http/logs.rs b/src/servers/src/http/logs.rs index b9f23b68089..ca7c061da05 100644 --- a/src/servers/src/http/logs.rs +++ b/src/servers/src/http/logs.rs @@ -19,24 +19,31 @@ use axum::extract::State; use axum::response::{IntoResponse, Response}; use axum::{Extension, Json}; use common_telemetry::tracing; -use log_query::LogQuery; +use log_query::{Limit, LogQuery}; use session::context::{Channel, QueryContext}; use crate::http::result::greptime_result_v1::GreptimedbV1Response; use crate::query_handler::LogQueryHandlerRef; +const DEFAULT_FETCH: usize = 1000; + +fn apply_default_fetch(limit: &mut Limit) { + limit.fetch.get_or_insert(DEFAULT_FETCH); +} + #[axum_macros::debug_handler] #[tracing::instrument(skip_all, fields(protocol = "http", request_type = "logs"))] pub async fn logs( State(handler): State, Extension(mut query_ctx): Extension, - Json(params): Json, + Json(mut params): Json, ) -> Response { let exec_start = Instant::now(); let db = query_ctx.get_db_string(); query_ctx.set_channel(Channel::Log); let query_ctx = Arc::new(query_ctx); + apply_default_fetch(&mut params.limit); let _timer = crate::metrics::METRIC_HTTP_LOGS_ELAPSED .with_label_values(&[db.as_str()]) @@ -48,3 +55,34 @@ pub async fn logs( resp.with_execution_time(exec_start.elapsed().as_millis() as u64) .into_response() } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_apply_default_fetch() { + let mut limit = Limit { + skip: None, + fetch: None, + }; + apply_default_fetch(&mut limit); + assert_eq!(limit.fetch, Some(DEFAULT_FETCH)); + + let mut limit = Limit { + skip: Some(10), + fetch: None, + }; + apply_default_fetch(&mut limit); + assert_eq!(limit.skip, Some(10)); + assert_eq!(limit.fetch, Some(DEFAULT_FETCH)); + + let mut limit = Limit { + skip: Some(10), + fetch: Some(42), + }; + apply_default_fetch(&mut limit); + assert_eq!(limit.skip, Some(10)); + assert_eq!(limit.fetch, Some(42)); + } +} diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index 6f2b1156862..f0d16894363 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -24,6 +24,7 @@ use axum::http::HeaderValue; use axum::response::{IntoResponse, Response}; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; +use common_query::prometheus::is_prometheus_stale_nan; use common_query::{Output, OutputData}; use common_recordbatch::RecordBatches; use datatypes::prelude::ConcreteDataType; @@ -268,8 +269,8 @@ impl PrometheusJsonResponse { // retrieve value if field_column.is_valid(row_index) { let v = field_column.value(row_index); - // ignore all NaN values to reduce the amount of data to be sent. - if v.is_nan() { + // Ignore Prometheus stale markers. + if is_prometheus_stale_nan(v) { continue; } @@ -350,3 +351,59 @@ impl PrometheusJsonResponse { Ok(data) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_recordbatch::{RecordBatch, RecordBatches}; + use datatypes::data_type::ConcreteDataType; + use datatypes::schema::{ColumnSchema, Schema}; + use datatypes::vectors::{Float64Vector, TimestampMillisecondVector}; + + use super::*; + + #[test] + fn matrix_response_preserves_ordinary_nan_and_filters_stale_markers() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ])); + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![ + 1_000, 2_000, 3_000, 4_000, + ])) as _, + Arc::new(Float64Vector::from(vec![ + Some(1.0), + Some(f64::from_bits(0x7ff8_0000_0000_0000)), + Some(f64::from_bits(0x7ff0_0000_0000_0002)), + None, + ])) as _, + ], + ) + .unwrap(); + let batches = RecordBatches::try_new(schema, vec![batch]).unwrap(); + + let response = + PrometheusJsonResponse::record_batches_to_data(batches, None, ValueType::Matrix) + .unwrap(); + let PrometheusResponse::PromData(data) = response else { + panic!("expected Prometheus data response"); + }; + let PromQueryResult::Matrix(series) = data.result else { + panic!("expected matrix result"); + }; + + assert_eq!(series.len(), 1); + assert_eq!( + series[0].values, + vec![(1.0, "1".to_string()), (2.0, "NaN".to_string())] + ); + } +} diff --git a/src/servers/src/otlp/coerce.rs b/src/servers/src/otlp/coerce.rs index febec0fda9b..2d87744c5c2 100644 --- a/src/servers/src/otlp/coerce.rs +++ b/src/servers/src/otlp/coerce.rs @@ -108,38 +108,35 @@ pub fn trace_value_datatype(value: &ValueData) -> Option { pub fn resolve_new_trace_column_type( observed_types: impl IntoIterator, ) -> Result, TraceCoerceError> { - let mut observed_types = observed_types.into_iter().collect::>(); - observed_types.dedup(); - - match observed_types.as_slice() { - [] => Ok(None), - [datatype] => Ok(Some(*datatype)), - [_, _] - if observed_types.contains(&ColumnDataType::String) - && observed_types.contains(&ColumnDataType::Boolean) => - { - Ok(Some(ColumnDataType::Boolean)) + let mut observed = Vec::new(); + for datatype in observed_types { + if !observed.contains(&datatype) { + observed.push(datatype); } - [_, _] - if observed_types.contains(&ColumnDataType::String) - && observed_types.contains(&ColumnDataType::Int64) => - { - Ok(Some(ColumnDataType::Int64)) - } - [_, _] - if observed_types.contains(&ColumnDataType::String) - && observed_types.contains(&ColumnDataType::Float64) => - { - Ok(Some(ColumnDataType::Float64)) - } - [_, _] - if observed_types.contains(&ColumnDataType::Int64) - && observed_types.contains(&ColumnDataType::Float64) => - { - Ok(Some(ColumnDataType::Float64)) - } - _ => Err(TraceCoerceError::Unsupported), } + + if observed.is_empty() { + return Ok(None); + } + if observed.len() == 1 { + return Ok(observed.first().copied()); + } + + [ + ColumnDataType::Boolean, + ColumnDataType::Int64, + ColumnDataType::Float64, + ColumnDataType::String, + ] + .into_iter() + .find(|target| { + observed.contains(target) + && observed + .iter() + .all(|source| source == target || is_supported_trace_coercion(*source, *target)) + }) + .map(Some) + .ok_or(TraceCoerceError::Unsupported) } #[cfg(test)] @@ -337,7 +334,15 @@ mod tests { ColumnDataType::Int64, ColumnDataType::Float64, ]), - Err(TraceCoerceError::Unsupported) + Ok(Some(ColumnDataType::Float64)) + ); + assert_eq!( + resolve_new_trace_column_type([ + ColumnDataType::Float64, + ColumnDataType::String, + ColumnDataType::Int64, + ]), + Ok(Some(ColumnDataType::Float64)) ); } } diff --git a/src/servers/src/otlp/trace/v1.rs b/src/servers/src/otlp/trace/v1.rs index e04d344e54e..eedc598ffba 100644 --- a/src/servers/src/otlp/trace/v1.rs +++ b/src/servers/src/otlp/trace/v1.rs @@ -12,10 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; +use api::v1::column_data_type_extension::TypeExt; use api::v1::value::ValueData; -use api::v1::{ColumnDataType, RowInsertRequests, Value}; +use api::v1::{ + ColumnDataType, ColumnDataTypeExtension, ColumnSchema, JsonTypeExtension, RowInsertRequests, + Value, +}; use common_catalog::consts::{trace_operations_table_name, trace_services_table_name}; use common_grpc::precision::Precision; use opentelemetry_proto::tonic::common::v1::any_value::Value as OtlpValue; @@ -40,6 +44,155 @@ const APPROXIMATE_COLUMN_COUNT: usize = 30; // Use a timestamp(2100-01-01 00:00:00) as large as possible. const MAX_TIMESTAMP: i64 = 4102444800000000000; +/// Distinguishes raw bytes from JSONB values that both use a binary protobuf value. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TraceBinaryType { + /// Raw OTLP bytes with no datatype extension. + Binary, + /// An OTLP array or key-value list encoded with the JSONB extension. + Json, +} + +impl TraceBinaryType { + /// Applies this logical binary type to a column schema. + pub fn apply_to_schema(self, schema: &mut ColumnSchema) { + schema.datatype = ColumnDataType::Binary as i32; + schema.datatype_extension = match self { + Self::Binary => None, + Self::Json => Some(ColumnDataTypeExtension { + type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), + }), + }; + } +} + +/// Per-column observations collected while building one trace chunk. +#[derive(Default)] +struct TraceBatchColumnSchema { + value_types: Vec, + present_rows: Vec, + binary_types: Vec<(usize, TraceBinaryType)>, +} + +impl TraceBatchColumnSchema { + /// Records a row once while preserving encounter order. + fn observe_row(&mut self, row_index: usize) { + if self.present_rows.last() != Some(&row_index) { + self.present_rows.push(row_index); + } + } + + /// Returns whether raw binary and JSONB values share this column. + fn has_incompatible_logical_types(&self) -> bool { + let Some((_, first_type)) = self.binary_types.first() else { + return false; + }; + self.binary_types + .iter() + .any(|(_, binary_type)| binary_type != first_type) + } +} + +/// Sparse column metadata retained for span-level fallback. +pub struct TraceRetryColumn { + /// Row indexes that contain this dynamic column. + pub present_rows: Vec, + /// Row-specific binary kinds needed to restore the logical schema. + pub binary_types: Vec<(usize, TraceBinaryType)>, +} + +/// Retry metadata keyed by dynamic trace column name. +pub type TraceRetryColumns = HashMap; + +/// Schema observations for dynamic columns in one converted trace chunk. +#[derive(Default)] +pub struct TraceBatchSchema { + columns: HashMap, +} + +impl TraceBatchSchema { + /// Records a scalar value type and the row where the column is present. + fn observe_value_type(&mut self, name: &str, row_index: usize, value_type: ColumnDataType) { + let column = self.columns.entry(name.to_string()).or_default(); + column.observe_row(row_index); + if !column.value_types.contains(&value_type) { + column.value_types.push(value_type); + } + } + + /// Records the logical kind of a binary protobuf value for one row. + fn observe_binary_type(&mut self, name: &str, row_index: usize, binary_type: TraceBinaryType) { + let column = self.columns.entry(name.to_string()).or_default(); + column.observe_row(row_index); + if !column.value_types.contains(&ColumnDataType::Binary) { + column.value_types.push(ColumnDataType::Binary); + } + if let Some((last_row_index, last_type)) = column.binary_types.last_mut() + && *last_row_index == row_index + { + *last_type = binary_type; + } else { + column.binary_types.push((row_index, binary_type)); + } + } + + /// Records a sparse column occurrence that has no dynamic value type. + fn observe_present_column(&mut self, name: &str, row_index: usize) { + self.columns + .entry(name.to_string()) + .or_default() + .observe_row(row_index); + } + + /// Returns the distinct value types in first-seen order for a column. + pub fn value_types(&self, name: &str) -> Option<&[ColumnDataType]> { + self.columns + .get(name) + .map(|column| column.value_types.as_slice()) + } + + /// Returns whether any column mixes raw binary and JSONB values. + pub fn has_incompatible_logical_types(&self) -> bool { + self.columns + .values() + .any(TraceBatchColumnSchema::has_incompatible_logical_types) + } + + /// Returns whether the named column mixes raw binary and JSONB values. + pub fn has_incompatible_logical_types_for(&self, name: &str) -> bool { + self.columns + .get(name) + .is_some_and(TraceBatchColumnSchema::has_incompatible_logical_types) + } + + /// Converts batch observations into sparse metadata for per-span retries. + pub fn into_retry_columns(self) -> TraceRetryColumns { + self.columns + .into_iter() + .filter(|(_, column)| !column.present_rows.is_empty()) + .map(|(name, column)| { + let binary_types = if column.has_incompatible_logical_types() + || column + .value_types + .iter() + .any(|datatype| *datatype != ColumnDataType::Binary) + { + column.binary_types + } else { + Vec::new() + }; + ( + name, + TraceRetryColumn { + present_rows: column.present_rows, + binary_types, + }, + ) + }) + .collect() + } +} + /// Converts trace spans into row insert requests for the main v1 trace table. /// /// Auxiliary service and operation table writes are built separately so the @@ -52,20 +205,58 @@ pub fn v1_to_grpc_main_insert_requests( _query_ctx: &QueryContextRef, _pipeline_handler: PipelineHandlerRef, ) -> Result<(RowInsertRequests, usize)> { + let requests = v1_to_grpc_main_insert_requests_from_iter(spans.iter().cloned(), table_name)?; + Ok((requests, spans.len())) +} + +/// Converts owned spans into unpadded main-table rows and schema observations. +pub fn v1_to_main_table_data_with_schema( + spans: Vec, +) -> Result<(TableData, TraceBatchSchema)> { + build_trace_table_data_with_schema(spans.into_iter()) +} + +/// Builds the main-table request without collecting batch schema observations. +fn v1_to_grpc_main_insert_requests_from_iter( + spans: impl ExactSizeIterator, + table_name: &str, +) -> Result { let mut multi_table_writer = MultiTableData::default(); - let trace_writer = build_trace_table_data(spans)?; + let trace_writer = build_trace_table_data_from_iter(spans, None)?; multi_table_writer.add_table_data(table_name, trace_writer); - Ok(multi_table_writer.into_row_insert_requests()) + Ok(multi_table_writer.into_row_insert_requests().0) } /// Builds the row-oriented payload for the main v1 trace table. pub fn build_trace_table_data(spans: &[TraceSpan]) -> Result { - let mut trace_writer = TableData::new(APPROXIMATE_COLUMN_COUNT, spans.len()); - for span in spans.iter().cloned() { - write_span_to_row(&mut trace_writer, span)?; - } + build_trace_table_data_from_iter(spans.iter().cloned(), None) +} +/// Builds trace rows while collecting dynamic column observations. +fn build_trace_table_data_with_schema( + spans: impl ExactSizeIterator, +) -> Result<(TableData, TraceBatchSchema)> { + let mut batch_schema = TraceBatchSchema::default(); + let trace_writer = build_trace_table_data_from_iter(spans, Some(&mut batch_schema))?; + Ok((trace_writer, batch_schema)) +} + +/// Shared row builder with optional batch schema observation. +fn build_trace_table_data_from_iter( + spans: impl ExactSizeIterator, + mut batch_schema: Option<&mut TraceBatchSchema>, +) -> Result { + let mut trace_writer = TableData::new(APPROXIMATE_COLUMN_COUNT, spans.len()); + for span in spans { + let row_index = trace_writer.num_rows(); + write_span_to_row_inner( + &mut trace_writer, + span, + row_index, + batch_schema.as_deref_mut(), + )?; + } Ok(trace_writer) } @@ -91,6 +282,17 @@ pub fn build_aux_table_requests( } pub fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> { + let row_index = writer.num_rows(); + write_span_to_row_inner(writer, span, row_index, None) +} + +/// Writes one span and optionally records its dynamic columns for reconciliation. +fn write_span_to_row_inner( + writer: &mut TableData, + span: TraceSpan, + row_index: usize, + mut batch_schema: Option<&mut TraceBatchSchema>, +) -> Result<()> { let mut row = writer.alloc_one_row(); // write ts @@ -132,6 +334,9 @@ pub fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> row_writer::write_fields(writer, fields.into_iter(), &mut row)?; if let Some(service_name) = span.service_name { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_present_column(SERVICE_NAME_COLUMN, row_index); + } row_writer::write_tags( writer, std::iter::once((SERVICE_NAME_COLUMN.to_string(), service_name)), @@ -139,13 +344,29 @@ pub fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> )?; } - write_attributes(writer, "span_attributes", span.span_attributes, &mut row)?; - write_attributes(writer, "scope_attributes", span.scope_attributes, &mut row)?; - write_attributes( + write_attributes_with_schema( + writer, + "span_attributes", + span.span_attributes, + &mut row, + row_index, + batch_schema.as_deref_mut(), + )?; + write_attributes_with_schema( + writer, + "scope_attributes", + span.scope_attributes, + &mut row, + row_index, + batch_schema.as_deref_mut(), + )?; + write_attributes_with_schema( writer, "resource_attributes", span.resource_attributes, &mut row, + row_index, + batch_schema, )?; row_writer::write_json( @@ -217,11 +438,25 @@ fn write_trace_operations_to_row( Ok(()) } +#[cfg(test)] pub(crate) fn write_attributes( writer: &mut TableData, prefix: &str, attributes: Attributes, row: &mut Vec, +) -> Result<()> { + let row_index = writer.num_rows(); + write_attributes_with_schema(writer, prefix, attributes, row, row_index, None) +} + +/// Writes flattened attributes without coercion and optionally records their actual types. +fn write_attributes_with_schema( + writer: &mut TableData, + prefix: &str, + attributes: Attributes, + row: &mut Vec, + row_index: usize, + mut batch_schema: Option<&mut TraceBatchSchema>, ) -> Result<()> { for attr in attributes.take().into_iter() { let key_suffix = attr.key; @@ -234,6 +469,9 @@ pub(crate) fn write_attributes( let key = format!("{}.{}", prefix, key_suffix); match attr.value.and_then(|v| v.value) { Some(OtlpValue::StringValue(v)) => { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_value_type(&key, row_index, ColumnDataType::String); + } // Keep the raw request value here. Mixed trace types are reconciled later // in the frontend once we can also see the existing table schema. writer.write_field_unchecked( @@ -244,6 +482,9 @@ pub(crate) fn write_attributes( ); } Some(OtlpValue::BoolValue(v)) => { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_value_type(&key, row_index, ColumnDataType::Boolean); + } // Do not coerce or promote types while building the request-local rows. writer.write_field_unchecked( &key, @@ -253,6 +494,9 @@ pub(crate) fn write_attributes( ); } Some(OtlpValue::IntValue(v)) => { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_value_type(&key, row_index, ColumnDataType::Int64); + } // Preserving the original value avoids order-dependent behavior inside one batch. writer.write_field_unchecked( &key, @@ -262,6 +506,9 @@ pub(crate) fn write_attributes( ); } Some(OtlpValue::DoubleValue(v)) => { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_value_type(&key, row_index, ColumnDataType::Float64); + } writer.write_field_unchecked( &key, ColumnDataType::Float64, @@ -269,28 +516,40 @@ pub(crate) fn write_attributes( row, ); } - Some(OtlpValue::ArrayValue(v)) => row_writer::write_json( - writer, - key, - any_value_to_jsonb(OtlpValue::ArrayValue(v)), - row, - )?, - Some(OtlpValue::KvlistValue(v)) => row_writer::write_json( - writer, - key, - any_value_to_jsonb(OtlpValue::KvlistValue(v)), - row, - )?, - Some(OtlpValue::BytesValue(v)) => { - row_writer::write_fields( - writer, - std::iter::once(make_column_data( - &key, - ColumnDataType::Binary, - Some(ValueData::BinaryValue(v)), + Some(OtlpValue::ArrayValue(v)) => { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_binary_type(&key, row_index, TraceBinaryType::Json); + } + writer.write_column_unchecked( + row_writer::build_json_column_schema(key), + Some(ValueData::BinaryValue( + any_value_to_jsonb(OtlpValue::ArrayValue(v)).to_vec(), )), row, - )?; + ); + } + Some(OtlpValue::KvlistValue(v)) => { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_binary_type(&key, row_index, TraceBinaryType::Json); + } + writer.write_column_unchecked( + row_writer::build_json_column_schema(key), + Some(ValueData::BinaryValue( + any_value_to_jsonb(OtlpValue::KvlistValue(v)).to_vec(), + )), + row, + ); + } + Some(OtlpValue::BytesValue(v)) => { + if let Some(batch_schema) = batch_schema.as_deref_mut() { + batch_schema.observe_binary_type(&key, row_index, TraceBinaryType::Binary); + } + writer.write_field_unchecked( + key, + ColumnDataType::Binary, + Some(ValueData::BinaryValue(v)), + row, + ); } None => {} } @@ -341,6 +600,129 @@ mod tests { } } + #[test] + fn test_batch_schema_preserves_column_order_and_attribute_types() { + let mut span1 = make_span("svc-a", "trace-a", "span-a"); + span1.span_attributes = Attributes::from(vec![make_kv("val", OtlpValue::IntValue(10))]); + let mut span2 = make_span("svc-a", "trace-a", "span-b"); + span2.span_attributes = Attributes::from(vec![make_kv("val", OtlpValue::DoubleValue(1.5))]); + + let (table_data, batch_schema) = + build_trace_table_data_with_schema(vec![span1, span2].into_iter()).unwrap(); + + assert_eq!( + batch_schema.value_types("span_attributes.val").unwrap(), + [ColumnDataType::Int64, ColumnDataType::Float64] + ); + let timestamp_index = table_data + .columns() + .iter() + .position(|column| column.column_name == TIMESTAMP_COLUMN) + .unwrap(); + let attribute_index = table_data + .columns() + .iter() + .position(|column| column.column_name == "span_attributes.val") + .unwrap(); + assert!(timestamp_index < attribute_index); + } + + #[test] + fn test_optional_batch_schema_observation_preserves_rows() { + let mut span = make_span("svc-a", "trace-a", "span-a"); + span.span_attributes = Attributes::from(vec![make_kv("val", OtlpValue::IntValue(10))]); + + let without_schema = build_trace_table_data(std::slice::from_ref(&span)).unwrap(); + let (with_schema, batch_schema) = + build_trace_table_data_with_schema(vec![span].into_iter()).unwrap(); + + assert_eq!( + without_schema.into_schema_and_rows(), + with_schema.into_schema_and_rows() + ); + assert_eq!( + batch_schema.value_types("span_attributes.val").unwrap(), + [ColumnDataType::Int64] + ); + let retry_columns = batch_schema.into_retry_columns(); + let retry_column = &retry_columns["span_attributes.val"]; + assert_eq!(retry_column.present_rows, [0]); + assert!(retry_column.binary_types.is_empty()); + } + + #[test] + fn test_batch_schema_tracks_binary_and_json_rows() { + let mut binary_span = make_span("svc-a", "trace-a", "span-a"); + binary_span.span_attributes = Attributes::from(vec![make_kv( + "val", + OtlpValue::BytesValue(vec![1_u8, 2, 3]), + )]); + let mut json_span = make_span("svc-a", "trace-a", "span-b"); + json_span.span_attributes = Attributes::from(vec![make_kv( + "val", + OtlpValue::ArrayValue(ArrayValue { + values: vec![AnyValue { + value: Some(OtlpValue::IntValue(1)), + }], + }), + )]); + + let (_, batch_schema) = + build_trace_table_data_with_schema(vec![binary_span, json_span].into_iter()).unwrap(); + + assert!(batch_schema.has_incompatible_logical_types()); + assert!(batch_schema.has_incompatible_logical_types_for("span_attributes.val")); + let retry_columns = batch_schema.into_retry_columns(); + let retry_column = &retry_columns["span_attributes.val"]; + assert_eq!(retry_column.present_rows, [0, 1]); + assert_eq!( + retry_column.binary_types, + [(0, TraceBinaryType::Binary), (1, TraceBinaryType::Json)] + ); + } + + #[test] + fn test_batch_schema_preserves_scalar_then_binary_values() { + let mut scalar_span = make_span("svc-a", "trace-a", "span-a"); + scalar_span.span_attributes = Attributes::from(vec![ + make_kv("bytes", OtlpValue::StringValue("text".to_string())), + make_kv("json", OtlpValue::StringValue("text".to_string())), + ]); + let mut binary_span = make_span("svc-a", "trace-a", "span-b"); + binary_span.span_attributes = Attributes::from(vec![ + make_kv("bytes", OtlpValue::BytesValue(vec![1_u8, 2, 3])), + make_kv( + "json", + OtlpValue::ArrayValue(ArrayValue { + values: vec![AnyValue { + value: Some(OtlpValue::IntValue(1)), + }], + }), + ), + ]); + + let (_, batch_schema) = + build_trace_table_data_with_schema(vec![scalar_span, binary_span].into_iter()).unwrap(); + + assert_eq!( + batch_schema.value_types("span_attributes.bytes").unwrap(), + [ColumnDataType::String, ColumnDataType::Binary] + ); + assert_eq!( + batch_schema.value_types("span_attributes.json").unwrap(), + [ColumnDataType::String, ColumnDataType::Binary] + ); + let retry_columns = batch_schema.into_retry_columns(); + assert_eq!( + retry_columns["span_attributes.bytes"].binary_types, + [(1, TraceBinaryType::Binary)] + ); + assert_eq!( + retry_columns["span_attributes.json"].binary_types, + [(1, TraceBinaryType::Json)] + ); + } + #[test] fn test_keep_mixed_numeric_values_until_frontend_reconciliation() { let mut writer = TableData::new(4, 2); diff --git a/src/servers/src/postgres/auth_handler.rs b/src/servers/src/postgres/auth_handler.rs index 286c39da64b..1788da89364 100644 --- a/src/servers/src/postgres/auth_handler.rs +++ b/src/servers/src/postgres/auth_handler.rs @@ -15,20 +15,28 @@ use std::fmt::Debug; use std::sync::Exclusive; -use ::auth::{Identity, Password, UserInfoRef, UserProviderRef, userinfo_by_name}; +use ::auth::{ + Identity, Password, PgAuthInfo, PgScramSha256Verifier, UserInfoRef, UserProviderRef, + userinfo_by_name, +}; use async_trait::async_trait; +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64; use common_catalog::parse_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; use common_time::Timezone; use futures::{Sink, SinkExt}; use pgwire::api::auth::StartupHandler; +use pgwire::api::auth::sasl::SCRAM_SHA_256_METHOD; use pgwire::api::{ClientInfo, PgWireConnectionState, auth}; use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; use pgwire::messages::response::ErrorResponse; -use pgwire::messages::startup::{Authentication, SecretKey}; +use pgwire::messages::startup::{Authentication, PasswordMessageFamily, SecretKey}; use pgwire::messages::{PgWireBackendMessage, PgWireFrontendMessage}; use session::Session; use snafu::IntoError; +use tokio::sync::Mutex; use crate::error::{AuthSnafu, Result}; use crate::metrics::METRIC_AUTH_FAILURE; @@ -39,14 +47,34 @@ use crate::query_handler::sql::ServerSqlQueryHandlerRef; pub(crate) struct PgLoginVerifier { user_provider: Option, + state: Mutex, } impl PgLoginVerifier { pub(crate) fn new(user_provider: Option) -> Self { - Self { user_provider } + Self { + user_provider, + state: Mutex::new(PgAuthenticationState::Initial), + } } } +enum PgAuthenticationState { + Initial, + Cleartext, + SaslInitial { + auth_info: PgAuthInfo, + }, + SaslFinal { + verifier: PgScramSha256Verifier, + user_info: Option, + channel_binding: String, + nonce: String, + client_first_bare: String, + server_first: String, + }, +} + #[allow(dead_code)] struct LoginInfo { user: Option, @@ -113,6 +141,57 @@ impl PgLoginVerifier { Ok(user_info) => Ok(Some(user_info)), } } + + async fn postgres_auth_info(&self, login: &LoginInfo) -> Result { + let user_provider = match &self.user_provider { + Some(provider) => provider, + None => return Ok(PgAuthInfo::Cleartext), + }; + + let user_name = match &login.user { + Some(name) => name, + None => return Ok(PgAuthInfo::Cleartext), + }; + + match user_provider + .postgres_auth_info(Identity::UserId(user_name, None)) + .await + { + Err(e) => { + METRIC_AUTH_FAILURE + .with_label_values(&[e.status_code().as_ref()]) + .inc(); + Err(AuthSnafu.into_error(e)) + } + Ok(auth_info) => Ok(auth_info), + } + } + + async fn authorize(&self, login: &LoginInfo, user_info: &UserInfoRef) -> Result<()> { + let user_provider = match &self.user_provider { + Some(provider) => provider, + None => return Ok(()), + }; + + let catalog = match &login.catalog { + Some(name) => name, + None => return Ok(()), + }; + let schema = match &login.schema { + Some(name) => name, + None => return Ok(()), + }; + + match user_provider.authorize(catalog, schema, user_info).await { + Err(e) => { + METRIC_AUTH_FAILURE + .with_label_values(&[e.status_code().as_ref()]) + .inc(); + Err(AuthSnafu.into_error(e)) + } + Ok(()) => Ok(()), + } + } } fn set_client_info(client: &mut C, session: &Session) @@ -190,12 +269,36 @@ impl StartupHandler for PostgresServerHandlerInner { } if self.login_verifier.user_provider.is_some() { + let login_info = LoginInfo::from_client_info(client); + let auth_info = match self.login_verifier.postgres_auth_info(&login_info).await + { + Ok(auth_info) => auth_info, + Err(_) => { + return send_password_authentication_failed(client).await; + } + }; + client.set_state(PgWireConnectionState::AuthenticationInProgress); - client - .send(PgWireBackendMessage::Authentication( - Authentication::CleartextPassword, - )) - .await?; + match auth_info { + PgAuthInfo::ScramSha256 { .. } => { + *self.login_verifier.state.lock().await = + PgAuthenticationState::SaslInitial { auth_info }; + client + .send(PgWireBackendMessage::Authentication(Authentication::SASL( + vec![SCRAM_SHA_256_METHOD.to_string()], + ))) + .await?; + } + PgAuthInfo::Cleartext => { + *self.login_verifier.state.lock().await = + PgAuthenticationState::Cleartext; + client + .send(PgWireBackendMessage::Authentication( + Authentication::CleartextPassword, + )) + .await?; + } + } } else { self.session.set_user_info(userinfo_by_name( client.metadata().get(super::METADATA_USER).cloned(), @@ -205,27 +308,20 @@ impl StartupHandler for PostgresServerHandlerInner { } } PgWireFrontendMessage::PasswordMessageFamily(pwd) => { - // the newer version of pgwire has a few variant password - // message like cleartext/md5 password, saslresponse, etc. Here - // we must manually coerce it into password - let pwd = pwd.into_password()?; - let login_info = LoginInfo::from_client_info(client); - - // do authenticate - let auth_result = self.login_verifier.auth(&login_info, &pwd.password).await; - - if let Ok(Some(user_info)) = auth_result { - self.session.set_user_info(user_info); - set_client_info(client, &self.session); - auth::finish_authentication(client, self.param_provider.as_ref()).await?; - } else { - return send_error( - client, - PgErrorCode::Ec28P01 - .to_err_info("password authentication failed".to_string()), - ) - .await; + match self + .authenticate_password_message(client, &login_info, pwd) + .await? + { + PgAuthenticationResult::Continue => {} + PgAuthenticationResult::Success(user_info) => { + self.session.set_user_info(user_info); + set_client_info(client, &self.session); + auth::finish_authentication(client, self.param_provider.as_ref()).await?; + } + PgAuthenticationResult::Failed => { + return send_password_authentication_failed(client).await; + } } } _ => {} @@ -234,6 +330,225 @@ impl StartupHandler for PostgresServerHandlerInner { } } +enum PgAuthenticationResult { + Continue, + Success(UserInfoRef), + Failed, +} + +impl PostgresServerHandlerInner { + /// Records a rejected SCRAM attempt in [`METRIC_AUTH_FAILURE`]. The label is + /// intentionally uniform (never `UserNotFound`), so the counter cannot be + /// used to distinguish a wrong password from an unknown user. + fn record_scram_failure(result: PgAuthenticationResult) -> PgAuthenticationResult { + if matches!(result, PgAuthenticationResult::Failed) { + METRIC_AUTH_FAILURE + .with_label_values(&[StatusCode::UserPasswordMismatch.as_ref()]) + .inc(); + } + result + } + + async fn authenticate_password_message( + &self, + client: &mut C, + login_info: &LoginInfo, + pwd: PasswordMessageFamily, + ) -> PgWireResult + where + C: ClientInfo + Sink + Unpin + Send, + C::Error: Debug, + PgWireError: From<>::Error>, + { + let mut state = self.login_verifier.state.lock().await; + let current_state = std::mem::replace(&mut *state, PgAuthenticationState::Initial); + + match current_state { + PgAuthenticationState::Cleartext => { + let pwd = pwd.into_password()?; + drop(state); + match self.login_verifier.auth(login_info, &pwd.password).await { + Ok(Some(user_info)) => Ok(PgAuthenticationResult::Success(user_info)), + Ok(None) | Err(_) => Ok(PgAuthenticationResult::Failed), + } + } + PgAuthenticationState::SaslInitial { auth_info } => { + // A labeled block funnels every rejection through a single + // `record_scram_failure`, so failed SCRAM attempts still show up + // in `METRIC_AUTH_FAILURE` without sprinkling the metric call + // across each early return. + let result = 'sasl: { + let sasl_initial = pwd.into_sasl_initial_response()?; + if sasl_initial.auth_method != SCRAM_SHA_256_METHOD { + break 'sasl PgAuthenticationResult::Failed; + } + let Some(data) = sasl_initial.data else { + break 'sasl PgAuthenticationResult::Failed; + }; + let Some(client_first) = ScramClientFirst::parse(&data) else { + break 'sasl PgAuthenticationResult::Failed; + }; + let PgAuthInfo::ScramSha256 { + verifier, + user_info, + } = auth_info + else { + break 'sasl PgAuthenticationResult::Failed; + }; + + let server_nonce = BASE64.encode(rand::random::<[u8; 18]>()); + let nonce = format!("{}{}", client_first.nonce, server_nonce); + let server_first = format!( + "r={},s={},i={}", + nonce, + BASE64.encode(verifier.salt()), + verifier.iterations() + ); + client + .send(PgWireBackendMessage::Authentication( + Authentication::SASLContinue(server_first.clone().into()), + )) + .await?; + *state = PgAuthenticationState::SaslFinal { + verifier, + user_info, + channel_binding: client_first.channel_binding, + nonce, + client_first_bare: client_first.bare, + server_first, + }; + PgAuthenticationResult::Continue + }; + Ok(Self::record_scram_failure(result)) + } + PgAuthenticationState::SaslFinal { + verifier, + user_info, + channel_binding, + nonce, + client_first_bare, + server_first, + } => { + let result = 'sasl: { + let sasl_response = pwd.into_sasl_response()?; + let Some(client_final) = ScramClientFinal::parse(&sasl_response.data) else { + break 'sasl PgAuthenticationResult::Failed; + }; + if client_final.channel_binding != channel_binding { + break 'sasl PgAuthenticationResult::Failed; + } + if client_final.nonce != nonce { + break 'sasl PgAuthenticationResult::Failed; + } + + let auth_message = format!( + "{},{},{}", + client_first_bare, server_first, client_final.without_proof + ); + let Ok(client_proof) = BASE64.decode(client_final.proof.as_bytes()) else { + break 'sasl PgAuthenticationResult::Failed; + }; + let Ok(Some(server_signature)) = + verifier.verify_client_proof(auth_message.as_bytes(), &client_proof) + else { + break 'sasl PgAuthenticationResult::Failed; + }; + let Some(user_info) = user_info else { + break 'sasl PgAuthenticationResult::Failed; + }; + + drop(state); + if self + .login_verifier + .authorize(login_info, &user_info) + .await + .is_err() + { + // `authorize` already recorded this failure; return early + // to bypass `record_scram_failure` and avoid double-counting. + return Ok(PgAuthenticationResult::Failed); + } + + client + .send(PgWireBackendMessage::Authentication( + Authentication::SASLFinal( + format!("v={}", BASE64.encode(server_signature)).into(), + ), + )) + .await?; + PgAuthenticationResult::Success(user_info) + }; + Ok(Self::record_scram_failure(result)) + } + PgAuthenticationState::Initial => Ok(PgAuthenticationResult::Failed), + } + } +} + +struct ScramClientFirst { + channel_binding: String, + bare: String, + nonce: String, +} + +impl ScramClientFirst { + fn parse(data: &[u8]) -> Option { + let message = std::str::from_utf8(data).ok()?; + let mut parts = message.splitn(3, ','); + let cbind = parts.next()?; + if !matches!(cbind, "n" | "y") { + return None; + } + let authzid = parts.next()?; + if !authzid.is_empty() { + return None; + } + let channel_binding = BASE64.encode(format!("{cbind},,")); + let bare = parts.next()?.to_string(); + let nonce = bare + .split(',') + .find_map(|chunk| chunk.strip_prefix("r="))? + .to_string(); + Some(Self { + channel_binding, + bare, + nonce, + }) + } +} + +struct ScramClientFinal { + channel_binding: String, + nonce: String, + without_proof: String, + proof: String, +} + +impl ScramClientFinal { + fn parse(data: &[u8]) -> Option { + let message = std::str::from_utf8(data).ok()?; + let proof_pos = message.rfind(",p=")?; + let without_proof = message[..proof_pos].to_string(); + let proof = message[proof_pos + 3..].to_string(); + let channel_binding = without_proof + .split(',') + .find_map(|chunk| chunk.strip_prefix("c="))?; + let nonce = without_proof + .split(',') + .find_map(|chunk| chunk.strip_prefix("r="))?; + if nonce.is_empty() || proof.is_empty() { + return None; + } + + Some(Self { + channel_binding: channel_binding.to_string(), + nonce: nonce.to_string(), + without_proof, + proof, + }) + } +} + async fn send_error(client: &mut C, err_info: ErrorInfo) -> PgWireResult<()> where C: ClientInfo + Sink + Unpin + Send, @@ -248,6 +563,19 @@ where Ok(()) } +async fn send_password_authentication_failed(client: &mut C) -> PgWireResult<()> +where + C: ClientInfo + Sink + Unpin + Send, + C::Error: Debug, + PgWireError: From<>::Error>, +{ + send_error( + client, + PgErrorCode::Ec28P01.to_err_info("password authentication failed".to_string()), + ) + .await +} + enum DbResolution { Resolved(String, String), NotFound(String), @@ -277,3 +605,43 @@ where Ok(DbResolution::NotFound("Database not specified".to_owned())) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_scram_client_first_parse() { + let message = b"n,,n=greptime,r=clientnonce"; + let client_first = ScramClientFirst::parse(message).unwrap(); + assert_eq!("biws", client_first.channel_binding); + assert_eq!("n=greptime,r=clientnonce", client_first.bare); + assert_eq!("clientnonce", client_first.nonce); + + let message = b"y,,n=greptime,r=clientnonce"; + let client_first = ScramClientFirst::parse(message).unwrap(); + assert_eq!("eSws", client_first.channel_binding); + assert_eq!("n=greptime,r=clientnonce", client_first.bare); + assert_eq!("clientnonce", client_first.nonce); + + assert!(ScramClientFirst::parse(b"p=tls-server-end-point,,n=greptime,r=nonce").is_none()); + assert!(ScramClientFirst::parse(b"n,a=authzid,n=greptime,r=nonce").is_none()); + } + + #[test] + fn test_scram_client_final_parse() { + let message = b"c=biws,r=clientnonceservernonce,p=dGVzdA=="; + let client_final = ScramClientFinal::parse(message).unwrap(); + assert_eq!("biws", client_final.channel_binding); + assert_eq!("clientnonceservernonce", client_final.nonce); + assert_eq!( + "c=biws,r=clientnonceservernonce", + client_final.without_proof + ); + assert_eq!("dGVzdA==", client_final.proof); + + assert!(ScramClientFinal::parse(b"r=nonce,p=dGVzdA==").is_none()); + assert!(ScramClientFinal::parse(b"c=biws,r=nonce").is_none()); + assert!(ScramClientFinal::parse(b"c=biws,r=,p=dGVzdA==").is_none()); + } +} diff --git a/src/servers/src/postgres/types.rs b/src/servers/src/postgres/types.rs index 74e230f83af..465f468f1d5 100644 --- a/src/servers/src/postgres/types.rs +++ b/src/servers/src/postgres/types.rs @@ -582,9 +582,10 @@ pub(super) fn parameters_to_scalar_values( Some(st @ ConcreteDataType::Timestamp(unit)) => { to_timestamp_scalar_value(data.and_then(|n| n.to_i64()), unit, st)? } - Some(ConcreteDataType::UInt64(_)) | None => { - ScalarValue::UInt64(data.and_then(|n| n.to_u64())) - } + Some(ConcreteDataType::UInt64(_)) | None => ScalarValue::UInt64( + data.map(|n| n.to_u64().ok_or_else(|| numeric_out_of_range_error(n))) + .transpose()?, + ), Some(st) => { return Err(invalid_parameter_error( "invalid_parameter_type", @@ -924,21 +925,28 @@ pub(super) fn parameters_to_scalar_values( &Type::NUMERIC_ARRAY => { let data = portal.parameter::>>(idx, &client_type)?; if let Some(data) = data { - let build_u64_list = |data: Vec>| { + let build_u64_list = |data: Vec>| -> PgWireResult { let values = data .into_iter() - .map(|n| ScalarValue::UInt64(n.and_then(|n| n.to_u64()))) - .collect::>(); - ScalarValue::List(ScalarValue::new_list( + .map(|n| { + Ok(ScalarValue::UInt64( + n.map(|n| { + n.to_u64().ok_or_else(|| numeric_out_of_range_error(n)) + }) + .transpose()?, + )) + }) + .collect::>>()?; + Ok(ScalarValue::List(ScalarValue::new_list( &values, &ArrowDataType::UInt64, true, - )) + ))) }; if let Some(server_type) = &server_type { match server_type { ConcreteDataType::List(list_type) => match list_type.item_type() { - ConcreteDataType::UInt64(_) => build_u64_list(data), + ConcreteDataType::UInt64(_) => build_u64_list(data)?, ConcreteDataType::Decimal128(dt) => { let values = data .into_iter() @@ -975,7 +983,7 @@ pub(super) fn parameters_to_scalar_values( } } else { // server type not provided - build_u64_list(data) + build_u64_list(data)? } } else { ScalarValue::Null @@ -1047,17 +1055,22 @@ pub(super) fn parameters_to_scalar_values( TimestampType::Nanosecond(_) => { let values = data .into_iter() - .filter_map(|ts| { - ts.and_then(|ts| { - ts.and_utc().timestamp_nanos_opt().map(|nanos| { + .map(|ts| match ts { + None => { + Ok(ScalarValue::TimestampNanosecond(None, None)) + } + Some(ts) => ts + .and_utc() + .timestamp_nanos_opt() + .map(|nanos| { ScalarValue::TimestampNanosecond( Some(nanos), None, ) }) - }) + .ok_or_else(|| numeric_out_of_range_error(ts)), }) - .collect::>(); + .collect::>>()?; ScalarValue::List(ScalarValue::new_list( &values, &ArrowDataType::Timestamp(TimeUnit::Nanosecond, None), @@ -2025,4 +2038,161 @@ mod test { let values = parameters_to_scalar_values(&plan, &portal).unwrap(); assert_eq!(values[0], ScalarValue::Int8(None)); } + + fn numeric_uint64_array_plan() -> LogicalPlan { + build_plan_with_params(vec![( + "$1", + DataType::List(Arc::new(Field::new("item", DataType::UInt64, true))), + )]) + } + + fn assert_numeric_out_of_range(result: PgWireResult>) { + match result.unwrap_err() { + PgWireError::UserError(error) => { + assert_eq!("22023", error.code); + assert_eq!("numeric_value_out_of_range", error.message); + } + error => panic!("expected numeric out-of-range error, got {error:?}"), + } + } + + fn qbs_timestamp_nanosecond_array_plan() -> LogicalPlan { + build_plan_with_params(vec![( + "$1", + DataType::List(Arc::new(Field::new( + "item", + DataType::Timestamp(TimeUnit::Nanosecond, None), + true, + ))), + )]) + } + + #[test] + fn test_qbs_pg_timestamp_nanosecond_array_preserves_null_slot() { + let portal = make_portal( + vec![Some(Type::TIMESTAMP_ARRAY)], + vec![s( + r#"{"2024-01-01 00:00:00.000001",NULL,"2024-01-01 00:00:00.000003"}"#, + )], + ); + + let values = + parameters_to_scalar_values(&qbs_timestamp_nanosecond_array_plan(), &portal).unwrap(); + let expected = ScalarValue::List(ScalarValue::new_list( + &[ + ScalarValue::TimestampNanosecond(Some(1_704_067_200_000_001_000), None), + ScalarValue::TimestampNanosecond(None, None), + ScalarValue::TimestampNanosecond(Some(1_704_067_200_000_003_000), None), + ], + &ArrowDataType::Timestamp(TimeUnit::Nanosecond, None), + true, + )); + assert_eq!(expected, values[0]); + } + + #[test] + fn test_qbs_pg_timestamp_nanosecond_array_out_of_range_rejected() { + let portal = make_portal( + vec![Some(Type::TIMESTAMP_ARRAY)], + vec![s(r#"{"3000-01-01 00:00:00"}"#)], + ); + + assert_numeric_out_of_range(parameters_to_scalar_values( + &qbs_timestamp_nanosecond_array_plan(), + &portal, + )); + } + + #[test] + fn test_numeric_uint64_scalar_negative_rejected() { + let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]); + let portal = make_portal(vec![Some(Type::NUMERIC)], vec![s("-1")]); + + assert_numeric_out_of_range(parameters_to_scalar_values(&plan, &portal)); + } + + #[test] + fn test_numeric_uint64_scalar_above_u64_max_rejected() { + let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]); + let portal = make_portal(vec![Some(Type::NUMERIC)], vec![s("18446744073709551616")]); + + assert_numeric_out_of_range(parameters_to_scalar_values(&plan, &portal)); + } + + #[test] + fn test_numeric_uint64_scalar_null_preserved() { + let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]); + let portal = make_portal(vec![Some(Type::NUMERIC)], vec![None]); + + let values = parameters_to_scalar_values(&plan, &portal).unwrap(); + assert_eq!(ScalarValue::UInt64(None), values[0]); + } + + #[test] + fn test_numeric_uint64_scalar_u64_max_preserved() { + let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]); + let portal = make_portal(vec![Some(Type::NUMERIC)], vec![s(&u64::MAX.to_string())]); + + let values = parameters_to_scalar_values(&plan, &portal).unwrap(); + assert_eq!(ScalarValue::UInt64(Some(u64::MAX)), values[0]); + } + + #[test] + fn test_numeric_uint64_array_outer_null_preserved() { + let portal = make_portal(vec![Some(Type::NUMERIC_ARRAY)], vec![None]); + + let values = parameters_to_scalar_values(&numeric_uint64_array_plan(), &portal).unwrap(); + assert_eq!(ScalarValue::Null, values[0]); + } + + #[test] + fn test_numeric_uint64_array_preserves_values_and_null_slots() { + let portal = make_portal( + vec![Some(Type::NUMERIC_ARRAY)], + vec![s("{42,NULL,18446744073709551615}")], + ); + + let values = parameters_to_scalar_values(&numeric_uint64_array_plan(), &portal).unwrap(); + let expected = ScalarValue::List(ScalarValue::new_list( + &[ + ScalarValue::UInt64(Some(42)), + ScalarValue::UInt64(None), + ScalarValue::UInt64(Some(u64::MAX)), + ], + &ArrowDataType::UInt64, + true, + )); + assert_eq!(expected, values[0]); + } + + #[test] + fn test_numeric_uint64_array_invalid_after_valid_and_null_prefix_rejected() { + let portal = make_portal(vec![Some(Type::NUMERIC_ARRAY)], vec![s("{42,NULL,-1}")]); + + assert_numeric_out_of_range(parameters_to_scalar_values( + &numeric_uint64_array_plan(), + &portal, + )); + } + + #[test] + fn test_numeric_uint64_array_above_u64_max_rejected() { + let portal = make_portal( + vec![Some(Type::NUMERIC_ARRAY)], + vec![s("{18446744073709551616}")], + ); + + assert_numeric_out_of_range(parameters_to_scalar_values( + &numeric_uint64_array_plan(), + &portal, + )); + } + + #[test] + fn test_numeric_uint64_uninferred_array_invalid_value_rejected() { + let plan = LogicalPlanBuilder::empty(true).build().unwrap(); + let portal = make_portal(vec![Some(Type::NUMERIC_ARRAY)], vec![s("{-1}")]); + + assert_numeric_out_of_range(parameters_to_scalar_values(&plan, &portal)); + } } diff --git a/src/servers/src/row_writer.rs b/src/servers/src/row_writer.rs index fb3351323c9..5b7abe151bf 100644 --- a/src/servers/src/row_writer.rs +++ b/src/servers/src/row_writer.rs @@ -96,17 +96,30 @@ impl TableData { value: Option, one_row: &mut Vec, ) { - let name = name.to_string(); - if let Some(index) = self.column_indexes.get(&name).copied() { - one_row[index].value_data = value; - } else { - let index = self.schema.len(); - self.schema.push(ColumnSchema { - column_name: name.clone(), + self.write_column_unchecked( + ColumnSchema { + column_name: name.to_string(), datatype: datatype as i32, semantic_type: SemanticType::Field as i32, ..Default::default() - }); + }, + value, + one_row, + ); + } + + pub fn write_column_unchecked( + &mut self, + column_schema: ColumnSchema, + value: Option, + one_row: &mut Vec, + ) { + if let Some(index) = self.column_indexes.get(&column_schema.column_name).copied() { + one_row[index].value_data = value; + } else { + let index = self.schema.len(); + let name = column_schema.column_name.clone(); + self.schema.push(column_schema); self.column_indexes.insert(name, index); one_row.push(Value { value_data: value }); } @@ -235,7 +248,7 @@ pub fn write_f64( ) } -fn build_json_column_schema(name: impl ToString) -> ColumnSchema { +pub(crate) fn build_json_column_schema(name: impl ToString) -> ColumnSchema { ColumnSchema { column_name: name.to_string(), datatype: ColumnDataType::Binary as i32, diff --git a/src/servers/tests/postgres/mod.rs b/src/servers/tests/postgres/mod.rs index 5a3e97e3e71..35489687d97 100644 --- a/src/servers/tests/postgres/mod.rs +++ b/src/servers/tests/postgres/mod.rs @@ -16,8 +16,8 @@ use std::net::SocketAddr; use std::sync::Arc; use std::time::Duration; -use auth::UserProviderRef; use auth::tests::{DatabaseAuthInfo, MockUserProvider}; +use auth::{UserProviderRef, format_pg_scram_sha256_password_verifier, user_provider_from_option}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_runtime::Builder as RuntimeBuilder; use common_runtime::runtime::BuilderBuild; @@ -43,12 +43,6 @@ fn create_postgres_server( tls: TlsOption, auth_info: Option, ) -> Result> { - let instance = Arc::new(create_testing_instance(table)); - let io_runtime = RuntimeBuilder::default() - .worker_threads(4) - .thread_name("postgres-io-handlers") - .build() - .unwrap(); let user_provider: Option = if check_pwd { let mut provider = MockUserProvider::default(); if let Some(info) = auth_info { @@ -58,6 +52,20 @@ fn create_postgres_server( } else { None }; + create_postgres_server_with_user_provider(table, tls, user_provider) +} + +fn create_postgres_server_with_user_provider( + table: TableRef, + tls: TlsOption, + user_provider: Option, +) -> Result> { + let instance = Arc::new(create_testing_instance(table)); + let io_runtime = RuntimeBuilder::default() + .worker_threads(4) + .thread_name("postgres-io-handlers") + .build() + .unwrap(); let tls_server_config = Arc::new( ReloadableTlsServerConfig::try_new(tls.clone()) @@ -75,6 +83,21 @@ fn create_postgres_server( ))) } +async fn start_test_server_with_user_provider( + user_provider: UserProviderRef, +) -> Result<(Box, u16)> { + common_telemetry::init_default_ut_logging(); + let _ = install_default_crypto_provider(); + + let table = MemTable::default_numbers_table(); + let mut postgres_server = + create_postgres_server_with_user_provider(table, Default::default(), Some(user_provider))?; + let listening = "127.0.0.1:0".parse::().unwrap(); + postgres_server.start(listening).await.unwrap(); + let server_addr = postgres_server.bind_addr().unwrap(); + Ok((postgres_server, server_addr.port())) +} + #[tokio::test] pub async fn test_start_postgres_server() -> Result<()> { let table = MemTable::default_numbers_table(); @@ -138,6 +161,61 @@ async fn test_schema_validating() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_pg_scram_sha256_auth() -> Result<()> { + let verifier = + format_pg_scram_sha256_password_verifier(b"greptime", b"pg-scram-salt", 4096).unwrap(); + let user_provider = + user_provider_from_option(&format!("static_user_provider:cmd:greptime={verifier}")) + .unwrap(); + let (postgres_server, server_port) = + start_test_server_with_user_provider(user_provider).await?; + + let client = create_plain_connection_with_credentials(server_port, "greptime", "greptime") + .await + .unwrap(); + let rows = client + .simple_query("SELECT uint32s FROM numbers LIMIT 1") + .await; + assert_eq!(unwrap_results(rows.unwrap().as_ref())[0], "0"); + + let wrong_password = + create_plain_connection_with_credentials(server_port, "greptime", "wrong").await; + assert!(wrong_password.is_err()); + + let unknown_user = + create_plain_connection_with_credentials(server_port, "not_found", "greptime").await; + assert!(unknown_user.is_err()); + + postgres_server.shutdown().await.unwrap(); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_pg_cleartext_auth_fallback() -> Result<()> { + common_telemetry::init_default_ut_logging(); + let table = MemTable::default_numbers_table(); + let mut postgres_server = create_postgres_server(table, true, Default::default(), None)?; + let listening = "127.0.0.1:0".parse::().unwrap(); + postgres_server.start(listening).await.unwrap(); + let server_port = postgres_server.bind_addr().unwrap().port(); + + let client = create_plain_connection_with_credentials(server_port, "greptime", "greptime") + .await + .unwrap(); + let rows = client + .simple_query("SELECT uint32s FROM numbers LIMIT 1") + .await; + assert_eq!(unwrap_results(rows.unwrap().as_ref())[0], "0"); + + let wrong_password = + create_plain_connection_with_credentials(server_port, "greptime", "wrong").await; + assert!(wrong_password.is_err()); + + postgres_server.shutdown().await.unwrap(); + Ok(()) +} + // #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_shutdown_pg_server(with_pwd: bool) -> Result<()> { common_telemetry::init_default_ut_logging(); @@ -433,6 +511,19 @@ async fn create_plain_connection( Ok(client) } +async fn create_plain_connection_with_credentials( + port: u16, + username: &str, + password: &str, +) -> std::result::Result { + let url = format!( + "host=127.0.0.1 port={port} user={username} password={password} connect_timeout=2 dbname={DEFAULT_SCHEMA_NAME}", + ); + let (client, conn) = tokio_postgres::connect(&url, NoTls).await?; + let _handle = tokio::spawn(conn); + Ok(client) +} + async fn create_connection_with_given_db( port: u16, db: &str, diff --git a/src/sql/src/statements/set_variables.rs b/src/sql/src/statements/set_variables.rs index 4c6eb16692f..346f1762e62 100644 --- a/src/sql/src/statements/set_variables.rs +++ b/src/sql/src/statements/set_variables.rs @@ -15,7 +15,7 @@ use std::fmt::Display; use serde::Serialize; -use sqlparser::ast::{Expr, ObjectName}; +use sqlparser::ast::{Expr, ObjectName, ObjectNamePart, Value, ValueWithSpan}; use sqlparser_derive::{Visit, VisitMut}; /// SET variables statement. @@ -25,6 +25,27 @@ pub struct SetVariables { pub value: Vec, } +impl SetVariables { + /// Returns the first supported `search_path` value. + pub fn search_path(&self) -> Option<&str> { + let [ObjectNamePart::Identifier(variable)] = self.variable.0.as_slice() else { + return None; + }; + if variable.quote_style.is_some() || !variable.value.eq_ignore_ascii_case("search_path") { + return None; + } + + match self.value.first()? { + Expr::Value(ValueWithSpan { + value: Value::SingleQuotedString(value) | Value::DoubleQuotedString(value), + .. + }) => Some(value), + Expr::Identifier(value) => Some(&value.value), + _ => None, + } + } +} + impl Display for SetVariables { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let variable = &self.variable; @@ -70,4 +91,26 @@ SET delayed_insert_timeout = 300"#, } } } + + #[test] + fn test_search_path() { + for (sql, expected) in [ + ("SET search_path TO public", Some("public")), + ("SET SEARCH_PATH TO 'private'", Some("private")), + ("SET search_path TO \"quoted\"", Some("quoted")), + ("SET \"search_path\" TO private", None), + ("SET timezone TO 'UTC'", None), + ] { + let statements = ParserContext::create_with_dialect( + sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + let [Statement::SetVariables(set)] = statements.as_slice() else { + unreachable!() + }; + assert_eq!(expected, set.search_path()); + } + } } diff --git a/src/sql/src/statements/statement.rs b/src/sql/src/statements/statement.rs index bac7e367cd9..ec7c73573b5 100644 --- a/src/sql/src/statements/statement.rs +++ b/src/sql/src/statements/statement.rs @@ -178,7 +178,6 @@ impl Statement { | Statement::ShowSearchPath(_) | Statement::ShowViews(_) | Statement::DescribeTable(_) - | Statement::Explain(_) | Statement::ShowVariables(_) | Statement::ShowProcesslist(_) | Statement::FetchCursor(_) @@ -189,6 +188,8 @@ impl Statement { #[cfg(feature = "enterprise")] Statement::ShowTriggers(_) => true, + Statement::Explain(explain) => !explain.analyze || explain.statement.is_readonly(), + // Write operations Statement::Insert(_) | Statement::Delete(_) diff --git a/src/sql/src/util.rs b/src/sql/src/util.rs index 6a7abe23bb5..5868be6aa0e 100644 --- a/src/sql/src/util.rs +++ b/src/sql/src/util.rs @@ -14,20 +14,19 @@ use std::collections::HashSet; use std::fmt::{Display, Formatter}; +use std::ops::ControlFlow; use itertools::Itertools; use promql_parser::label::{METRIC_NAME, MatchOp}; use promql_parser::parser::{ - AggregateExpr as PromAggregateExpr, BinaryExpr as PromBinaryExpr, Call as PromCall, - Expr as PromExpr, MatrixSelector as PromMatrixSelector, ParenExpr as PromParenExpr, - SubqueryExpr as PromSubqueryExpr, UnaryExpr as PromUnaryExpr, - VectorSelector as PromVectorSelector, + Expr as PromExpr, MatrixSelector as PromMatrixSelector, VectorSelector as PromVectorSelector, }; +use promql_parser::util::{ExprVisitor, walk_expr}; use serde::Serialize; use snafu::ensure; use sqlparser::ast::{ - Array, Expr, Ident, ObjectName, ObjectNamePart, SetExpr, SqlOption, TableFactor, - TableWithJoins, Value, ValueWithSpan, + Array, Expr, Ident, ObjectName, ObjectNamePart, SqlOption, Value, ValueWithSpan, + Visit as AstVisit, Visitor, }; use sqlparser_derive::{Visit, VisitMut}; @@ -184,48 +183,60 @@ pub fn parse_option_string(option: SqlOption) -> Result<(String, OptionValue)> { /// Walk through a [Query] and extract all the tables referenced in it. pub fn extract_tables_from_query(query: &SqlOrTql) -> impl Iterator { - let mut names = HashSet::new(); - - match query { - SqlOrTql::Sql(query, _) => { - extract_tables_from_sql_query(&query.inner, &mut names); - extract_tables_from_hybrid_cte_query(query, &mut names); - } - SqlOrTql::Tql(tql, _) => extract_tables_from_tql(tql, &mut names), - } - - names.into_iter() + extract_tables_from_query_inner(query).0.into_iter() } -fn extract_tables_from_hybrid_cte_query(query: &Query, sql_names: &mut HashSet) { - if let Some(hybrid_cte) = &query.hybrid_cte { - let mut cte_names: HashSet = hybrid_cte - .cte_tables - .iter() - .map(|cte| ParserContext::canonicalize_identifier(cte.name.clone()).value) - .collect(); - remove_cte_names(sql_names, &cte_names); +/// Walk through a [Query] and extract its referenced tables, returning `None` +/// if any table reference cannot be resolved statically. +pub fn extract_tables_from_query_checked( + query: &SqlOrTql, +) -> Option> { + let (names, complete) = extract_tables_from_query_inner(query); + complete.then_some(names.into_iter()) +} - cte_names.clear(); - for cte in &hybrid_cte.cte_tables { - let cte_name = ParserContext::canonicalize_identifier(cte.name.clone()).value; - let mut cte_query_names = HashSet::new(); - match &cte.content { - CteContent::Sql(cte_query) => { - extract_tables_from_sql_query(cte_query, &mut cte_query_names) - } - CteContent::Tql(tql) => extract_tables_from_tql(tql, &mut cte_query_names), - } - if hybrid_cte.recursive { - cte_names.insert(cte_name.clone()); - } - remove_cte_names(&mut cte_query_names, &cte_names); - sql_names.extend(cte_query_names); - if !hybrid_cte.recursive { - cte_names.insert(cte_name); - } +fn extract_tables_from_query_inner(query: &SqlOrTql) -> (HashSet, bool) { + let mut names = HashSet::new(); + + let complete = match query { + SqlOrTql::Sql(query, _) => { + extract_tables_from_sql_query(&query.inner, &mut names); + extract_tables_from_hybrid_cte_query(query, &mut names) } + SqlOrTql::Tql(tql, _) => extract_tables_from_tql(tql, &mut names), + }; + + (names, complete) +} + +fn extract_tables_from_hybrid_cte_query( + query: &Query, + sql_names: &mut HashSet, +) -> bool { + let Some(hybrid_cte) = &query.hybrid_cte else { + return true; + }; + + let mut complete = true; + let cte_names: HashSet = hybrid_cte + .cte_tables + .iter() + .map(|cte| ParserContext::canonicalize_identifier(cte.name.clone()).value) + .collect(); + remove_cte_names(sql_names, &cte_names); + + for cte in &hybrid_cte.cte_tables { + let mut cte_query_names = HashSet::new(); + match &cte.content { + CteContent::Sql(cte_query) => { + extract_tables_from_sql_query(cte_query, &mut cte_query_names) + } + CteContent::Tql(tql) => complete &= extract_tables_from_tql(tql, &mut cte_query_names), + } + sql_names.extend(cte_query_names); } + + complete } fn remove_cte_names(names: &mut HashSet, cte_names: &HashSet) { @@ -246,55 +257,55 @@ fn remove_cte_names(names: &mut HashSet, cte_names: &HashSet }); } -fn extract_tables_from_tql(tql: &Tql, names: &mut HashSet) { +fn extract_tables_from_tql(tql: &Tql, names: &mut HashSet) -> bool { let promql = match tql { Tql::Eval(eval) => &eval.query, Tql::Explain(explain) => &explain.query, Tql::Analyze(analyze) => &analyze.query, }; - if let Ok(expr) = promql_parser::parser::parse(promql) { - extract_tables_from_prom_expr(&expr, names); - } + let Ok(expr) = promql_parser::parser::parse(promql) else { + return false; + }; + extract_tables_from_prom_expr(&expr, names) } -fn extract_tables_from_prom_expr(expr: &PromExpr, names: &mut HashSet) { - match expr { - PromExpr::Aggregate(PromAggregateExpr { expr, .. }) => { - extract_tables_from_prom_expr(expr, names); - } - PromExpr::Unary(PromUnaryExpr { expr, .. }) => { - extract_tables_from_prom_expr(expr, names); - } - PromExpr::Binary(PromBinaryExpr { lhs, rhs, .. }) => { - extract_tables_from_prom_expr(lhs, names); - extract_tables_from_prom_expr(rhs, names); - } - PromExpr::Paren(PromParenExpr { expr }) => { - extract_tables_from_prom_expr(expr, names); - } - PromExpr::Subquery(PromSubqueryExpr { expr, .. }) => { - extract_tables_from_prom_expr(expr, names); - } - PromExpr::VectorSelector(selector) => { - extract_metric_name_from_vector_selector(selector, names); - } - PromExpr::MatrixSelector(PromMatrixSelector { vs, .. }) => { - extract_metric_name_from_vector_selector(vs, names); - } - PromExpr::Call(PromCall { args, .. }) => { - for arg in &args.args { - extract_tables_from_prom_expr(arg, names); - } - } - PromExpr::NumberLiteral(_) | PromExpr::StringLiteral(_) | PromExpr::Extension(_) => {} +fn extract_tables_from_prom_expr(expr: &PromExpr, names: &mut HashSet) -> bool { + struct TableCollector<'a> { + names: &'a mut HashSet, + complete: bool, } + + impl ExprVisitor for TableCollector<'_> { + type Error = (); + + fn pre_visit(&mut self, expr: &PromExpr) -> std::result::Result { + self.complete &= match expr { + PromExpr::VectorSelector(selector) => { + extract_metric_name_from_vector_selector(selector, self.names) + } + PromExpr::MatrixSelector(PromMatrixSelector { vs, .. }) => { + extract_metric_name_from_vector_selector(vs, self.names) + } + PromExpr::Extension(_) => false, + _ => true, + }; + Ok(true) + } + } + + let mut collector = TableCollector { + names, + complete: true, + }; + let _ = walk_expr(&mut collector, expr); + collector.complete } fn extract_metric_name_from_vector_selector( selector: &PromVectorSelector, names: &mut HashSet, -) { +) -> bool { let metric_name = selector.name.clone().or_else(|| { let mut metric_name_matchers = selector.matchers.find_matchers(METRIC_NAME); if metric_name_matchers.len() == 1 && metric_name_matchers[0].op == MatchOp::Equal { @@ -304,9 +315,16 @@ fn extract_metric_name_from_vector_selector( } }); let Some(metric_name) = metric_name else { - return; + return false; }; + if selector.matchers.matchers.iter().any(|matcher| { + (matcher.name == SCHEMA_MATCHER || matcher.name == DATABASE_MATCHER) + && matcher.op != MatchOp::Equal + }) { + return false; + } + let schema_matcher = selector.matchers.matchers.iter().rev().find(|matcher| { matcher.op == MatchOp::Equal && (matcher.name == SCHEMA_MATCHER || matcher.name == DATABASE_MATCHER) @@ -322,6 +340,7 @@ fn extract_metric_name_from_vector_selector( metric_name, ))])); } + true } /// translate the start location to the index in the sql string @@ -340,96 +359,116 @@ pub fn location_to_index(sql: &str, location: &sqlparser::tokenizer::Location) - index - 1 } +// Tracks CTE aliases as sqlparser's visitor enters their query bodies. +struct QueryScope { + cte_names: Vec, + next_cte: usize, + recursive: bool, + scope_start: usize, + add_to_parent_after: Option, +} + +struct RelationCollector<'a> { + names: &'a mut HashSet, + ctes_in_scope: Vec, + query_scopes: Vec, + #[cfg(test)] + query_visits: usize, +} + +impl<'a> RelationCollector<'a> { + fn new(names: &'a mut HashSet) -> Self { + Self { + names, + ctes_in_scope: Vec::new(), + query_scopes: Vec::new(), + #[cfg(test)] + query_visits: 0, + } + } +} + +impl Visitor for RelationCollector<'_> { + type Break = (); + + fn pre_visit_query(&mut self, query: &sqlparser::ast::Query) -> ControlFlow { + #[cfg(test)] + { + self.query_visits += 1; + } + + // Query::visit enters WITH definitions in declaration order before its body. + let parent_cte = self.query_scopes.last_mut().and_then(|scope| { + let cte_name = scope.cte_names.get(scope.next_cte)?.clone(); + scope.next_cte += 1; + Some((cte_name, scope.recursive)) + }); + + let add_to_parent_after = match parent_cte { + Some((cte_name, true)) => { + self.ctes_in_scope.push(cte_name); + None + } + Some((cte_name, false)) => Some(cte_name), + None => None, + }; + + let scope_start = self.ctes_in_scope.len(); + let (cte_names, recursive) = if let Some(with) = &query.with { + ( + with.cte_tables + .iter() + .map(|cte| ParserContext::canonicalize_identifier(cte.alias.name.clone()).value) + .collect(), + with.recursive, + ) + } else { + (Vec::new(), false) + }; + + self.query_scopes.push(QueryScope { + cte_names, + next_cte: 0, + recursive, + scope_start, + add_to_parent_after, + }); + ControlFlow::Continue(()) + } + + fn post_visit_query(&mut self, _query: &sqlparser::ast::Query) -> ControlFlow { + let Some(scope) = self.query_scopes.pop() else { + return ControlFlow::Break(()); + }; + self.ctes_in_scope.truncate(scope.scope_start); + if let Some(cte_name) = scope.add_to_parent_after { + self.ctes_in_scope.push(cte_name); + } + ControlFlow::Continue(()) + } + + fn pre_visit_relation(&mut self, relation: &ObjectName) -> ControlFlow { + let is_cte = matches!( + relation.0.as_slice(), + [part] + if part.as_ident().is_some_and(|ident| { + self.ctes_in_scope.contains( + &ParserContext::canonicalize_identifier(ident.clone()).value, + ) + }) + ); + if !is_cte { + self.names.insert(relation.clone()); + } + ControlFlow::Continue(()) + } +} + /// Helper function for [extract_tables_from_query]. /// /// Handle [sqlparser::ast::Query]. fn extract_tables_from_sql_query(query: &sqlparser::ast::Query, names: &mut HashSet) { - let mut cte_names = HashSet::new(); - if let Some(with) = &query.with { - for cte in &with.cte_tables { - let cte_name = ParserContext::canonicalize_identifier(cte.alias.name.clone()).value; - let mut cte_query_names = HashSet::new(); - extract_tables_from_sql_query(&cte.query, &mut cte_query_names); - if with.recursive { - cte_names.insert(cte_name.clone()); - } - remove_cte_names(&mut cte_query_names, &cte_names); - names.extend(cte_query_names); - if !with.recursive { - cte_names.insert(cte_name); - } - } - } - - let mut body_names = HashSet::new(); - extract_tables_from_set_expr(&query.body, &mut body_names); - remove_cte_names(&mut body_names, &cte_names); - names.extend(body_names); -} - -/// Helper function for [extract_tables_from_query]. -/// -/// Handle [SetExpr]. -fn extract_tables_from_set_expr(set_expr: &SetExpr, names: &mut HashSet) { - match set_expr { - SetExpr::Select(select) => { - for from in &select.from { - extract_tables_from_table_with_joins(from, names); - } - } - SetExpr::Query(query) => { - extract_tables_from_sql_query(query, names); - } - SetExpr::SetOperation { left, right, .. } => { - extract_tables_from_set_expr(left, names); - extract_tables_from_set_expr(right, names); - } - _ => {} - }; -} - -/// Helper function for [extract_tables_from_query]. -/// -/// Handle [TableWithJoins]. -fn extract_tables_from_table_with_joins( - table_with_joins: &TableWithJoins, - names: &mut HashSet, -) { - table_factor_to_object_name(&table_with_joins.relation, names); - for join in &table_with_joins.joins { - table_factor_to_object_name(&join.relation, names); - } -} - -/// Helper function for [extract_tables_from_query]. -/// -/// Handle [TableFactor]. -fn table_factor_to_object_name(table_factor: &TableFactor, names: &mut HashSet) { - match table_factor { - TableFactor::Table { name, .. } => { - names.insert(name.to_owned()); - } - TableFactor::Derived { subquery, .. } => { - extract_tables_from_sql_query(subquery, names); - } - TableFactor::NestedJoin { - table_with_joins, .. - } => { - extract_tables_from_table_with_joins(table_with_joins, names); - } - TableFactor::Pivot { table, .. } - | TableFactor::Unpivot { table, .. } - | TableFactor::MatchRecognize { table, .. } => { - table_factor_to_object_name(table, names); - } - TableFactor::TableFunction { .. } - | TableFactor::Function { .. } - | TableFactor::UNNEST { .. } - | TableFactor::JsonTable { .. } - | TableFactor::OpenJsonTable { .. } - | TableFactor::XmlTable { .. } - | TableFactor::SemanticView { .. } => {} - } + let _ = query.visit(&mut RelationCollector::new(names)); } #[cfg(test)] @@ -518,6 +557,65 @@ TQL EVAL (now() - '15s'::interval, now(), '5s') count_values("status_code", {__n } } + #[test] + fn test_extract_tables_from_tql_query_completeness() { + let testcases = [ + ("cpu", true, vec!["cpu"]), + (r#"{__name__="cpu"}"#, true, vec!["cpu"]), + (r#"cpu + {job="api"}"#, false, vec!["cpu"]), + (r#"{__name__=~"cpu.*"}"#, false, vec![]), + (r#"cpu{__schema__=~"private.*"}"#, false, vec![]), + ]; + + for (promql, expected_complete, expected_tables) in testcases { + let sql = format!("TQL EVAL (0, 10, '5s') {promql}"); + let mut stmts = ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + let Statement::Tql(tql) = stmts.pop().unwrap() else { + unreachable!() + }; + + let query = SqlOrTql::Tql(tql, sql); + let (tables, complete) = extract_tables_from_query_inner(&query); + let mut tables = tables + .into_iter() + .map(|table| format_raw_object_name(&table)) + .collect_vec(); + tables.sort(); + assert_eq!(expected_complete, complete, "{promql}"); + assert_eq!(expected_tables, tables, "{promql}"); + assert_eq!( + expected_complete, + extract_tables_from_query_checked(&query).is_some(), + "{promql}" + ); + } + } + + #[test] + fn test_extract_tables_from_chained_tql_ctes() { + let sql = "WITH denied AS (TQL EVAL (0, 10, '5s') allowed), \ + leak AS (TQL EVAL (0, 10, '5s') denied) \ + SELECT * FROM leak"; + let mut stmts = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap(); + let Statement::Query(query) = stmts.pop().unwrap() else { + unreachable!() + }; + + let mut tables = extract_tables_from_query_checked(&SqlOrTql::Sql(*query, sql.to_string())) + .unwrap() + .map(|table| format_raw_object_name(&table)) + .collect_vec(); + tables.sort(); + assert_eq!(vec!["allowed".to_string(), "denied".to_string()], tables); + } + #[test] fn test_extract_tables_from_sql_query_with_derived_join() { let sql = r#" @@ -556,6 +654,37 @@ LEFT JOIN ( ); } + #[test] + fn test_extract_tables_from_sql_query_with_expression_subqueries() { + let sql = r#" +SELECT + (SELECT max(value) FROM scalar_source) +FROM outer_source +WHERE EXISTS (SELECT 1 FROM exists_source) +ORDER BY (SELECT max(value) FROM order_source); +"#; + let mut stmts = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap(); + let Statement::Query(query) = stmts.pop().unwrap() else { + unreachable!() + }; + + let mut tables = extract_tables_from_query(&SqlOrTql::Sql(*query, sql.to_string())) + .map(|table| format_raw_object_name(&table)) + .collect_vec(); + tables.sort(); + assert_eq!( + vec![ + "exists_source".to_string(), + "order_source".to_string(), + "outer_source".to_string(), + "scalar_source".to_string(), + ], + tables + ); + } + #[test] fn test_extract_tables_from_sql_query_with_cte_scopes() { let testcases = vec![ @@ -570,6 +699,15 @@ SELECT * FROM source; ), ( r#" +WITH RECURSIVE source AS ( + SELECT * FROM source +) +SELECT * FROM source; +"#, + vec![], + ), + ( + r#" WITH first_cte AS ( SELECT * FROM physical_source ), second_cte AS ( @@ -579,6 +717,24 @@ SELECT * FROM second_cte; "#, vec!["physical_source".to_string()], ), + ( + r#" +SELECT * FROM ( + WITH nested_cte AS (SELECT * FROM nested_source) + SELECT * FROM nested_cte +); +"#, + vec!["nested_source".to_string()], + ), + ( + r#" +SELECT * FROM ( + WITH nested_cte AS (SELECT * FROM nested_source) + SELECT (SELECT * FROM nested_cte) +); +"#, + vec!["nested_source".to_string()], + ), ]; for (sql, expected_tables) in testcases { @@ -603,6 +759,40 @@ SELECT * FROM second_cte; } } + #[test] + fn test_extract_tables_from_deeply_nested_ctes_visits_each_query_once() { + let mut sql = "SELECT * FROM physical_source".to_string(); + for depth in 0..20 { + sql = format!("WITH cte_{depth} AS ({sql}) SELECT * FROM cte_{depth}"); + } + + let mut stmts = ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + let Statement::Query(query) = stmts.pop().unwrap() else { + unreachable!() + }; + + let mut tables = HashSet::new(); + let query_visits = { + let mut collector = RelationCollector::new(&mut tables); + let _ = query.inner.visit(&mut collector); + collector.query_visits + }; + + assert_eq!(21, query_visits); + assert_eq!( + vec!["physical_source"], + tables + .into_iter() + .map(|table| format_raw_object_name(&table)) + .collect_vec() + ); + } + #[test] fn test_extract_tables_from_tql_query_with_schema_matcher() { let sql = r#" diff --git a/src/store-api/src/data_source.rs b/src/store-api/src/data_source.rs index 3fbef8e08c8..2e047cad6c0 100644 --- a/src/store-api/src/data_source.rs +++ b/src/store-api/src/data_source.rs @@ -16,6 +16,7 @@ use std::sync::Arc; use common_error::ext::BoxedError; use common_recordbatch::SendableRecordBatchStream; +use datafusion_physical_plan::ExecutionPlan; use crate::storage::ScanRequest; @@ -24,6 +25,15 @@ use crate::storage::ScanRequest; pub trait DataSource { /// Retrieves a stream of record batches based on the provided scan request. fn get_stream(&self, request: ScanRequest) -> Result; + + /// Builds a physical plan for the scan request if the data source needs to + /// expose its execution semantics to DataFusion. + fn get_physical_plan( + &self, + _request: ScanRequest, + ) -> Result>, BoxedError> { + Ok(None) + } } pub type DataSourceRef = Arc; diff --git a/src/store-api/src/logstore.rs b/src/store-api/src/logstore.rs index 19bbe195f0b..a2c1e3ee803 100644 --- a/src/store-api/src/logstore.rs +++ b/src/store-api/src/logstore.rs @@ -86,6 +86,14 @@ pub trait LogStore: Send + Sync + 'static + std::fmt::Debug { entry_id: EntryId, ) -> Result<(), Self::Error>; + /// Marks all entries of a region as obsolete and removes its dedicated namespace when + /// supported by the backend. + async fn obsolete_all( + &self, + provider: &Provider, + region_id: RegionId, + ) -> Result<(), Self::Error>; + /// Makes an entry instance of the associated Entry type fn entry( &self, diff --git a/src/store-api/src/mito_engine_options.rs b/src/store-api/src/mito_engine_options.rs index 8762ed69779..4320a40d8de 100644 --- a/src/store-api/src/mito_engine_options.rs +++ b/src/store-api/src/mito_engine_options.rs @@ -19,6 +19,8 @@ pub use common_wal::options::WAL_OPTIONS_KEY; /// Option key for append mode. pub const APPEND_MODE_KEY: &str = "append_mode"; +/// Option key for the per-region write buffer size. +pub const WRITE_BUFFER_SIZE_KEY: &str = "write_buffer_size"; /// Option key for merge mode. pub const MERGE_MODE_KEY: &str = "merge_mode"; /// Option key for TTL(time-to-live) @@ -72,6 +74,7 @@ pub const SST_FORMAT_KEY: &str = "sst_format"; pub fn is_mito_engine_option_key(key: &str) -> bool { [ "ttl", + WRITE_BUFFER_SIZE_KEY, AUTO_FLUSH_INTERVAL_KEY, COMPACTION_TYPE, COMPACTION_OVERRIDE, @@ -107,6 +110,7 @@ mod tests { #[test] fn test_is_mito_engine_option_key() { assert!(is_mito_engine_option_key("ttl")); + assert!(is_mito_engine_option_key("write_buffer_size")); assert!(is_mito_engine_option_key("auto_flush_interval")); assert!(is_mito_engine_option_key("compaction.type")); assert!(is_mito_engine_option_key("compaction.override")); diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index e046e9df568..d63c9058a33 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -23,16 +23,17 @@ use api::v1::column_def::{ }; use api::v1::region::bulk_insert_request::Body; use api::v1::region::{ - AlterRequest, AlterRequests, BuildIndexRequest, BulkInsertRequest, CloseRequest, - CompactRequest, CreateRequest, CreateRequests, DeleteRequests, DropRequest, DropRequests, - FlushRequest, InsertRequests, OpenRequest, TruncateRequest, alter_request, compact_request, - region_request, truncate_request, + AlterRequest, AlterRequests, BuildIndexRequest, BulkInsertRequest, + CleanUpRequest as PbCleanUpRequest, CloseRequest, CompactRequest, CreateRequest, + CreateRequests, DeleteRequests, DropRequest, DropRequests, FlushRequest, InsertRequests, + OpenRequest, TruncateRequest, alter_request, compact_request, region_request, truncate_request, }; use api::v1::{ self, Analyzer, ArrowIpc, FulltextBackend as PbFulltextBackend, Option as PbOption, Rows, SemanticType, SkippingIndexType as PbSkippingIndexType, WriteHint, }; pub use common_base::AffectedRows; +use common_base::readable_size::ReadableSize; use common_grpc::flight::FlightDecoder; use common_recordbatch::DfRecordBatch; use common_time::{TimeToLive, Timestamp}; @@ -54,7 +55,7 @@ use crate::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY; use crate::metrics; use crate::mito_engine_options::{ APPEND_MODE_KEY, AUTO_FLUSH_INTERVAL_KEY, SST_FORMAT_KEY, TTL_KEY, TWCS_MAX_OUTPUT_FILE_SIZE, - TWCS_TIME_WINDOW, TWCS_TRIGGER_FILE_NUM, + TWCS_TIME_WINDOW, TWCS_TRIGGER_FILE_NUM, WRITE_BUFFER_SIZE_KEY, }; use crate::path_utils::table_dir; use crate::storage::{ColumnId, RegionId, ScanRequest}; @@ -145,6 +146,7 @@ pub enum RegionRequest { Create(RegionCreateRequest), Drop(RegionDropRequest), Open(RegionOpenRequest), + CleanUp(RegionCleanUpRequest), Close(RegionCloseRequest), Alter(RegionAlterRequest), Flush(RegionFlushRequest), @@ -167,6 +169,7 @@ impl RegionRequest { region_request::Body::Create(create) => make_region_create(create), region_request::Body::Drop(drop) => make_region_drop(drop), region_request::Body::Open(open) => make_region_open(open), + region_request::Body::CleanUp(clean_up) => make_region_clean_up(clean_up), region_request::Body::Close(close) => make_region_close(close), region_request::Body::Alter(alter) => make_region_alter(alter), region_request::Body::Flush(flush) => make_region_flush(flush), @@ -325,11 +328,27 @@ fn make_region_open(open: OpenRequest) -> Result> )]) } +fn make_region_clean_up(clean_up: PbCleanUpRequest) -> Result> { + let region_id = RegionId::from(clean_up.region_id); + let table_dir = table_dir(&clean_up.path, region_id.table_id()); + Ok(vec![( + region_id, + RegionRequest::CleanUp(RegionCleanUpRequest { + engine: clean_up.engine, + table_dir, + path_type: PathType::Bare, + options: clean_up.options, + }), + )]) +} + fn make_region_close(close: CloseRequest) -> Result> { let region_id = close.region_id.into(); Ok(vec![( region_id, - RegionRequest::Close(RegionCloseRequest {}), + RegionRequest::Close(RegionCloseRequest { + flush_on_close: close.flush_on_close, + }), )]) } @@ -646,9 +665,32 @@ impl RegionOpenRequest { } } +/// Offline region cleanup request. +#[derive(Debug, Clone)] +pub struct RegionCleanUpRequest { + /// Region engine name + pub engine: String, + /// Directory for table's data home. Usually is composed by catalog and table id + pub table_dir: String, + /// Path type for generating paths + pub path_type: PathType, + /// Options of the cleaned region. + pub options: HashMap, +} + +impl RegionCleanUpRequest { + /// Returns true when the region belongs to the metric engine's physical table. + pub fn is_physical_table(&self) -> bool { + self.options.contains_key(PHYSICAL_TABLE_METADATA_KEY) + } +} + /// Close region request. -#[derive(Debug)] -pub struct RegionCloseRequest {} +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct RegionCloseRequest { + /// Whether to flush the region before closing it. + pub flush_on_close: bool, +} /// Alter metadata of a region. #[derive(Debug, PartialEq, Eq, Clone)] @@ -1364,6 +1406,7 @@ impl From for ModifyColumnType { #[derive(Debug, Eq, PartialEq, Clone, Serialize, Deserialize)] pub enum SetRegionOption { + WriteBufferSize(Option), Ttl(Option), // Modifying TwscOptions with values as (option name, new value). Twsc(String, String), @@ -1381,6 +1424,12 @@ impl TryFrom<&PbOption> for SetRegionOption { fn try_from(value: &PbOption) -> std::result::Result { let PbOption { key, value } = value; match key.as_str() { + WRITE_BUFFER_SIZE_KEY => { + let size = value + .parse::() + .map_err(|_| InvalidSetRegionOptionRequestSnafu { key, value }.build())?; + Ok(Self::WriteBufferSize(Some(size))) + } TTL_KEY => { let ttl = TimeToLive::from_humantime_or_str(value) .map_err(|_| InvalidSetRegionOptionRequestSnafu { key, value }.build())?; @@ -1429,6 +1478,7 @@ impl From<&UnsetRegionOption> for SetRegionOption { SetRegionOption::Twsc(unset_option.to_string(), String::new()) } UnsetRegionOption::Ttl => SetRegionOption::Ttl(Default::default()), + UnsetRegionOption::WriteBufferSize => SetRegionOption::WriteBufferSize(None), } } } @@ -1439,6 +1489,7 @@ impl TryFrom<&str> for UnsetRegionOption { fn try_from(key: &str) -> Result { match key.to_ascii_lowercase().as_str() { TTL_KEY => Ok(Self::Ttl), + WRITE_BUFFER_SIZE_KEY => Ok(Self::WriteBufferSize), TWCS_TRIGGER_FILE_NUM => Ok(Self::TwcsTriggerFileNum), TWCS_MAX_OUTPUT_FILE_SIZE => Ok(Self::TwcsMaxOutputFileSize), TWCS_TIME_WINDOW => Ok(Self::TwcsTimeWindow), @@ -1453,12 +1504,14 @@ pub enum UnsetRegionOption { TwcsMaxOutputFileSize, TwcsTimeWindow, Ttl, + WriteBufferSize, } impl UnsetRegionOption { pub fn as_str(&self) -> &str { match self { Self::Ttl => TTL_KEY, + Self::WriteBufferSize => WRITE_BUFFER_SIZE_KEY, Self::TwcsTriggerFileNum => TWCS_TRIGGER_FILE_NUM, Self::TwcsMaxOutputFileSize => TWCS_MAX_OUTPUT_FILE_SIZE, Self::TwcsTimeWindow => TWCS_TIME_WINDOW, @@ -1633,6 +1686,7 @@ impl fmt::Display for RegionRequest { RegionRequest::Create(_) => write!(f, "Create"), RegionRequest::Drop(_) => write!(f, "Drop"), RegionRequest::Open(_) => write!(f, "Open"), + RegionRequest::CleanUp(_) => write!(f, "CleanUp"), RegionRequest::Close(_) => write!(f, "Close"), RegionRequest::Alter(_) => write!(f, "Alter"), RegionRequest::Flush(_) => write!(f, "Flush"), @@ -1789,6 +1843,41 @@ mod tests { ); } + #[test] + fn test_write_buffer_size_region_options() { + let option = PbOption { + key: WRITE_BUFFER_SIZE_KEY.to_string(), + value: "128MiB".to_string(), + }; + assert_eq!( + SetRegionOption::WriteBufferSize(Some(ReadableSize::mb(128))), + SetRegionOption::try_from(&option).unwrap() + ); + + let option = PbOption { + key: WRITE_BUFFER_SIZE_KEY.to_string(), + value: "invalid".to_string(), + }; + SetRegionOption::try_from(&option).unwrap_err(); + + // Clearing the option uses UNSET. Empty SET values, including SQL NULL, + // remain invalid because the SQL layer does not distinguish NULL from ''. + let option = PbOption { + key: WRITE_BUFFER_SIZE_KEY.to_string(), + value: String::new(), + }; + SetRegionOption::try_from(&option).unwrap_err(); + + assert_eq!( + UnsetRegionOption::WriteBufferSize, + UnsetRegionOption::try_from(WRITE_BUFFER_SIZE_KEY).unwrap() + ); + assert_eq!( + SetRegionOption::WriteBufferSize(None), + SetRegionOption::from(&UnsetRegionOption::WriteBufferSize) + ); + } + /// Returns a new region metadata for testing. Metadata: /// `[(ts, ms, 1), (tag_0, string, 2), (field_0, string, 3), (field_1, bool, 4)]` fn new_metadata() -> RegionMetadata { @@ -2194,6 +2283,28 @@ mod tests { assert_eq!(request.requirements, RegionRequirements::object_storage()); } + #[test] + fn test_parse_region_cleanup_from_proto() { + let clean_up = api::v1::region::CleanUpRequest { + region_id: RegionId::new(42, 3).as_u64(), + engine: "mito".to_string(), + path: "test".to_string(), + options: HashMap::from([("k".to_string(), "v".to_string())]), + }; + + let requests = + RegionRequest::try_from_request_body(region_request::Body::CleanUp(clean_up)).unwrap(); + let RegionRequest::CleanUp(request) = &requests[0].1 else { + unreachable!() + }; + + assert_eq!(requests[0].0, RegionId::new(42, 3)); + assert_eq!(request.engine, "mito"); + assert_eq!(request.table_dir, "data/test/42/"); + assert_eq!(request.path_type, PathType::Bare); + assert_eq!(request.options.get("k"), Some(&"v".to_string())); + } + #[test] fn test_validate_modify_column_fulltext_options() { let kind = AlterKind::SetIndexes { diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs index 006cfb34962..b2801dd8654 100644 --- a/src/table/src/metadata.rs +++ b/src/table/src/metadata.rs @@ -350,6 +350,9 @@ impl TableMeta { for request in requests { match request { + SetRegionOption::WriteBufferSize(new_write_buffer_size) => { + new_options.write_buffer_size = *new_write_buffer_size; + } SetRegionOption::Ttl(new_ttl) => { new_options.ttl = *new_ttl; } diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index 12214a92f9d..12331bf8bb9 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -145,7 +145,7 @@ pub fn validate_table_option(key: &str) -> bool { #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] #[serde(default)] pub struct TableOptions { - /// Memtable size of memtable. + /// Per-region write buffer stall threshold. Writes are rejected at twice this size. pub write_buffer_size: Option, /// Time-to-live of table. Expired data will be automatically purged. pub ttl: Option, @@ -155,7 +155,7 @@ pub struct TableOptions { pub extra_options: HashMap, } -pub const WRITE_BUFFER_SIZE_KEY: &str = "write_buffer_size"; +pub const WRITE_BUFFER_SIZE_KEY: &str = store_api::mito_engine_options::WRITE_BUFFER_SIZE_KEY; pub const TTL_KEY: &str = store_api::mito_engine_options::TTL_KEY; pub const STORAGE_KEY: &str = "storage"; pub const COMMENT_KEY: &str = "comment"; diff --git a/src/table/src/table.rs b/src/table/src/table.rs index 0ae7d580d8a..542a5066771 100644 --- a/src/table/src/table.rs +++ b/src/table/src/table.rs @@ -19,6 +19,7 @@ use common_recordbatch::SendableRecordBatchStream; use datafusion::execution::FunctionRegistry; use datafusion::logical_expr::Cast; use datafusion::logical_expr::expr::ScalarFunction; +use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionContext; use datafusion_expr::expr::Expr; use datatypes::data_type::DataType; @@ -142,6 +143,12 @@ impl Table { .context(TablesRecordBatchSnafu) } + pub fn scan_to_plan(&self, request: ScanRequest) -> Result>> { + self.data_source + .get_physical_plan(request) + .context(TablesRecordBatchSnafu) + } + /// Tests whether the table provider can make use of any or all filter expressions /// to optimise data retrieval. pub fn supports_filters_pushdown(&self, filters: &[&Expr]) -> Result> { diff --git a/src/table/src/table/adapter.rs b/src/table/src/table/adapter.rs index 339c82dddeb..8907140c312 100644 --- a/src/table/src/table/adapter.rs +++ b/src/table/src/table/adapter.rs @@ -113,6 +113,11 @@ impl TableProvider for DfTableProviderAdapter { request.limit = limit; request.clone() }; + + if let Some(plan) = self.table.scan_to_plan(request.clone())? { + return Ok(plan); + } + let stream = self.table.scan_to_stream(request).await?; // build sort physical expr diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 825d930a6fc..9ce752553e5 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -541,7 +541,7 @@ impl Stream for StreamWithMetricWrapper { Ok(record_batch) => { // we don't record elapsed time here // since it's calling storage api involving I/O ops - let batch_bytes = record_batch.buffer_memory_size(); + let batch_bytes = record_batch.logical_slice_memory_size(); this.metric.record_output_bytes(batch_bytes); this.metric.record_output(record_batch.num_rows()); Poll::Ready(Some(Ok(record_batch.into_df_record_batch()))) diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 122318d45f1..37ca78f66f3 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -251,6 +251,7 @@ impl GreptimeDbStandaloneBuilder { flow_metadata_allocator, region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl), soft_drop_enabled: false, + soft_drop_retention: None, }, procedure_manager.clone(), repartition_procedure_factory, diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index fa9bad0fac0..90c3708b534 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -28,6 +28,7 @@ use datatypes::arrow::array::{ use frontend::error::Error; use frontend::instance::Instance; use operator::error::Error as OperatorError; +use query::datafusion::QUERY_PARALLELISM_HINT; use rstest::rstest; use rstest_reuse::apply; use servers::error as server_error; @@ -89,6 +90,98 @@ async fn test_create_database_and_insert_query(instance: Arc) } } +#[tokio::test(flavor = "multi_thread")] +async fn test_distributed_scalar_latest_with_query_parallelism_below_regions() { + common_telemetry::init_default_ut_logging(); + + let distributed = crate::tests::create_distributed_instance( + "test_distributed_scalar_latest_with_query_parallelism_below_regions", + ) + .await; + let frontend = distributed.frontend(); + + execute_sql( + &frontend, + r#" +CREATE TABLE cpu ( + rack STRING NULL, + os STRING NULL, + usage_user BIGINT NULL, + greptime_timestamp TIMESTAMP(9) NOT NULL, + TIME INDEX (greptime_timestamp) +) +PARTITION ON COLUMNS (rack) ( + rack < '2', + rack >= '2' AND rack < '4', + rack >= '4' AND rack < '6', + rack >= '6' AND rack < '8', + rack >= '8' +) +ENGINE = mito +WITH (append_mode = 'true', sst_format = 'flat') +"#, + ) + .await; + + execute_sql( + &frontend, + r#" +INSERT INTO cpu VALUES + ('1', 'linux', 10, '2023-06-12 01:04:49'), + ('1', 'linux', 15, '2023-06-12 01:04:50'), + ('3', 'windows', 25, '2023-06-12 01:05:00'), + ('5', 'mac', 30, '2023-06-12 01:03:00'), + ('7', 'linux', 45, '2023-06-12 02:00:00'), + ('2', 'linux', 20, '2023-06-12 01:04:51'), + ('2', 'windows', 22, '2023-06-12 01:06:00'), + ('4', 'mac', 12, '2023-06-12 00:59:00'), + ('6', 'linux', 35, '2023-06-12 01:04:55'), + ('8', 'windows', 50, '2023-06-12 02:10:00') +"#, + ) + .await; + + let latest_sql = r#" +SELECT rack, os, greptime_timestamp +FROM cpu +WHERE greptime_timestamp = ( + SELECT greptime_timestamp + FROM cpu + ORDER BY greptime_timestamp DESC + LIMIT 1 +) +"#; + + let result = execute_sql_with_query_parallelism(&frontend, latest_sql, 1) + .await + .data + .pretty_print() + .await; + assert_eq!( + result, + r#"+------+---------+---------------------+ +| rack | os | greptime_timestamp | ++------+---------+---------------------+ +| 8 | windows | 2023-06-12T02:10:00 | ++------+---------+---------------------+"# + ); + + let explain = + execute_sql_with_query_parallelism(&frontend, &format!("EXPLAIN {latest_sql}"), 1) + .await + .data + .pretty_print() + .await; + assert!( + explain.contains("SortExec"), + "query_parallelism=1 with five regions should insert SortExec below MergeSortExec; explain:\n{explain}" + ); + assert!( + explain.contains("MergeSortExec"), + "query_parallelism=1 with five regions should keep the distributed merge stage stable; explain:\n{explain}" + ); +} + #[apply(both_instances_cases)] async fn test_show_create_table(instance: Arc) { let frontend = instance.frontend(); @@ -1090,6 +1183,138 @@ async fn test_execute_copy_from_headerless_csv(instance: Arc) check_output_stream(output, expect).await; } +#[apply(both_instances_cases)] +async fn test_execute_copy_from_csv_strict_headers(instance: Arc) { + let instance = instance.frontend(); + let tmp_dir = temp_dir::create_temp_dir("test_execute_copy_from_csv_strict_headers"); + let matching_path = tmp_dir.path().join("matching.csv"); + let unknown_path = tmp_dir.path().join("unknown.csv"); + let missing_path = tmp_dir.path().join("missing.csv"); + let duplicate_path = tmp_dir.path().join("duplicate.csv"); + let matching_location = prepare_path(matching_path.to_str().unwrap()); + let unknown_location = prepare_path(unknown_path.to_str().unwrap()); + let missing_location = prepare_path(missing_path.to_str().unwrap()); + let duplicate_location = prepare_path(duplicate_path.to_str().unwrap()); + std::fs::write( + &matching_path, + "ts,host_id,reading_value\n2024-01-01T00:00:00,1,10.5\n", + ) + .unwrap(); + std::fs::write( + &unknown_path, + "host_id,reading_value,ts,extra\n1,10.5,2024-01-01T00:00:00,ignored\n", + ) + .unwrap(); + std::fs::write(&missing_path, "host_id,ts\n1,2024-01-01T00:00:00\n").unwrap(); + std::fs::write( + &duplicate_path, + "host_id,reading_value,ts,host_id\n1,10.5,2024-01-01T00:00:00,2\n", + ) + .unwrap(); + + let output = execute_sql( + &instance, + "CREATE TABLE csv_strict_headers(host_id INT, reading_value DOUBLE, ts TIMESTAMP TIME INDEX);", + ) + .await + .data; + assert!(matches!(output, OutputData::AffectedRows(0))); + + let output = execute_sql( + &instance, + &format!( + "COPY csv_strict_headers FROM '{}' WITH (FORMAT='csv', STRICT_HEADERS='true');", + matching_location + ), + ) + .await + .data; + assert!(matches!(output, OutputData::AffectedRows(1))); + + let output = execute_sql(&instance, "SELECT * FROM csv_strict_headers;") + .await + .data; + let expect = "\ ++---------+---------------+---------------------+ +| host_id | reading_value | ts | ++---------+---------------+---------------------+ +| 1 | 10.5 | 2024-01-01T00:00:00 | ++---------+---------------+---------------------+"; + check_output_stream(output, expect).await; + + let err = try_execute_sql( + &instance, + &format!( + "COPY csv_strict_headers FROM '{}' WITH (FORMAT='csv', STRICT_HEADERS='true');", + unknown_location + ), + ) + .await + .unwrap_err(); + assert!(matches!( + unwrap_frontend_error(&err), + Error::TableOperation { + source: OperatorError::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + }, + .. + } if unknown_columns == &vec!["extra".to_string()] + && missing_columns.is_empty() + && duplicate_columns.is_empty() + )); + + let err = try_execute_sql( + &instance, + &format!( + "COPY csv_strict_headers FROM '{}' WITH (FORMAT='csv', STRICT_HEADERS='true');", + missing_location + ), + ) + .await + .unwrap_err(); + assert!(matches!( + unwrap_frontend_error(&err), + Error::TableOperation { + source: OperatorError::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + }, + .. + } if unknown_columns.is_empty() + && missing_columns == &vec!["reading_value".to_string()] + && duplicate_columns.is_empty() + )); + + let err = try_execute_sql( + &instance, + &format!( + "COPY csv_strict_headers FROM '{}' WITH (FORMAT='csv', STRICT_HEADERS='true');", + duplicate_location + ), + ) + .await + .unwrap_err(); + assert!(matches!( + unwrap_frontend_error(&err), + Error::TableOperation { + source: OperatorError::CsvHeaderMismatch { + unknown_columns, + missing_columns, + duplicate_columns, + .. + }, + .. + } if unknown_columns.is_empty() + && missing_columns.is_empty() + && duplicate_columns == &vec!["host_id".to_string()] + )); +} + #[apply(both_instances_cases)] async fn test_execute_query_external_table_json(instance: Arc) { unsafe { @@ -2744,6 +2969,16 @@ async fn execute_sql(instance: &Arc, sql: &str) -> Output { execute_sql_with(instance, sql, QueryContext::arc()).await } +async fn execute_sql_with_query_parallelism( + instance: &Arc, + sql: &str, + parallelism: usize, +) -> Output { + let mut query_ctx = QueryContext::with_db_name(None); + query_ctx.set_extension(QUERY_PARALLELISM_HINT, parallelism.to_string()); + execute_sql_with(instance, sql, Arc::new(query_ctx)).await +} + async fn try_execute_sql(instance: &Arc, sql: &str) -> server_error::Result { try_execute_sql_with(instance, sql, QueryContext::arc()).await } diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index c6366f48efe..5f4e5886d8a 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -44,7 +44,7 @@ use common_memory_manager::OnExhaustedPolicy; use common_options::plugin_options::StandaloneFlag; use flate2::Compression; use flate2::write::GzEncoder; -use log_query::{Context, Limit, LogQuery, TimeFilter}; +use log_query::{AggFunc, Context, Limit, LogExpr, LogQuery, TimeFilter}; use loki_proto::logproto::{EntryAdapter, LabelPairAdapter, PushRequest, StreamAdapter}; use loki_proto::prost_types::Timestamp; use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; @@ -1939,6 +1939,7 @@ compress_manifest = false experimental_compaction_memory_limit = "unlimited" experimental_compaction_on_exhausted = "wait" auto_flush_interval = "30m" +default_region_write_buffer_size = "0KiB" enable_write_cache = false write_cache_path = "" write_cache_size = "5GiB" @@ -6996,10 +6997,8 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { "unexpected error body: {body:?}" ); assert!( - body.message.contains("Chunk fallback triggered by") - || body - .message - .contains("Discarded 2 spans after ambiguous chunk failure"), + body.message + .contains("Discarded 2 spans after pre-write request failure"), "unexpected error body: {body:?}" ); @@ -7810,14 +7809,16 @@ pub async fn test_log_query(store_type: StorageType) { .await; assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); let res = client - .post("/v1/sql?sql=insert into logs values ('2024-11-07 10:53:50', 'hello');") + .post( + "/v1/sql?sql=insert into logs values ('2024-11-06 23:59:59', 'before-date-end'), ('2024-11-07 10:53:50', 'before-explicit-end'), ('2024-11-07 10:53:51', 'at-explicit-end'), ('2024-11-07 10:53:52', 'after-explicit-end');", + ) .header("Content-Type", "application/x-www-form-urlencoded") .send() .await; assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); // test log query - let log_query = LogQuery { + let mut log_query = LogQuery { table: TableName { catalog_name: "greptime".to_string(), schema_name: "public".to_string(), @@ -7825,12 +7826,12 @@ pub async fn test_log_query(store_type: StorageType) { }, time_filter: TimeFilter { start: Some("2024-11-07".to_string()), - end: None, + end: Some("2024-11-07T10:53:51Z".to_string()), span: None, }, limit: Limit { skip: None, - fetch: Some(1), + fetch: Some(3), }, columns: vec!["ts".to_string(), "message".to_string()], filters: Default::default(), @@ -7847,7 +7848,114 @@ pub async fn test_log_query(store_type: StorageType) { assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); let resp = res.text().await; let v = get_rows_from_output(&resp); - assert_eq!(v, "[[1730976830000,\"hello\"]]"); + assert_eq!(v, "[[1730976830000,\"before-explicit-end\"]]"); + + log_query.time_filter.start = Some("2024-11-06".to_string()); + log_query.time_filter.end = Some("2024-11-07".to_string()); + log_query.limit.fetch = Some(4); + let res = client + .post("/v1/logs") + .header("Content-Type", "application/json") + .body(serde_json::to_string(&log_query).unwrap()) + .send() + .await; + + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + let resp = res.text().await; + let output = serde_json::from_str::(&resp).unwrap(); + let rows = output["output"][0]["records"]["rows"].as_array().unwrap(); + let mut messages = rows + .iter() + .map(|row| row[1].as_str().unwrap()) + .collect::>(); + messages.sort_unstable(); + assert_eq!( + messages, + vec![ + "after-explicit-end", + "at-explicit-end", + "before-date-end", + "before-explicit-end", + ] + ); + + let res = client + .get("/v1/sql?sql=create table logs_limit (`ts` timestamp time index, `message` string);") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + let res = client + .get("/v1/sql?sql=insert into logs_limit select number, 'hello' from numbers limit 1001;") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + + let aggregate_query = LogQuery { + table: TableName::new("greptime", "public", "logs_limit"), + time_filter: TimeFilter { + start: Some("1970-01-01T00:00:00Z".to_string()), + end: Some("1970-01-01T00:00:02Z".to_string()), + span: None, + }, + limit: Limit { + skip: None, + fetch: None, + }, + columns: vec![], + filters: Default::default(), + context: Context::None, + exprs: vec![LogExpr::AggrFunc { + expr: vec![AggFunc::new( + "count".to_string(), + vec![LogExpr::NamedIdent("message".to_string())], + Some("count_result".to_string()), + )], + by: vec![], + }], + }; + let res = client + .post("/v1/logs") + .header("Content-Type", "application/json") + .body(serde_json::to_string(&aggregate_query).unwrap()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + assert_eq!(get_rows_from_output(&res.text().await), "[[1001]]"); + + let mut output_query = LogQuery { + table: TableName::new("greptime", "public", "logs_limit"), + time_filter: TimeFilter { + start: Some("1970-01-01T00:00:00Z".to_string()), + end: Some("1970-01-01T00:00:02Z".to_string()), + span: None, + }, + limit: Limit { + skip: None, + fetch: None, + }, + columns: vec!["ts".to_string(), "message".to_string()], + filters: Default::default(), + context: Context::None, + exprs: vec![], + }; + let res = client + .post("/v1/logs") + .header("Content-Type", "application/json") + .body(serde_json::to_string(&output_query).unwrap()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + assert_eq!(get_output_row_count(&res.text().await), 1000); + + output_query.limit.fetch = Some(17); + let res = client + .post("/v1/logs") + .header("Content-Type", "application/json") + .body(serde_json::to_string(&output_query).unwrap()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + assert_eq!(get_output_row_count(&res.text().await), 17); guard.remove_all().await; } @@ -9631,6 +9739,18 @@ fn get_rows_from_output(output: &str) -> String { .to_string() } +fn get_output_row_count(output: &str) -> usize { + let resp: Value = serde_json::from_str(output).unwrap(); + resp.get("output") + .and_then(Value::as_array) + .and_then(|v| v.first()) + .and_then(|v| v.get("records")) + .and_then(|v| v.get("rows")) + .and_then(Value::as_array) + .unwrap() + .len() +} + fn compress_vec_with_gzip(data: Vec) -> Vec { let mut encoder = GzEncoder::new(Vec::new(), Compression::default()); encoder.write_all(&data).unwrap(); diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index 32f80a32cba..cc55cb1bfdc 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -345,6 +345,17 @@ pub async fn test_partition_unpartitioned_mito(store_type: StorageType) { .with_datanodes(datanodes as u32) .with_store_config(store_config) .with_datanode_wal_config(DatanodeWalConfig::Noop) + .with_metasrv_gc_config(GcSchedulerOptions { + enable: true, + gc_cooldown_period: Duration::from_nanos(1), + ..Default::default() + }) + .with_datanode_gc_config(GcConfig { + enable: true, + lingering_time: Some(Duration::from_secs(0)), + unknown_file_lingering_time: Duration::from_secs(0), + ..Default::default() + }) .build(true) .await; @@ -472,6 +483,17 @@ pub async fn test_repartition_on_columns_metadata_mito(store_type: StorageType) .with_datanodes(datanodes as u32) .with_store_config(store_config) .with_datanode_wal_config(DatanodeWalConfig::Noop) + .with_metasrv_gc_config(GcSchedulerOptions { + enable: true, + gc_cooldown_period: Duration::from_nanos(1), + ..Default::default() + }) + .with_datanode_gc_config(GcConfig { + enable: true, + lingering_time: Some(Duration::from_secs(0)), + unknown_file_lingering_time: Duration::from_secs(0), + ..Default::default() + }) .build(true) .await; @@ -557,6 +579,17 @@ pub async fn test_repartition_on_columns_data_correctness_mito(store_type: Stora .with_datanodes(datanodes as u32) .with_store_config(store_config) .with_datanode_wal_config(DatanodeWalConfig::Noop) + .with_metasrv_gc_config(GcSchedulerOptions { + enable: true, + gc_cooldown_period: Duration::from_nanos(1), + ..Default::default() + }) + .with_datanode_gc_config(GcConfig { + enable: true, + lingering_time: Some(Duration::from_secs(0)), + unknown_file_lingering_time: Duration::from_secs(0), + ..Default::default() + }) .build(true) .await; @@ -612,6 +645,17 @@ pub async fn test_partition_unpartitioned_metric(store_type: StorageType) { .with_datanodes(datanodes as u32) .with_store_config(store_config) .with_datanode_wal_config(DatanodeWalConfig::Noop) + .with_metasrv_gc_config(GcSchedulerOptions { + enable: true, + gc_cooldown_period: Duration::from_nanos(1), + ..Default::default() + }) + .with_datanode_gc_config(GcConfig { + enable: true, + lingering_time: Some(Duration::from_secs(0)), + unknown_file_lingering_time: Duration::from_secs(0), + ..Default::default() + }) .build(true) .await; diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 156de7c415c..0e477061bb7 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashMap; +use std::str::FromStr; use std::time::Duration; use auth::user_provider_from_option; @@ -1353,6 +1354,68 @@ pub async fn test_postgres_uint64_parameter(store_type: StorageType) { let count: i64 = rows[0].get(0); assert_eq!(count, 1); + let scalar_statement = client + .prepare("SELECT arrow_cast($1, 'UInt64')") + .await + .unwrap(); + let scalar_null_statement = client + .prepare("SELECT arrow_cast($1, 'UInt64') IS NULL") + .await + .unwrap(); + let array_statement = client + .prepare("SELECT arrow_cast($1, 'List(UInt64)')") + .await + .unwrap(); + let array_null_statement = client + .prepare("SELECT arrow_cast($1, 'List(UInt64)') IS NULL") + .await + .unwrap(); + + let ordinary = Decimal::from(42); + let max = Decimal::from_str("18446744073709551615").unwrap(); + for value in [&ordinary, &max] { + let row = client.query_one(&scalar_statement, &[value]).await.unwrap(); + assert_eq!(*value, row.get::<_, Decimal>(0)); + } + + let scalar_null: Option = None; + let row = client + .query_one(&scalar_null_statement, &[&scalar_null]) + .await + .unwrap(); + assert!(row.get::<_, bool>(0)); + + let array_values = vec![Some(ordinary), None, Some(max)]; + let row = client + .query_one(&array_statement, &[&array_values]) + .await + .unwrap(); + assert_eq!(array_values, row.get::<_, Vec>>(0)); + + let array_null: Option>> = None; + let row = client + .query_one(&array_null_statement, &[&array_null]) + .await + .unwrap(); + assert!(row.get::<_, bool>(0)); + + for value in [ + Decimal::from(-1), + Decimal::from_str("18446744073709551616").unwrap(), + ] { + let error = client + .query(&scalar_statement, &[&value]) + .await + .unwrap_err(); + assert_pg_numeric_range_error(error); + + let error = client + .query(&array_statement, &[&vec![Some(value)]]) + .await + .unwrap_err(); + assert_pg_numeric_range_error(error); + } + drop(client); rx.await.unwrap(); @@ -1360,6 +1423,12 @@ pub async fn test_postgres_uint64_parameter(store_type: StorageType) { guard.remove_all().await; } +fn assert_pg_numeric_range_error(error: tokio_postgres::Error) { + let error = error.as_db_error().expect("expected PostgreSQL user error"); + assert_eq!("22023", error.code().code()); + assert_eq!("numeric_value_out_of_range", error.message()); +} + pub async fn test_postgres_explain_bind_parameter(store_type: StorageType) { // Regression test for #8029: EXPLAIN / EXPLAIN ANALYZE must accept bind // parameters over the Postgres extended query protocol. diff --git a/tests/cases/distributed/explain/multi_partitions.result b/tests/cases/distributed/explain/multi_partitions.result index 0c3d4fa4b70..47d807542d8 100644 --- a/tests/cases/distributed/explain/multi_partitions.result +++ b/tests/cases/distributed/explain/multi_partitions.result @@ -30,7 +30,7 @@ explain SELECT * FROM multi_partitions_test_table WHERE ts > cast(1000000000 as |_|_Filter: multi_partitions_test_table.ts > TimestampMillisecond(1000000000, None)_| |_|_TableScan: multi_partitions_test_table, partial_filters=[multi_partitions_test_table.ts > TimestampMillisecond(1000000000, None)]_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ diff --git a/tests/cases/distributed/explain/step_aggr_advance.result b/tests/cases/distributed/explain/step_aggr_advance.result index c2304df5eb9..7ea86dbccbf 100644 --- a/tests/cases/distributed/explain/step_aggr_advance.result +++ b/tests/cases/distributed/explain/step_aggr_advance.result @@ -37,7 +37,7 @@ tql explain (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr | | Filter: aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None) AND aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None) | | | TableScan: aggr_optimize_not, partial_filters=[aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None), aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None)] | | | ]] | -| physical_plan | SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] | +| physical_plan | MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] | | | MergeScanExec: REDACTED | | | +---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -54,7 +54,7 @@ tql analyze (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED @@ -244,7 +244,7 @@ tql explain (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time | | Filter: aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None) AND aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None) | | | TableScan: aggr_optimize_not, partial_filters=[aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None), aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None)] | | | ]] | -| physical_plan | SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] | +| physical_plan | MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] | | | MergeScanExec: REDACTED | | | +---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -261,7 +261,7 @@ tql analyze (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] REDACTED @@ -485,9 +485,10 @@ tql explain (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize | | PromSeriesDivideExec: tags=["a", "b", "c", "d"] | | | SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC, greptime_timestamp@4 ASC], preserve_partitioning=[true] | | | MergeScanExec: REDACTED -| | SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] | -| | CooperativeExec | -| | MergeScanExec: REDACTED +| | RepartitionExec: partitioning=REDACTED +| | MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] | +| | SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] | +| | MergeScanExec: REDACTED | | | +---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -514,8 +515,9 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED |_|_|_SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC, greptime_timestamp@4 ASC], preserve_partitioning=[true] REDACTED |_|_|_MergeScanExec: REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_CooperativeExec REDACTED @@ -974,7 +976,7 @@ EXPLAIN SELECT pk_col_1, pk_col_2, sum(val_col_1) FROM step_aggr_extended GROUP |_|_Aggregate: groupBy=[[step_aggr_extended.pk_col_1, step_aggr_extended.pk_col_2]], aggr=[[sum(step_aggr_extended.val_col_1)]] | |_|_TableScan: step_aggr_extended_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [pk_col_1@0 ASC NULLS LAST, pk_col_2@1 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [pk_col_1@0 ASC NULLS LAST, pk_col_2@1 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ diff --git a/tests/cases/distributed/explain/subqueries.result b/tests/cases/distributed/explain/subqueries.result index 49bf5d97566..e6f7286535e 100644 --- a/tests/cases/distributed/explain/subqueries.result +++ b/tests/cases/distributed/explain/subqueries.result @@ -287,7 +287,7 @@ EXPLAIN SELECT x FROM (SELECT a AS x FROM t) sq ORDER BY x; |_|_Projection: t.a AS x_| |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -326,7 +326,7 @@ EXPLAIN SELECT x, COUNT(*) AS c FROM (SELECT a AS x FROM t) sq GROUP BY x ORDER |_|_Projection: t.a AS x_| |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -397,7 +397,7 @@ EXPLAIN SELECT sq.x FROM (SELECT a AS x FROM t) sq ORDER BY sq.x; |_|_Projection: t.a AS x_| |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -433,7 +433,7 @@ EXPLAIN SELECT y FROM (SELECT x AS y FROM (SELECT a AS x FROM t) sq1) sq2 ORDER |_|_Projection: t.a AS x_| |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [y@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [y@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -471,7 +471,7 @@ EXPLAIN SELECT x, x + 1 AS y FROM (SELECT a AS x FROM t) sq ORDER BY x; |_|_Projection: t.a AS x_| |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -625,7 +625,7 @@ EXPLAIN SELECT x FROM (SELECT a AS x FROM t) sq ORDER BY x LIMIT 2; |_|_Projection: t.a AS x_| |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST], fetch=2_| +| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST], fetch=2_| |_|_SortExec: TopK(fetch=2), expr=[x@0 ASC NULLS LAST], preserve_partitioning=[true]_| |_|_MergeScanExec: REDACTED |_|_| diff --git a/tests/cases/distributed/information_schema/ssts_limit.result b/tests/cases/distributed/information_schema/ssts_limit.result index 98a19842995..0b64a3c0fa0 100644 --- a/tests/cases/distributed/information_schema/ssts_limit.result +++ b/tests/cases/distributed/information_schema/ssts_limit.result @@ -63,6 +63,37 @@ FROM ( | 1 | +-----------------------+ +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +EXPLAIN SELECT COUNT(*) AS filtered_limited_rows +FROM ( + SELECT region_id + FROM information_schema.ssts_manifest + WHERE table_id > 0 + LIMIT 1 +); + ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: count(Int64(1)) AS count(*) AS filtered_limited_rows | +| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] | +| | Limit: skip=0, fetch=1 | +| | Projection: information_schema.ssts_manifest.region_id | +| | Filter: information_schema.ssts_manifest.table_id > UInt32(0) | +| | TableScan: information_schema.ssts_manifest, partial_filters=[information_schema.ssts_manifest.table_id > UInt32(0)] | +| | ]] | +| physical_plan | ProjectionExec: expr=[count(Int64(1))@0 as filtered_limited_rows] | +| | AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] | +| | CoalescePartitionsExec | +| | AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] | +| | CoalescePartitionsExec: fetch=1 | +| | FilterExec: table_id@0 > 0, projection=[], fetch=1 | +| | RepartitionExec: REDACTED +| | DistributedInspectExec: kind=SstManifest, scan=ScanRequest { projection: ProjectionInput { projection: [2], nested_paths: [] }, filters: [table_id > UInt32(0)] }, schema=Schema { fields: [Field { name: "table_id", data_type: UInt32 }], metadata: {"greptime:version": "0"} } | +| | | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + DROP TABLE ssts_limit_case; Affected Rows: 0 diff --git a/tests/cases/distributed/information_schema/ssts_limit.sql b/tests/cases/distributed/information_schema/ssts_limit.sql index c02f31d63d7..d548bd283a5 100644 --- a/tests/cases/distributed/information_schema/ssts_limit.sql +++ b/tests/cases/distributed/information_schema/ssts_limit.sql @@ -35,4 +35,13 @@ FROM ( LIMIT 1 ); +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +EXPLAIN SELECT COUNT(*) AS filtered_limited_rows +FROM ( + SELECT region_id + FROM information_schema.ssts_manifest + WHERE table_id > 0 + LIMIT 1 +); + DROP TABLE ssts_limit_case; diff --git a/tests/cases/distributed/optimizer/first_value_advance.result b/tests/cases/distributed/optimizer/first_value_advance.result index ac977b6fdfb..08e1b8a4112 100644 --- a/tests/cases/distributed/optimizer/first_value_advance.result +++ b/tests/cases/distributed/optimizer/first_value_advance.result @@ -240,7 +240,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t.host]], aggr=[[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]]] | |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -267,7 +267,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST]@0 as ordered_host, first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST]@1 as first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]@2 as first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -606,7 +606,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t1.host]], aggr=[[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]]] | |_|_TableScan: t1_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -632,7 +632,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST]@0 as ordered_host, first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]@1 as first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED diff --git a/tests/cases/distributed/optimizer/last_value_advance.result b/tests/cases/distributed/optimizer/last_value_advance.result index d01ea184300..7b131d44f37 100644 --- a/tests/cases/distributed/optimizer/last_value_advance.result +++ b/tests/cases/distributed/optimizer/last_value_advance.result @@ -240,7 +240,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t.host]], aggr=[[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]]] | |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -267,7 +267,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST]@0 as ordered_host, last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST]@1 as last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]@2 as last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -606,7 +606,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t1.host]], aggr=[[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]]] | |_|_TableScan: t1_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -632,7 +632,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST]@0 as ordered_host, last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]@1 as last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED diff --git a/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.result b/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.result new file mode 100644 index 00000000000..91c0e3e851d --- /dev/null +++ b/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.result @@ -0,0 +1,118 @@ +CREATE TABLE pd_sub_l( + ts TIMESTAMP TIME INDEX, + device_id INTEGER, + area STRING, + lv INTEGER, + PRIMARY KEY(device_id, area) +) PARTITION ON COLUMNS (device_id, area) ( + device_id < 2 AND area < 'm', + device_id < 2 AND area >= 'm', + device_id >= 2 +) ENGINE=mito; + +Affected Rows: 0 + +CREATE TABLE pd_sub_r( + ts TIMESTAMP TIME INDEX, + device_id INTEGER, + area STRING, + rv INTEGER, + PRIMARY KEY(device_id, area) +) PARTITION ON COLUMNS (device_id, area) ( + device_id < 2 AND area < 'm', + device_id < 2 AND area >= 'm', + device_id >= 2 +) ENGINE=mito; + +Affected Rows: 0 + +INSERT INTO pd_sub_l VALUES + (1, 1, 'a', 10), + (2, 1, 'z', 20); + +Affected Rows: 2 + +INSERT INTO pd_sub_r VALUES + (1, 1, 'a', 100), + (2, 1, 'z', 200); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('pd_sub_l'); + ++-------------------------------+ +| ADMIN FLUSH_TABLE('pd_sub_l') | ++-------------------------------+ +| 0 | ++-------------------------------+ + +ADMIN FLUSH_TABLE('pd_sub_r'); + ++-------------------------------+ +| ADMIN FLUSH_TABLE('pd_sub_r') | ++-------------------------------+ +| 0 | ++-------------------------------+ + +-- Joining on only a subset of a multi-column partition key must still compare +-- rows across all matching partitions for that subset key. +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \d+\(\d+,\s+\d+\) REDACTED_REGION +-- SQLNESS REPLACE Hash\(\[device_id@0\],\s+\d+\) Hash([device_id@0],REDACTED_PARTITIONS) +-- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED +-- SQLNESS REPLACE input_partitions=REDACTED\s+\| input_partitions=REDACTED| +EXPLAIN SELECT l.area, r.area +FROM pd_sub_l l JOIN pd_sub_r r ON l.device_id = r.device_id +ORDER BY l.area, r.area; + ++---------------+-------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-------------------------------------------------------------------------------------------------------------------+ +| logical_plan | Sort: l.area ASC NULLS LAST, r.area ASC NULLS LAST | +| | Projection: l.area, r.area | +| | Inner Join: l.device_id = r.device_id | +| | Projection: l.device_id, l.area | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: l | +| | Filter: pd_sub_l.device_id IS NOT NULL | +| | TableScan: pd_sub_l, partial_filters=[pd_sub_l.device_id IS NOT NULL] | +| | ]] | +| | Projection: r.device_id, r.area | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: r | +| | Filter: pd_sub_r.device_id IS NOT NULL | +| | TableScan: pd_sub_r, partial_filters=[pd_sub_r.device_id IS NOT NULL] | +| | ]] | +| physical_plan | SortPreservingMergeExec: [area@0 ASC NULLS LAST, area@1 ASC NULLS LAST] | +| | SortExec: expr=[area@0 ASC NULLS LAST, area@1 ASC NULLS LAST], preserve_partitioning=[true] | +| | HashJoinExec: mode=Partitioned, join_type=Inner, on=[(device_id@0, device_id@0)], projection=[area@1, area@3] | +| | RepartitionExec: partitioning=Hash([device_id@0],REDACTED_PARTITIONS), input_partitions=REDACTED| +| | ProjectionExec: expr=[device_id@1 as device_id, area@2 as area] | +| | MergeScanExec: peers=[REDACTED_REGION, REDACTED_REGION, REDACTED_REGION, ] | +| | RepartitionExec: partitioning=Hash([device_id@0],REDACTED_PARTITIONS), input_partitions=REDACTED| +| | ProjectionExec: expr=[device_id@1 as device_id, area@2 as area] | +| | MergeScanExec: peers=[REDACTED_REGION, REDACTED_REGION, REDACTED_REGION, ] | +| | | ++---------------+-------------------------------------------------------------------------------------------------------------------+ + +SELECT l.area, r.area +FROM pd_sub_l l JOIN pd_sub_r r ON l.device_id = r.device_id +ORDER BY l.area, r.area; + ++------+------+ +| area | area | ++------+------+ +| a | a | +| a | z | +| z | a | +| z | z | ++------+------+ + +DROP TABLE pd_sub_l; + +Affected Rows: 0 + +DROP TABLE pd_sub_r; + +Affected Rows: 0 + diff --git a/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.sql b/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.sql new file mode 100644 index 00000000000..41337e74a0f --- /dev/null +++ b/tests/cases/distributed/optimizer/pass_distribution_partition_subset_join.sql @@ -0,0 +1,52 @@ +CREATE TABLE pd_sub_l( + ts TIMESTAMP TIME INDEX, + device_id INTEGER, + area STRING, + lv INTEGER, + PRIMARY KEY(device_id, area) +) PARTITION ON COLUMNS (device_id, area) ( + device_id < 2 AND area < 'm', + device_id < 2 AND area >= 'm', + device_id >= 2 +) ENGINE=mito; + +CREATE TABLE pd_sub_r( + ts TIMESTAMP TIME INDEX, + device_id INTEGER, + area STRING, + rv INTEGER, + PRIMARY KEY(device_id, area) +) PARTITION ON COLUMNS (device_id, area) ( + device_id < 2 AND area < 'm', + device_id < 2 AND area >= 'm', + device_id >= 2 +) ENGINE=mito; + +INSERT INTO pd_sub_l VALUES + (1, 1, 'a', 10), + (2, 1, 'z', 20); + +INSERT INTO pd_sub_r VALUES + (1, 1, 'a', 100), + (2, 1, 'z', 200); + +ADMIN FLUSH_TABLE('pd_sub_l'); +ADMIN FLUSH_TABLE('pd_sub_r'); + +-- Joining on only a subset of a multi-column partition key must still compare +-- rows across all matching partitions for that subset key. +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE \d+\(\d+,\s+\d+\) REDACTED_REGION +-- SQLNESS REPLACE Hash\(\[device_id@0\],\s+\d+\) Hash([device_id@0],REDACTED_PARTITIONS) +-- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED +-- SQLNESS REPLACE input_partitions=REDACTED\s+\| input_partitions=REDACTED| +EXPLAIN SELECT l.area, r.area +FROM pd_sub_l l JOIN pd_sub_r r ON l.device_id = r.device_id +ORDER BY l.area, r.area; + +SELECT l.area, r.area +FROM pd_sub_l l JOIN pd_sub_r r ON l.device_id = r.device_id +ORDER BY l.area, r.area; + +DROP TABLE pd_sub_l; +DROP TABLE pd_sub_r; diff --git a/tests/cases/distributed/optimizer/time_index_filter_pushdown.result b/tests/cases/distributed/optimizer/time_index_filter_pushdown.result index 22e96793747..016c49deea0 100644 --- a/tests/cases/distributed/optimizer/time_index_filter_pushdown.result +++ b/tests/cases/distributed/optimizer/time_index_filter_pushdown.result @@ -84,6 +84,82 @@ group by | windows | 2 | +---------+----------+ +-- Query the row at the latest timestamp through a scalar subquery. +-- Keep the distributed merge sort visible, but redact the optional local SortExec +-- that appears when query parallelism is lower than the region count. +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (?m)^\|_\|_SortExec:.*\n +EXPLAIN SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| Projection: cpu.rack, cpu.os, cpu.greptime_timestamp_| +|_|_Inner Join: cpu.greptime_timestamp = __scalar_sq_1.greptime_timestamp_| +|_|_Projection: cpu.rack, cpu.os, cpu.greptime_timestamp_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| TableScan: cpu_| +|_| ]]_| +|_|_SubqueryAlias: __scalar_sq_1_| +|_|_Limit: skip=0, fetch=1_| +|_|_MergeSort: cpu.greptime_timestamp DESC NULLS FIRST_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Limit: skip=0, fetch=1_| +|_|_Sort: cpu.greptime_timestamp DESC NULLS FIRST_| +|_|_Projection: cpu.greptime_timestamp_| +|_|_TableScan: cpu_| +|_| ]]_| +| physical_plan | HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(greptime_timestamp@0, greptime_timestamp@2)], projection=[rack@1, os@2, greptime_timestamp@3]_| +|_|_MergeSortExec: [greptime_timestamp@0 DESC], fetch=1_| +|_|_MergeScanExec: REDACTED +|_|_ProjectionExec: expr=[rack@0 as rack, os@1 as os, greptime_timestamp@3 as greptime_timestamp]_| +|_|_CooperativeExec_| +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + ++------+---------+---------------------+ +| rack | os | greptime_timestamp | ++------+---------+---------------------+ +| 8 | windows | 2023-06-12T02:10:00 | ++------+---------+---------------------+ + drop table cpu; Affected Rows: 0 diff --git a/tests/cases/distributed/optimizer/time_index_filter_pushdown.sql b/tests/cases/distributed/optimizer/time_index_filter_pushdown.sql index f5fa24efe77..c9c80c24a61 100644 --- a/tests/cases/distributed/optimizer/time_index_filter_pushdown.sql +++ b/tests/cases/distributed/optimizer/time_index_filter_pushdown.sql @@ -59,4 +59,47 @@ where group by os; -drop table cpu; \ No newline at end of file +-- Query the row at the latest timestamp through a scalar subquery. +-- Keep the distributed merge sort visible, but redact the optional local SortExec +-- that appears when query parallelism is lower than the region count. +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (?m)^\|_\|_SortExec:.*\n +EXPLAIN SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + +SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + +drop table cpu; diff --git a/tests/cases/distributed/repartition/repartition.result b/tests/cases/distributed/repartition/repartition.result index 72212699af8..7b25a13a061 100644 --- a/tests/cases/distributed/repartition/repartition.result +++ b/tests/cases/distributed/repartition/repartition.result @@ -12,6 +12,7 @@ CREATE TABLE alter_repartition_table( Affected Rows: 0 +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE alter_repartition_table REPARTITION ( device_id < 100 ) INTO ( @@ -19,7 +20,7 @@ ALTER TABLE alter_repartition_table REPARTITION ( device_id < 100 AND area >= 'South' ); -Affected Rows: 0 +Error: 1004(InvalidArguments), Invalid arguments: Repartition requires metasrv GC to be enabled SHOW CREATE TABLE alter_repartition_table; @@ -35,21 +36,21 @@ SHOW CREATE TABLE alter_repartition_table; | | PRIMARY KEY ("device_id") | | | ) | | | PARTITION ON COLUMNS ("device_id", "area") ( | -| | device_id < 100 AND area < 'South', | +| | device_id < 100, | | | device_id >= 100 AND device_id < 200, | -| | device_id >= 200, | -| | device_id < 100 AND area >= 'South' | +| | device_id >= 200 | | | ) | | | ENGINE=mito | | | | +-------------------------+--------------------------------------------------------+ +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE alter_repartition_table MERGE PARTITION ( - device_id < 100 AND area < 'South', - device_id < 100 AND area >= 'South' + device_id < 100, + device_id >= 100 AND device_id < 200, ); -Affected Rows: 0 +Error: 1004(InvalidArguments), Invalid arguments: Repartition requires metasrv GC to be enabled SHOW CREATE TABLE alter_repartition_table; @@ -124,11 +125,7 @@ ALTER TABLE alter_repartition_table_with_options REPARTITION ( WAIT = false ); -+--------------------------------------+ -| Procedure ID | -+--------------------------------------+ -| PROC_ID | -+--------------------------------------+ +Error: 1004(InvalidArguments), Invalid arguments: Repartition requires metasrv GC to be enabled DROP TABLE alter_repartition_table_with_options; @@ -177,7 +174,7 @@ WITH ( Affected Rows: 0 --- Split physical table partition +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE metric_physical_table SPLIT PARTITION ( host < 'h1' ) INTO ( @@ -185,7 +182,7 @@ ALTER TABLE metric_physical_table SPLIT PARTITION ( host >= 'h0' AND host < 'h1' ); -Affected Rows: 0 +Error: 1004(InvalidArguments), Invalid arguments: Repartition requires metasrv GC to be enabled SHOW CREATE TABLE metric_physical_table; @@ -200,10 +197,9 @@ SHOW CREATE TABLE metric_physical_table; | | PRIMARY KEY ("host") | | | ) | | | PARTITION ON COLUMNS ("host") ( | -| | host < 'h0', | +| | host < 'h1', | | | host >= 'h1' AND host < 'h2', | -| | host >= 'h2', | -| | host >= 'h0' AND host < 'h1' | +| | host >= 'h2' | | | ) | | | ENGINE=metric | | | WITH( | @@ -227,13 +223,13 @@ SELECT * FROM logical_table_v2; ++ ++ --- Merge physical table partition +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE metric_physical_table MERGE PARTITION ( - host < 'h0', - host >= 'h0' AND host < 'h1' + host < 'h1', + host >= 'h1' AND host < 'h2' ); -Affected Rows: 0 +Error: 1004(InvalidArguments), Invalid arguments: Repartition requires metasrv GC to be enabled SHOW CREATE TABLE metric_physical_table; diff --git a/tests/cases/distributed/repartition/repartition.sql b/tests/cases/distributed/repartition/repartition.sql index 92b1d6ad97b..180969f96fa 100644 --- a/tests/cases/distributed/repartition/repartition.sql +++ b/tests/cases/distributed/repartition/repartition.sql @@ -10,6 +10,7 @@ CREATE TABLE alter_repartition_table( device_id >= 200 ); +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE alter_repartition_table REPARTITION ( device_id < 100 ) INTO ( @@ -19,9 +20,10 @@ ALTER TABLE alter_repartition_table REPARTITION ( SHOW CREATE TABLE alter_repartition_table; +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE alter_repartition_table MERGE PARTITION ( - device_id < 100 AND area < 'South', - device_id < 100 AND area >= 'South' + device_id < 100, + device_id >= 100 AND device_id < 200, ); SHOW CREATE TABLE alter_repartition_table; @@ -109,7 +111,7 @@ WITH ( on_physical_table = "metric_physical_table" ); --- Split physical table partition +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE metric_physical_table SPLIT PARTITION ( host < 'h1' ) INTO ( @@ -126,10 +128,10 @@ SELECT * FROM logical_table_v1; SELECT * FROM logical_table_v2; --- Merge physical table partition +-- Repartition requires metasrv GC. Verify split requests are rejected when it is disabled. ALTER TABLE metric_physical_table MERGE PARTITION ( - host < 'h0', - host >= 'h0' AND host < 'h1' + host < 'h1', + host >= 'h1' AND host < 'h2' ); SHOW CREATE TABLE metric_physical_table; diff --git a/tests/cases/standalone/common/alter/alter_table_options.result b/tests/cases/standalone/common/alter/alter_table_options.result index f91fd9da88e..b169562c894 100644 --- a/tests/cases/standalone/common/alter/alter_table_options.result +++ b/tests/cases/standalone/common/alter/alter_table_options.result @@ -259,6 +259,79 @@ DROP TABLE ato; Affected Rows: 0 +CREATE TABLE write_buffer_size_options( + i INTEGER, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY(i) +); + +Affected Rows: 0 + +ALTER TABLE write_buffer_size_options SET 'write_buffer_size' = '128MiB'; + +Affected Rows: 0 + +SHOW CREATE TABLE write_buffer_size_options; + ++---------------------------+----------------------------------------------------------+ +| Table | Create Table | ++---------------------------+----------------------------------------------------------+ +| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( | +| | "i" INT NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("i") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | write_buffer_size = '128.0MiB' | +| | ) | ++---------------------------+----------------------------------------------------------+ + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE write_buffer_size_options; + ++---------------------------+----------------------------------------------------------+ +| Table | Create Table | ++---------------------------+----------------------------------------------------------+ +| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( | +| | "i" INT NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("i") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | write_buffer_size = '128.0MiB' | +| | ) | ++---------------------------+----------------------------------------------------------+ + +ALTER TABLE write_buffer_size_options UNSET 'write_buffer_size'; + +Affected Rows: 0 + +SHOW CREATE TABLE write_buffer_size_options; + ++---------------------------+----------------------------------------------------------+ +| Table | Create Table | ++---------------------------+----------------------------------------------------------+ +| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( | +| | "i" INT NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("i") | +| | ) | +| | | +| | ENGINE=mito | +| | | ++---------------------------+----------------------------------------------------------+ + +DROP TABLE write_buffer_size_options; + +Affected Rows: 0 + CREATE TABLE phy (ts timestamp time index, val double) engine=metric with ("physical_metric_table" = ""); Affected Rows: 0 diff --git a/tests/cases/standalone/common/alter/alter_table_options.sql b/tests/cases/standalone/common/alter/alter_table_options.sql index badf1275cdb..c3f67faeee5 100644 --- a/tests/cases/standalone/common/alter/alter_table_options.sql +++ b/tests/cases/standalone/common/alter/alter_table_options.sql @@ -55,6 +55,25 @@ SHOW CREATE TABLE ato; DROP TABLE ato; +CREATE TABLE write_buffer_size_options( + i INTEGER, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY(i) +); + +ALTER TABLE write_buffer_size_options SET 'write_buffer_size' = '128MiB'; + +SHOW CREATE TABLE write_buffer_size_options; + +-- SQLNESS ARG restart=true +SHOW CREATE TABLE write_buffer_size_options; + +ALTER TABLE write_buffer_size_options UNSET 'write_buffer_size'; + +SHOW CREATE TABLE write_buffer_size_options; + +DROP TABLE write_buffer_size_options; + CREATE TABLE phy (ts timestamp time index, val double) engine=metric with ("physical_metric_table" = ""); ALTER TABLE phy set ttl='2years'; diff --git a/tests/cases/standalone/common/information_schema/region_info.result b/tests/cases/standalone/common/information_schema/region_info.result index ed9a07add60..91cab62336d 100644 --- a/tests/cases/standalone/common/information_schema/region_info.result +++ b/tests/cases/standalone/common/information_schema/region_info.result @@ -41,7 +41,7 @@ ADMIN FLUSH_TABLE('region_info_case'); +---------------------------------------+ -- SQLNESS REPLACE (\s+\d+\s+) --- SQLNESS REPLACE (\{".*"\}) +-- SQLNESS REPLACE (\{.*\}) -- SQLNESS REPLACE (-{40,}) ---------------- -- SQLNESS REPLACE (region_options\s+\|) region_options | SELECT region_id, state, role, writable, committed_sequence, flushed_sequence, manifest_version, compaction_time_window, region_options, sst_format diff --git a/tests/cases/standalone/common/information_schema/region_info.sql b/tests/cases/standalone/common/information_schema/region_info.sql index 1aa682393f2..6c901387ce9 100644 --- a/tests/cases/standalone/common/information_schema/region_info.sql +++ b/tests/cases/standalone/common/information_schema/region_info.sql @@ -11,7 +11,7 @@ INSERT INTO region_info_case VALUES (1, 1), (2, 2); ADMIN FLUSH_TABLE('region_info_case'); -- SQLNESS REPLACE (\s+\d+\s+) --- SQLNESS REPLACE (\{".*"\}) +-- SQLNESS REPLACE (\{.*\}) -- SQLNESS REPLACE (-{40,}) ---------------- -- SQLNESS REPLACE (region_options\s+\|) region_options | SELECT region_id, state, role, writable, committed_sequence, flushed_sequence, manifest_version, compaction_time_window, region_options, sst_format diff --git a/tests/cases/standalone/common/order/order_by.result b/tests/cases/standalone/common/order/order_by.result index ceebb94ff5d..aed1a491c27 100644 --- a/tests/cases/standalone/common/order/order_by.result +++ b/tests/cases/standalone/common/order/order_by.result @@ -296,7 +296,7 @@ explain analyze select tag from t where num > 6 order by ts desc limit 2; | stage | node | plan_| +-+-+-+ | 0_| 0_|_ProjectionExec: expr=[tag@0 as tag] REDACTED -|_|_|_SortPreservingMergeExec: [ts@1 DESC], fetch=2 REDACTED +|_|_|_MergeSortExec: [ts@1 DESC], fetch=2 REDACTED |_|_|_SortExec: TopK(fetch=2), expr=[ts@1 DESC], preserve_partitioning=[true] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| diff --git a/tests/cases/standalone/common/promql/functions.result b/tests/cases/standalone/common/promql/functions.result index 02ff599ced9..325fbccd17d 100644 --- a/tests/cases/standalone/common/promql/functions.result +++ b/tests/cases/standalone/common/promql/functions.result @@ -62,16 +62,6 @@ tql eval (10, 10, '1s') double_exponential_smoothing(prom_series[10s], 0.4 + 0.1 | 1970-01-01T00:00:10 | 47.0806953125 | p | +---------------------+------------------------------------------------------------------------------------------+------+ --- holt_winters (backward compatibility) --- SQLNESS SORT_RESULT 3 1 -tql eval (10, 10, '1s') holt_winters(prom_series[10s], 0.4 + 0.1, 0.1); - -+---------------------+------------------------------------------------------------------------------------------+------+ -| ts | prom_double_exponential_smoothing(ts_range,val,Float64(0.4) + Float64(0.1),Float64(0.1)) | host | -+---------------------+------------------------------------------------------------------------------------------+------+ -| 1970-01-01T00:00:10 | 47.0806953125 | p | -+---------------------+------------------------------------------------------------------------------------------+------+ - DROP TABLE prom_series; Affected Rows: 0 diff --git a/tests/cases/standalone/common/promql/functions.sql b/tests/cases/standalone/common/promql/functions.sql index 6657fc49e81..0a257f29da8 100644 --- a/tests/cases/standalone/common/promql/functions.sql +++ b/tests/cases/standalone/common/promql/functions.sql @@ -34,10 +34,6 @@ tql eval (3, 3, '1s') predict_linear(prom_series[3s], 40 + 2); -- SQLNESS SORT_RESULT 3 1 tql eval (10, 10, '1s') double_exponential_smoothing(prom_series[10s], 0.4 + 0.1, 0.1); --- holt_winters (backward compatibility) --- SQLNESS SORT_RESULT 3 1 -tql eval (10, 10, '1s') holt_winters(prom_series[10s], 0.4 + 0.1, 0.1); - DROP TABLE prom_series; CREATE TABLE diff --git a/tests/cases/standalone/common/promql/scalar.result b/tests/cases/standalone/common/promql/scalar.result index c3292b4f5c5..ec1c768f693 100644 --- a/tests/cases/standalone/common/promql/scalar.result +++ b/tests/cases/standalone/common/promql/scalar.result @@ -586,8 +586,8 @@ INSERT INTO TABLE presence_metric VALUES Affected Rows: 15 --- NaN drops `cpu0` from the grouped count, while the NULL sample on `cpu2` --- still leaves a zero-valued row in `count(...) by (cpu)`. +-- Ordinary NaN keeps `cpu0` present in the grouped count, while the NULL sample +-- on `cpu2` still leaves a zero-valued row in `count(...) by (cpu)`. -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 600, '200s') count(presence_metric{instance="i1"}) by (cpu); @@ -595,6 +595,8 @@ TQL EVAL (0, 600, '200s') count(presence_metric{instance="i1"}) by (cpu); | cpu | ts | count(presence_metric.val) | +------+---------------------+----------------------------+ | cpu0 | 1970-01-01T00:00:00 | 2 | +| cpu0 | 1970-01-01T00:03:20 | 2 | +| cpu0 | 1970-01-01T00:06:40 | 2 | | cpu0 | 1970-01-01T00:10:00 | 2 | | cpu1 | 1970-01-01T00:00:00 | 1 | | cpu1 | 1970-01-01T00:03:20 | 1 | @@ -605,7 +607,8 @@ TQL EVAL (0, 600, '200s') count(presence_metric{instance="i1"}) by (cpu); | cpu2 | 1970-01-01T00:06:40 | 0 | +------+---------------------+----------------------------+ --- Nested-count rewrite should preserve grouped presence after stale-NaN filtering and null-value pruning. +-- Nested-count rewrite should preserve grouped presence after stale-marker filtering +-- and null-value pruning. -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 600, '200s') scalar(count(count(presence_metric{instance="i1"}) by (cpu))); @@ -613,8 +616,8 @@ TQL EVAL (0, 600, '200s') scalar(count(count(presence_metric{instance="i1"}) by | ts | scalar(count(count(presence_metric.val))) | +---------------------+-------------------------------------------+ | 1970-01-01T00:00:00 | 3.0 | -| 1970-01-01T00:03:20 | 2.0 | -| 1970-01-01T00:06:40 | 2.0 | +| 1970-01-01T00:03:20 | 3.0 | +| 1970-01-01T00:06:40 | 3.0 | | 1970-01-01T00:10:00 | 2.0 | +---------------------+-------------------------------------------+ @@ -626,8 +629,8 @@ TQL EVAL (0, 600, '200s') scalar(count(sum(presence_metric{instance="i1"}) by (c | ts | scalar(count(sum(presence_metric.val))) | +---------------------+-----------------------------------------+ | 1970-01-01T00:00:00 | 3.0 | -| 1970-01-01T00:03:20 | 1.0 | -| 1970-01-01T00:06:40 | 1.0 | +| 1970-01-01T00:03:20 | 2.0 | +| 1970-01-01T00:06:40 | 2.0 | | 1970-01-01T00:10:00 | 2.0 | +---------------------+-----------------------------------------+ diff --git a/tests/cases/standalone/common/promql/scalar.sql b/tests/cases/standalone/common/promql/scalar.sql index 662f9665fee..b0c85d89e35 100644 --- a/tests/cases/standalone/common/promql/scalar.sql +++ b/tests/cases/standalone/common/promql/scalar.sql @@ -186,12 +186,13 @@ INSERT INTO TABLE presence_metric VALUES (600000, 'i1', 'cpu0', 'b', 8.0), (600000, 'i2', 'cpu9', 'a', 103.0); --- NaN drops `cpu0` from the grouped count, while the NULL sample on `cpu2` --- still leaves a zero-valued row in `count(...) by (cpu)`. +-- Ordinary NaN keeps `cpu0` present in the grouped count, while the NULL sample +-- on `cpu2` still leaves a zero-valued row in `count(...) by (cpu)`. -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 600, '200s') count(presence_metric{instance="i1"}) by (cpu); --- Nested-count rewrite should preserve grouped presence after stale-NaN filtering and null-value pruning. +-- Nested-count rewrite should preserve grouped presence after stale-marker filtering +-- and null-value pruning. -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 600, '200s') scalar(count(count(presence_metric{instance="i1"}) by (cpu))); diff --git a/tests/cases/standalone/common/promql/set_operation.result b/tests/cases/standalone/common/promql/set_operation.result index 83af6fc8755..50a0b6cf3b6 100644 --- a/tests/cases/standalone/common/promql/set_operation.result +++ b/tests/cases/standalone/common/promql/set_operation.result @@ -466,12 +466,12 @@ tql eval (0, 2000, '400') sum(t1{job="a"}); -- SQLNESS SORT_RESULT 3 1 tql eval (0, 2000, '400') sum(t1{job="a"}) - sum(t1{job="e"} or vector(1)); -+---------------------+-----------------------------------------------------+ -| ts | t1.sum(t1.greptime_value) - .sum(t1.greptime_value) | -+---------------------+-----------------------------------------------------+ -| 1970-01-01T00:00:00 | 0.0 | -| 1970-01-01T00:20:00 | 2.0 | -+---------------------+-----------------------------------------------------+ ++---------------------+---------------------------------------------------------+ +| ts | lhs.sum(t1.greptime_value) - rhs.sum(t1.greptime_value) | ++---------------------+---------------------------------------------------------+ +| 1970-01-01T00:00:00 | 0.0 | +| 1970-01-01T00:20:00 | 2.0 | ++---------------------+---------------------------------------------------------+ drop table t1; @@ -545,12 +545,12 @@ tql eval (0, 2000, '400') count(max by (namespace) (stats_used_bytes{namespace=~ +---------------------+---------------------------------------------+ | ts | count(max(stats_used_bytes.greptime_value)) | +---------------------+---------------------------------------------+ -| 1970-01-01T00:00:00 | 0 | -| 1970-01-01T00:06:40 | 0 | -| 1970-01-01T00:13:20 | 0 | -| 1970-01-01T00:20:00 | 2 | -| 1970-01-01T00:26:40 | 0 | -| 1970-01-01T00:33:20 | 0 | +| 1970-01-01T00:00:00 | 0.0 | +| 1970-01-01T00:06:40 | 0.0 | +| 1970-01-01T00:13:20 | 0.0 | +| 1970-01-01T00:20:00 | 2.0 | +| 1970-01-01T00:26:40 | 0.0 | +| 1970-01-01T00:33:20 | 0.0 | +---------------------+---------------------------------------------+ -- SQLNESS SORT_RESULT 3 1 diff --git a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result index d5313ba41c6..8c90d97a325 100644 --- a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result +++ b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result @@ -512,10 +512,10 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_ProjectionExec: expr=[host@2 as host, job@3 as job, ts@5 as ts, __tsid@4 as __tsid, greptime_value@0 / greptime_value@1 as tsid_binary_join_right.greptime_value / tsid_binary_join_left.greptime_value] REDACTED +| 0_| 0_|_ProjectionExec: expr=[host@2 as host, job@3 as job, ts@5 as ts, __tsid@4 as __tsid, greptime_value@0 / greptime_value@1 as lhs.greptime_value / rhs.greptime_value] REDACTED |_|_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(__tsid@1, __tsid@3), (ts@0, ts@4)], projection=[greptime_value@2, greptime_value@3, host@4, job@5, __tsid@6, ts@7], NullsEqual: true REDACTED |_|_|_ProjectionExec: expr=[ts@0 as ts, __tsid@1 as __tsid, greptime_value@2 as greptime_value] REDACTED -|_|_|_UnionDistinctOnExec: on col=[["host", "job"]], ts_col=[ts] REDACTED +|_|_|_UnionDistinctOnExec: on col=[5, 6], ts_col=0 REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_CoalescePartitionsExec REDACTED @@ -523,14 +523,14 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi |_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| -| 1_| 0_|_ProjectionExec: expr=[ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job] REDACTED +| 1_| 0_|_ProjectionExec: expr=[ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job, CASE WHEN host@1 IS NOT NULL THEN host@1 ELSE_END as __promql_or_match_0, CASE WHEN job@2 IS NOT NULL THEN job@2 ELSE_END as __promql_or_match_1] REDACTED |_|_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED |_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED |_|_|_| -| 1_| 0_|_ProjectionExec: expr=[ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job] REDACTED +| 1_| 0_|_ProjectionExec: expr=[ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job, CASE WHEN host@1 IS NOT NULL THEN host@1 ELSE_END as __promql_or_match_0, CASE WHEN job@2 IS NOT NULL THEN job@2 ELSE_END as __promql_or_match_1] REDACTED |_|_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED @@ -672,14 +672,14 @@ TQL EVAL (0, 5, '5s') ((tsid_binary_join_left > bool tsid_binary_join_right) + t -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsid_binary_join_left; -+-------+------+---------------------+------------------------------------------------------------------------------+ -| host | job | ts | tsid_binary_join_right.greptime_value / tsid_binary_join_left.greptime_value | -+-------+------+---------------------+------------------------------------------------------------------------------+ -| host1 | job1 | 1970-01-01T00:00:00 | 1.0 | -| host1 | job1 | 1970-01-01T00:00:05 | 1.0 | -| host2 | job2 | 1970-01-01T00:00:00 | 1.0 | -| host2 | job2 | 1970-01-01T00:00:05 | 1.0 | -+-------+------+---------------------+------------------------------------------------------------------------------+ ++-------+------+---------------------+-----------------------------------------+ +| host | job | ts | lhs.greptime_value / rhs.greptime_value | ++-------+------+---------------------+-----------------------------------------+ +| host1 | job1 | 1970-01-01T00:00:00 | 1.0 | +| host1 | job1 | 1970-01-01T00:00:05 | 1.0 | +| host2 | job2 | 1970-01-01T00:00:00 | 1.0 | +| host2 | job2 | 1970-01-01T00:00:05 | 1.0 | ++-------+------+---------------------+-----------------------------------------+ -- Range functions are outside the first island version; the range selector subtree must -- remain a barrier. diff --git a/tests/cases/standalone/common/show/show_create.result b/tests/cases/standalone/common/show/show_create.result index 72392a0997d..18879ed95e8 100644 --- a/tests/cases/standalone/common/show/show_create.result +++ b/tests/cases/standalone/common/show/show_create.result @@ -15,7 +15,7 @@ PARTITION ON COLUMNS (`id`) ( ENGINE=mito WITH( ttl = '7d', - write_buffer_size = 1024 + write_buffer_size = '1KiB' ); Affected Rows: 0 diff --git a/tests/cases/standalone/common/show/show_create.sql b/tests/cases/standalone/common/show/show_create.sql index eb0f90c6746..fd8b1c16ea2 100644 --- a/tests/cases/standalone/common/show/show_create.sql +++ b/tests/cases/standalone/common/show/show_create.sql @@ -15,7 +15,7 @@ PARTITION ON COLUMNS (`id`) ( ENGINE=mito WITH( ttl = '7d', - write_buffer_size = 1024 + write_buffer_size = '1KiB' ); SHOW CREATE TABLE system_metrics; diff --git a/tests/cases/standalone/common/tql/aggr_over_time.result b/tests/cases/standalone/common/tql/aggr_over_time.result index c846d73c051..e39c9d52384 100644 --- a/tests/cases/standalone/common/tql/aggr_over_time.result +++ b/tests/cases/standalone/common/tql/aggr_over_time.result @@ -295,6 +295,7 @@ tql eval (60, 60, '1s') min_over_time(data[2m]); | 1970-01-01T00:01:00 | 0.0 | some_nan | | 1970-01-01T00:01:00 | 0.0 | some_nan3 | | 1970-01-01T00:01:00 | 1.0 | some_nan2 | +| 1970-01-01T00:01:00 | NaN | only_nan | +---------------------+----------------------------------+-----------+ -- eval instant at 1m max_over_time(data[2m]) @@ -313,6 +314,7 @@ tql eval (60, 60, '1s') max_over_time(data[2m]); | 1970-01-01T00:01:00 | 2.0 | some_nan | | 1970-01-01T00:01:00 | 2.0 | some_nan2 | | 1970-01-01T00:01:00 | 3.0 | numbers | +| 1970-01-01T00:01:00 | NaN | only_nan | +---------------------+----------------------------------+-----------+ -- eval instant at 1m last_over_time(data[2m]) @@ -327,10 +329,11 @@ tql eval (60, 60, '1s') last_over_time(data[2m]); +---------------------+-----------------------------------+-----------+ | ts | prom_last_over_time(ts_range,val) | ty | +---------------------+-----------------------------------+-----------+ -| 1970-01-01T00:01:00 | 0.0 | some_nan | | 1970-01-01T00:01:00 | 1.0 | some_nan2 | | 1970-01-01T00:01:00 | 1.0 | some_nan3 | | 1970-01-01T00:01:00 | 3.0 | numbers | +| 1970-01-01T00:01:00 | NaN | only_nan | +| 1970-01-01T00:01:00 | NaN | some_nan | +---------------------+-----------------------------------+-----------+ drop table data; diff --git a/tests/cases/standalone/common/tql/operator.result b/tests/cases/standalone/common/tql/operator.result index bc105b14549..bc0971c3f3b 100644 --- a/tests/cases/standalone/common/tql/operator.result +++ b/tests/cases/standalone/common/tql/operator.result @@ -40,11 +40,14 @@ tql eval (60, 60, '1s') trigy atan2 trigx; -- eval instant at 1m trigy atan2 trigNaN -- trigy{} NaN --- This query doesn't have result because `trignan` is NaN and will be filtered out. +-- Ordinary NaN is a valid sample and propagates through `atan2`. tql eval (60, 60, '1s') trigy atan2 trignan; -++ -++ ++---------------------+------------------------------+ +| ts | atan2(trigy.val,trignan.val) | ++---------------------+------------------------------+ +| 1970-01-01T00:01:00 | NaN | ++---------------------+------------------------------+ -- eval instant at 1m 10 atan2 20 -- 0.4636476090008061 diff --git a/tests/cases/standalone/common/tql/operator.sql b/tests/cases/standalone/common/tql/operator.sql index 2ab77dceb9d..8d10f62bcc7 100644 --- a/tests/cases/standalone/common/tql/operator.sql +++ b/tests/cases/standalone/common/tql/operator.sql @@ -23,7 +23,7 @@ tql eval (60, 60, '1s') trigy atan2 trigx; -- eval instant at 1m trigy atan2 trigNaN -- trigy{} NaN --- This query doesn't have result because `trignan` is NaN and will be filtered out. +-- Ordinary NaN is a valid sample and propagates through `atan2`. tql eval (60, 60, '1s') trigy atan2 trignan; -- eval instant at 1m 10 atan2 20 @@ -73,4 +73,3 @@ VALUES tql eval (1743465600.5, 1743465610, '1s') irate(t[2s]); drop table t; - diff --git a/tests/cases/standalone/common/tql/partition.result b/tests/cases/standalone/common/tql/partition.result index 5d0fd981238..45927f8abbd 100644 --- a/tests/cases/standalone/common/tql/partition.result +++ b/tests/cases/standalone/common/tql/partition.result @@ -133,8 +133,8 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[1000], eval range=[60000], time index=[j] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED -|_|_|_SortExec: expr=[k@2 ASC, l@3 ASC, j@1 ASC], preserve_partitioning=[true] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeSortExec: [k@2 ASC, l@3 ASC, j@1 ASC] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [k@2 ASC, l@3 ASC, j@1 ASC] REDACTED diff --git a/tests/cases/standalone/common/window/latest_per_series.result b/tests/cases/standalone/common/window/latest_per_series.result new file mode 100644 index 00000000000..c389c8aeb8a --- /dev/null +++ b/tests/cases/standalone/common/window/latest_per_series.result @@ -0,0 +1,187 @@ +CREATE TABLE netflow_raw ( + greptime_timestamp TIMESTAMP TIME INDEX, + src STRING, + dst_port INT, + dst_addr STRING, + byte_count DOUBLE, + flow_id INT, + PRIMARY KEY (src, dst_port) +); + +Affected Rows: 0 + +INSERT INTO netflow_raw VALUES + (1000, 'source-a', 80, 'target-80-old', 100.0, 1), + (5000, 'source-a', 80, 'target-80-new', 500.0, 2), + (2000, 'source-a', 443, 'target-443-old', 200.0, 3), + (7000, 'source-a', 443, 'target-443-new', 700.0, 4), + (3000, 'source-a', 22, 'target-22-only', 300.0, 5), + (4000, 'source-a', 8080, 'target-8080-only', 400.0, 6), + (8000, 'source-b', 9999, 'target-other-src', 999.0, 7); + +Affected Rows: 7 + +-- Latest row per dst_port using a window subquery. The window result must be +-- filtered outside the subquery, not directly in the WHERE clause. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id +FROM ( + SELECT + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id, + ROW_NUMBER() OVER ( + PARTITION BY dst_port + ORDER BY greptime_timestamp DESC + ) AS rn + FROM netflow_raw + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +) latest_per_port +WHERE rn = 1 +ORDER BY greptime_timestamp DESC, dst_port +LIMIT 3; + ++---------------------+----------+------------------+------------+---------+ +| greptime_timestamp | dst_port | dst_addr | byte_count | flow_id | ++---------------------+----------+------------------+------------+---------+ +| 1970-01-01T00:00:07 | 443 | target-443-new | 700.0 | 4 | +| 1970-01-01T00:00:05 | 80 | target-80-new | 500.0 | 2 | +| 1970-01-01T00:00:04 | 8080 | target-8080-only | 400.0 | 6 | ++---------------------+----------+------------------+------------+---------+ + +-- DISTINCT ON keeps the first row per dst_port according to its own ORDER BY. +-- Wrap it to apply the final "latest groups first" ordering and LIMIT. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id +FROM ( + SELECT DISTINCT ON (dst_port) + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id + FROM netflow_raw + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp + ORDER BY dst_port, greptime_timestamp DESC +) latest_per_port +ORDER BY greptime_timestamp DESC, dst_port +LIMIT 3; + ++---------------------+----------+------------------+------------+---------+ +| greptime_timestamp | dst_port | dst_addr | byte_count | flow_id | ++---------------------+----------+------------------+------------+---------+ +| 1970-01-01T00:00:07 | 443 | target-443-new | 700.0 | 4 | +| 1970-01-01T00:00:05 | 80 | target-80-new | 500.0 | 2 | +| 1970-01-01T00:00:04 | 8080 | target-8080-only | 400.0 | 6 | ++---------------------+----------+------------------+------------+---------+ + +-- RANK returns all rows tied at the latest timestamp for each dst_port. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id, r +FROM ( + SELECT + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id, + RANK() OVER ( + PARTITION BY dst_port + ORDER BY greptime_timestamp DESC + ) AS r + FROM netflow_raw + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +) latest_per_port +WHERE r = 1 +ORDER BY greptime_timestamp DESC, dst_port +LIMIT 3; + ++---------------------+----------+------------------+------------+---------+---+ +| greptime_timestamp | dst_port | dst_addr | byte_count | flow_id | r | ++---------------------+----------+------------------+------------+---------+---+ +| 1970-01-01T00:00:07 | 443 | target-443-new | 700.0 | 4 | 1 | +| 1970-01-01T00:00:05 | 80 | target-80-new | 500.0 | 2 | 1 | +| 1970-01-01T00:00:04 | 8080 | target-8080-only | 400.0 | 6 | 1 | ++---------------------+----------+------------------+------------+---------+---+ + +-- Ordered aggregate variant when the caller can enumerate the columns to fetch. +SELECT + dst_port, + last_value(greptime_timestamp ORDER BY greptime_timestamp) AS latest_ts, + last_value(dst_addr ORDER BY greptime_timestamp) AS latest_dst_addr, + last_value(byte_count ORDER BY greptime_timestamp) AS latest_byte_count, + last_value(flow_id ORDER BY greptime_timestamp) AS latest_flow_id +FROM netflow_raw +WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +GROUP BY dst_port +ORDER BY latest_ts DESC, dst_port +LIMIT 3; + ++----------+---------------------+------------------+-------------------+----------------+ +| dst_port | latest_ts | latest_dst_addr | latest_byte_count | latest_flow_id | ++----------+---------------------+------------------+-------------------+----------------+ +| 443 | 1970-01-01T00:00:07 | target-443-new | 700.0 | 4 | +| 80 | 1970-01-01T00:00:05 | target-80-new | 500.0 | 2 | +| 8080 | 1970-01-01T00:00:04 | target-8080-only | 400.0 | 6 | ++----------+---------------------+------------------+-------------------+----------------+ + +CREATE TABLE netflow_raw_ties ( + greptime_timestamp TIMESTAMP TIME INDEX, + src STRING, + dst_port INT, + dst_addr STRING, + byte_count DOUBLE, + flow_id INT, + PRIMARY KEY (src, dst_port, flow_id) +); + +Affected Rows: 0 + +INSERT INTO netflow_raw_ties VALUES + (1000, 'source-a', 443, 'target-443-old', 100.0, 1), + (9000, 'source-a', 443, 'target-443-tie-a', 900.0, 2), + (9000, 'source-a', 443, 'target-443-tie-b', 901.0, 3), + (8000, 'source-a', 80, 'target-80-new', 800.0, 4), + (9500, 'source-b', 443, 'target-other-src', 950.0, 5); + +Affected Rows: 5 + +-- RANK with r = 1 keeps all rows tied at the latest timestamp per dst_port. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id, r +FROM ( + SELECT + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id, + RANK() OVER ( + PARTITION BY dst_port + ORDER BY greptime_timestamp DESC + ) AS r + FROM netflow_raw_ties + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +) latest_per_port +WHERE r = 1 +ORDER BY greptime_timestamp DESC, dst_port, flow_id; + ++---------------------+----------+------------------+------------+---------+---+ +| greptime_timestamp | dst_port | dst_addr | byte_count | flow_id | r | ++---------------------+----------+------------------+------------+---------+---+ +| 1970-01-01T00:00:09 | 443 | target-443-tie-a | 900.0 | 2 | 1 | +| 1970-01-01T00:00:09 | 443 | target-443-tie-b | 901.0 | 3 | 1 | +| 1970-01-01T00:00:08 | 80 | target-80-new | 800.0 | 4 | 1 | ++---------------------+----------+------------------+------------+---------+---+ + +DROP TABLE netflow_raw_ties; + +Affected Rows: 0 + +DROP TABLE netflow_raw; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/window/latest_per_series.sql b/tests/cases/standalone/common/window/latest_per_series.sql new file mode 100644 index 00000000000..f809c8777ee --- /dev/null +++ b/tests/cases/standalone/common/window/latest_per_series.sql @@ -0,0 +1,134 @@ +CREATE TABLE netflow_raw ( + greptime_timestamp TIMESTAMP TIME INDEX, + src STRING, + dst_port INT, + dst_addr STRING, + byte_count DOUBLE, + flow_id INT, + PRIMARY KEY (src, dst_port) +); + +INSERT INTO netflow_raw VALUES + (1000, 'source-a', 80, 'target-80-old', 100.0, 1), + (5000, 'source-a', 80, 'target-80-new', 500.0, 2), + (2000, 'source-a', 443, 'target-443-old', 200.0, 3), + (7000, 'source-a', 443, 'target-443-new', 700.0, 4), + (3000, 'source-a', 22, 'target-22-only', 300.0, 5), + (4000, 'source-a', 8080, 'target-8080-only', 400.0, 6), + (8000, 'source-b', 9999, 'target-other-src', 999.0, 7); + +-- Latest row per dst_port using a window subquery. The window result must be +-- filtered outside the subquery, not directly in the WHERE clause. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id +FROM ( + SELECT + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id, + ROW_NUMBER() OVER ( + PARTITION BY dst_port + ORDER BY greptime_timestamp DESC + ) AS rn + FROM netflow_raw + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +) latest_per_port +WHERE rn = 1 +ORDER BY greptime_timestamp DESC, dst_port +LIMIT 3; + +-- DISTINCT ON keeps the first row per dst_port according to its own ORDER BY. +-- Wrap it to apply the final "latest groups first" ordering and LIMIT. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id +FROM ( + SELECT DISTINCT ON (dst_port) + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id + FROM netflow_raw + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp + ORDER BY dst_port, greptime_timestamp DESC +) latest_per_port +ORDER BY greptime_timestamp DESC, dst_port +LIMIT 3; + +-- RANK returns all rows tied at the latest timestamp for each dst_port. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id, r +FROM ( + SELECT + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id, + RANK() OVER ( + PARTITION BY dst_port + ORDER BY greptime_timestamp DESC + ) AS r + FROM netflow_raw + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +) latest_per_port +WHERE r = 1 +ORDER BY greptime_timestamp DESC, dst_port +LIMIT 3; + +-- Ordered aggregate variant when the caller can enumerate the columns to fetch. +SELECT + dst_port, + last_value(greptime_timestamp ORDER BY greptime_timestamp) AS latest_ts, + last_value(dst_addr ORDER BY greptime_timestamp) AS latest_dst_addr, + last_value(byte_count ORDER BY greptime_timestamp) AS latest_byte_count, + last_value(flow_id ORDER BY greptime_timestamp) AS latest_flow_id +FROM netflow_raw +WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +GROUP BY dst_port +ORDER BY latest_ts DESC, dst_port +LIMIT 3; + +CREATE TABLE netflow_raw_ties ( + greptime_timestamp TIMESTAMP TIME INDEX, + src STRING, + dst_port INT, + dst_addr STRING, + byte_count DOUBLE, + flow_id INT, + PRIMARY KEY (src, dst_port, flow_id) +); + +INSERT INTO netflow_raw_ties VALUES + (1000, 'source-a', 443, 'target-443-old', 100.0, 1), + (9000, 'source-a', 443, 'target-443-tie-a', 900.0, 2), + (9000, 'source-a', 443, 'target-443-tie-b', 901.0, 3), + (8000, 'source-a', 80, 'target-80-new', 800.0, 4), + (9500, 'source-b', 443, 'target-other-src', 950.0, 5); + +-- RANK with r = 1 keeps all rows tied at the latest timestamp per dst_port. +SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id, r +FROM ( + SELECT + greptime_timestamp, + dst_port, + dst_addr, + byte_count, + flow_id, + RANK() OVER ( + PARTITION BY dst_port + ORDER BY greptime_timestamp DESC + ) AS r + FROM netflow_raw_ties + WHERE src = 'source-a' + AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp +) latest_per_port +WHERE r = 1 +ORDER BY greptime_timestamp DESC, dst_port, flow_id; + +DROP TABLE netflow_raw_ties; + +DROP TABLE netflow_raw; diff --git a/tests/cases/standalone/optimizer/first_value_advance.result b/tests/cases/standalone/optimizer/first_value_advance.result index ffc0b341f0a..afcefae4cff 100644 --- a/tests/cases/standalone/optimizer/first_value_advance.result +++ b/tests/cases/standalone/optimizer/first_value_advance.result @@ -240,7 +240,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t.host]], aggr=[[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]]] | |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -267,7 +267,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED @@ -603,7 +603,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t1.host]], aggr=[[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]]] | |_|_TableScan: t1_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -629,7 +629,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED diff --git a/tests/cases/standalone/optimizer/last_value_advance.result b/tests/cases/standalone/optimizer/last_value_advance.result index 0af561e6185..a1882a75887 100644 --- a/tests/cases/standalone/optimizer/last_value_advance.result +++ b/tests/cases/standalone/optimizer/last_value_advance.result @@ -240,7 +240,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t.host]], aggr=[[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]]] | |_|_TableScan: t_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -267,7 +267,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED @@ -603,7 +603,7 @@ order by ordered_host; |_|_Aggregate: groupBy=[[t1.host]], aggr=[[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]]] | |_|_TableScan: t1_| |_| ]]_| -| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_| +| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -629,7 +629,7 @@ order by ordered_host; +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED +| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED diff --git a/tests/cases/standalone/optimizer/time_index_filter_pushdown.result b/tests/cases/standalone/optimizer/time_index_filter_pushdown.result index de6222ae28e..fd0d869a33f 100644 --- a/tests/cases/standalone/optimizer/time_index_filter_pushdown.result +++ b/tests/cases/standalone/optimizer/time_index_filter_pushdown.result @@ -197,6 +197,31 @@ group by | windows | 2 | +---------+----------+ +-- Query the row at the latest timestamp through a scalar subquery. +SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + ++------+---------+---------------------+ +| rack | os | greptime_timestamp | ++------+---------+---------------------+ +| 8 | windows | 2023-06-12T02:10:00 | ++------+---------+---------------------+ + drop table cpu; Affected Rows: 0 diff --git a/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql b/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql index c0636e26a4b..7a7030d0925 100644 --- a/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql +++ b/tests/cases/standalone/optimizer/time_index_filter_pushdown.sql @@ -120,4 +120,23 @@ where group by os; +-- Query the row at the latest timestamp through a scalar subquery. +SELECT + rack, + os, + greptime_timestamp +FROM + cpu +WHERE + greptime_timestamp = ( + SELECT + greptime_timestamp + FROM + cpu + ORDER BY + greptime_timestamp DESC + LIMIT + 1 + ); + drop table cpu; diff --git a/tests/compatibility/cases/legacy_json2_non_append_table/case.toml b/tests/compatibility/cases/legacy_json2_non_append_table/case.toml new file mode 100644 index 00000000000..8bb090ea788 --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_non_append_table/case.toml @@ -0,0 +1,8 @@ +name = "legacy_json2_non_append_table" +reason = "Verify current binaries can read a non-append JSON2 table created by old binaries and migrate it to append_mode=true." +introduced_by = "PR #8434" +topologies = ["distributed"] +from_range = ["=v1.1.0", "=v1.1.1", "=v1.1.2"] +to_range = [">v1.1.2"] +features = ["json", "json2", "table", "append_mode", "read"] +owner = "query" diff --git a/tests/compatibility/cases/legacy_json2_non_append_table/setup.sql b/tests/compatibility/cases/legacy_json2_non_append_table/setup.sql new file mode 100644 index 00000000000..6eeeed1f908 --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_non_append_table/setup.sql @@ -0,0 +1,10 @@ +CREATE TABLE t_legacy_json2_non_append_table ( + ts TIMESTAMP TIME INDEX, + j JSON2 +); + +INSERT INTO t_legacy_json2_non_append_table (ts, j) VALUES + ('2026-07-08 00:00:00+0000', '{"a": 1, "nested": {"s": "old-1"}}'), + ('2026-07-08 00:01:00+0000', '{"a": 2, "nested": {"s": "old-2"}}'); + +ADMIN FLUSH_TABLE('t_legacy_json2_non_append_table'); diff --git a/tests/compatibility/cases/legacy_json2_non_append_table/verify.result b/tests/compatibility/cases/legacy_json2_non_append_table/verify.result new file mode 100644 index 00000000000..2cc3c230e2b --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_non_append_table/verify.result @@ -0,0 +1,63 @@ +SELECT ts, j.a AS a, j.nested.s AS nested_s +FROM t_legacy_json2_non_append_table +ORDER BY ts; + ++---------------------+---+----------+ +| ts | a | nested_s | ++---------------------+---+----------+ +| 2026-07-08T00:00:00 | 1 | old-1 | +| 2026-07-08T00:01:00 | 2 | old-2 | ++---------------------+---+----------+ + +SHOW CREATE TABLE t_legacy_json2_non_append_table; + ++---------------------------------+----------------------------------------------------------------+ +| Table | Create Table | ++---------------------------------+----------------------------------------------------------------+ +| t_legacy_json2_non_append_table | CREATE TABLE IF NOT EXISTS "t_legacy_json2_non_append_table" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2 NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | | ++---------------------------------+----------------------------------------------------------------+ + +ALTER TABLE t_legacy_json2_non_append_table SET 'append_mode' = 'true'; + +Affected Rows: 0 + +SHOW CREATE TABLE t_legacy_json2_non_append_table; + ++---------------------------------+----------------------------------------------------------------+ +| Table | Create Table | ++---------------------------------+----------------------------------------------------------------+ +| t_legacy_json2_non_append_table | CREATE TABLE IF NOT EXISTS "t_legacy_json2_non_append_table" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2 NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++---------------------------------+----------------------------------------------------------------+ + +INSERT INTO t_legacy_json2_non_append_table (ts, j) VALUES + ('2026-07-08 00:02:00+0000', '{"a": 3, "nested": {"s": "current-3"}}'); + +Affected Rows: 1 + +SELECT ts, j.a AS a, j.nested.s AS nested_s +FROM t_legacy_json2_non_append_table +ORDER BY ts; + ++---------------------+---+-----------+ +| ts | a | nested_s | ++---------------------+---+-----------+ +| 2026-07-08T00:00:00 | 1 | old-1 | +| 2026-07-08T00:01:00 | 2 | old-2 | +| 2026-07-08T00:02:00 | 3 | current-3 | ++---------------------+---+-----------+ diff --git a/tests/compatibility/cases/legacy_json2_non_append_table/verify.sql b/tests/compatibility/cases/legacy_json2_non_append_table/verify.sql new file mode 100644 index 00000000000..2dbc7fc729e --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_non_append_table/verify.sql @@ -0,0 +1,16 @@ +SELECT ts, j.a AS a, j.nested.s AS nested_s +FROM t_legacy_json2_non_append_table +ORDER BY ts; + +SHOW CREATE TABLE t_legacy_json2_non_append_table; + +ALTER TABLE t_legacy_json2_non_append_table SET 'append_mode' = 'true'; + +SHOW CREATE TABLE t_legacy_json2_non_append_table; + +INSERT INTO t_legacy_json2_non_append_table (ts, j) VALUES + ('2026-07-08 00:02:00+0000', '{"a": 3, "nested": {"s": "current-3"}}'); + +SELECT ts, j.a AS a, j.nested.s AS nested_s +FROM t_legacy_json2_non_append_table +ORDER BY ts; diff --git a/tests/compatibility/cases/legacy_json2_variant_payload/case.toml b/tests/compatibility/cases/legacy_json2_variant_payload/case.toml new file mode 100644 index 00000000000..dd46e28ba4d --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_variant_payload/case.toml @@ -0,0 +1,8 @@ +name = "legacy_json2_variant_payload" +reason = "Verify current binaries can read JSON2 variant payloads written by old binaries after the JSON2 encoding change." +introduced_by = "PR #8435" +topologies = ["distributed"] +from_range = ["=v1.1.0", "=v1.1.1", "=v1.1.2"] +to_range = [">v1.1.1"] +features = ["json", "json2", "table", "read"] +owner = "query" diff --git a/tests/compatibility/cases/legacy_json2_variant_payload/setup.sql b/tests/compatibility/cases/legacy_json2_variant_payload/setup.sql new file mode 100644 index 00000000000..cef3a913c93 --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_variant_payload/setup.sql @@ -0,0 +1,29 @@ +CREATE TABLE t_legacy_json2_variant_payload ( + ts TIMESTAMP TIME INDEX, + j JSON2 +) WITH ( + 'append_mode' = 'true' +); + +INSERT INTO t_legacy_json2_variant_payload (ts, j) VALUES + ('2026-07-08 00:00:00+0000', '{"a": {"b": 1}, "c": "s1", "d": [{"e": {"f": 0.1}}]}'), + ('2026-07-08 00:01:00+0000', '{"a": {"b": -2}, "c": [1], "d": [{"e": {"f": 0.2}}]}'), + ('2026-07-08 00:02:00+0000', '{"a": {"b": "s3"}, "c": "s3", "d": [{"e": {"g": -0.3}}]}'), + ('2026-07-08 00:03:00+0000', '{"a": {"b": null}, "c": null}'), + ('2026-07-08 00:04:00+0000', '{"c": "s5", "d": {"e": true}}'), + ('2026-07-08 00:05:00+0000', '{"a": {"b": 6}}'); + +ADMIN FLUSH_TABLE('t_legacy_json2_variant_payload'); + +CREATE TABLE t_legacy_json2_variant_payload_object_d ( + ts TIMESTAMP TIME INDEX, + j JSON2 +) WITH ( + 'append_mode' = 'true' +); + +INSERT INTO t_legacy_json2_variant_payload_object_d (ts, j) VALUES + ('2026-07-08 01:00:00+0000', '{"d": {"e": true}}'), + ('2026-07-08 01:01:00+0000', '{"d": {"e": false}}'); + +ADMIN FLUSH_TABLE('t_legacy_json2_variant_payload_object_d'); diff --git a/tests/compatibility/cases/legacy_json2_variant_payload/verify.result b/tests/compatibility/cases/legacy_json2_variant_payload/verify.result new file mode 100644 index 00000000000..a474cc84345 --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_variant_payload/verify.result @@ -0,0 +1,60 @@ +SELECT ts, j.a.b AS a_b, j.c AS c +FROM t_legacy_json2_variant_payload +ORDER BY ts; + ++---------------------+-----+-----+ +| ts | a_b | c | ++---------------------+-----+-----+ +| 2026-07-08T00:00:00 | 1 | s1 | +| 2026-07-08T00:01:00 | -2 | [1] | +| 2026-07-08T00:02:00 | s3 | s3 | +| 2026-07-08T00:03:00 | | | +| 2026-07-08T00:04:00 | | s5 | +| 2026-07-08T00:05:00 | 6 | | ++---------------------+-----+-----+ + +SELECT + ts, + j.d.e AS d_e +FROM t_legacy_json2_variant_payload_object_d +ORDER BY ts; + ++---------------------+-------+ +| ts | d_e | ++---------------------+-------+ +| 2026-07-08T01:00:00 | true | +| 2026-07-08T01:01:00 | false | ++---------------------+-------+ + +-- JSON2 projection rendering is a user-visible compatibility contract. +SELECT ts, j.d AS d +FROM t_legacy_json2_variant_payload +ORDER BY ts; + ++---------------------+--------------------+ +| ts | d | ++---------------------+--------------------+ +| 2026-07-08T00:00:00 | [{"e":{"f":0.1}}] | +| 2026-07-08T00:01:00 | [{"e":{"f":0.2}}] | +| 2026-07-08T00:02:00 | [{"e":{"g":-0.3}}] | +| 2026-07-08T00:03:00 | | +| 2026-07-08T00:04:00 | {"e":true} | +| 2026-07-08T00:05:00 | | ++---------------------+--------------------+ + +SHOW CREATE TABLE t_legacy_json2_variant_payload; + ++--------------------------------+---------------------------------------------------------------+ +| Table | Create Table | ++--------------------------------+---------------------------------------------------------------+ +| t_legacy_json2_variant_payload | CREATE TABLE IF NOT EXISTS "t_legacy_json2_variant_payload" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2 NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++--------------------------------+---------------------------------------------------------------+ diff --git a/tests/compatibility/cases/legacy_json2_variant_payload/verify.sql b/tests/compatibility/cases/legacy_json2_variant_payload/verify.sql new file mode 100644 index 00000000000..a3c2cab1b92 --- /dev/null +++ b/tests/compatibility/cases/legacy_json2_variant_payload/verify.sql @@ -0,0 +1,16 @@ +SELECT ts, j.a.b AS a_b, j.c AS c +FROM t_legacy_json2_variant_payload +ORDER BY ts; + +SELECT + ts, + j.d.e AS d_e +FROM t_legacy_json2_variant_payload_object_d +ORDER BY ts; + +-- JSON2 projection rendering is a user-visible compatibility contract. +SELECT ts, j.d AS d +FROM t_legacy_json2_variant_payload +ORDER BY ts; + +SHOW CREATE TABLE t_legacy_json2_variant_payload; diff --git a/tests/perf/AGENTS.md b/tests/perf/AGENTS.md index eb95544d22d..7f64b283bdc 100644 --- a/tests/perf/AGENTS.md +++ b/tests/perf/AGENTS.md @@ -11,14 +11,14 @@ docs together. - `[case]` is report metadata only. `[scenario]` is the executable regression configuration and must include `kind`, data layout, tables, queries, and - thresholds. Keep scenario parsing explicit in both - `tests/perf/query_regression_runner.py` and - `src/cmd/src/bin/query_perf_fixture.rs`. + thresholds. Rust owns case schema, defaults, validation, and normalized plan + output through `query_perf_fixture plan`; `tests/perf/query_regression_runner.py` + should only orchestrate from that normalized JSON. - Keep the direct-SST generator generic. Issue-specific behavior belongs in case files and thresholds, not in Rust generator logic. - Before pushing perf harness changes, run at least: - `uv run --no-project python -m py_compile .github/scripts/query-regression-run.py .github/scripts/query-regression-summary.py .github/scripts/query-regression-pr-metadata.py tests/perf/query_regression_runner.py` - `cargo fmt --all -- --check` - - `cargo build -p cmd --bin query_perf_fixture` + - `cargo build -p cmd --bin query_perf_fixture --features dev-tools` - dry-run the Python runner and Rust fixture generator against all built-in cases when the DSL or workflow case selection changes. diff --git a/tests/perf/README.md b/tests/perf/README.md index a59081b5681..91bc6c1637e 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -4,6 +4,11 @@ This directory is for query performance cases that compare a base build with a candidate build. It is not a replacement for sqlness: the goal is to measure the effect of optimizer/query-engine changes on realistic scan work. +The Prometheus remote-write cases are end-to-end regression coverage for the +automatic write, flush, SST, and query path. They do not replace controlled +encoding experiments that explicitly compare `plain`/dictionary, +`no_dictionary`, BYTE_STREAM_SPLIT, or Auto policies. + ## Phase 1: direct readable SST fixtures Phase 1 should generate data by writing readable Mito SST files and matching @@ -40,8 +45,10 @@ The runner also supports `scenario.kind = "prom_remote_write_then_query"` for a bounded write-path smoke/regression flow. This path is explicit and separate from the direct-SST fixture path: it starts the base and candidate distributed clusters, writes deterministic Prometheus remote-write v1 samples through -`/v1/prometheus/write`, flushes the configured physical metric table, then runs -the configured SQL/TQL queries against the logical metric table. +`/v1/prometheus/write`, flushes the configured physical metric table, checks +visibility, runs the configured SQL/TQL queries, stops and awaits the datanode, +then optionally inspects the landed SST/Parquet footer, encoding, and size and +runs the read-bench against the quiescent data directory. Remote-write cases configure one database, one logical metric, and one physical table under `[scenario.remote_write]`: @@ -52,34 +59,121 @@ kind = "prom_remote_write_then_query" [scenario.remote_write] database = "public" -metric = "prom_remote_write_smoke" +metric = "prom_remote_write_seeded_random" physical_table = "greptime_physical_table" -series_count = 8 -samples_per_series = 30 +series_count = 512 +samples_per_series = 1440 +sample_chunk_size = 480 +flush_every_sample_chunks = 1 start_unix_millis = 1_704_067_200_000 -step_millis = 15_000 -chunk_series_count = 4 -visibility_timeout_seconds = 30 +step_millis = 60_000 +chunk_series_count = 128 +timeout_seconds = 180 +visibility_timeout_seconds = 120 # Optional: split by time so each helper invocation writes only this many samples # per series, then periodically flush the physical metric table to produce # multiple time-interval SSTs. -# sample_chunk_size = 300 -# flush_every_sample_chunks = 1 [scenario.remote_write.prom_store] pending_rows_flush_interval = "1s" +max_batch_rows = 100000 ``` +Cases can optionally add `[scenario.remote_write.value]` to control the generated +sample values without changing label cardinality or the runner lifecycle: + +```toml +[scenario.remote_write.value] +pattern = "quantized_signal" # linear, constant, modulo, unique, seeded_random, + # run_length, quantized_signal, + # signal_with_sporadic_stalls, mixed_signal_repeated +base = 0.0 +step = 0.125 +cardinality = 4096 # buckets for modulo/seeded_random/quantized_signal/run_length +seed = 12345 # deterministic seeded_random input +run_length = 8 # adjacent samples per bucket for run_length/quantized_signal +stall_every = 100 # interval for signal_with_sporadic_stalls +stall_length = 16 # held samples inside each stall interval +mixed_every = 5 # every Nth sample becomes the repeated base value +``` + +The default `linear` pattern preserves the helper's historical formula. Use +`constant` or low-cardinality `modulo`/`seeded_random` values for repeated-value +data shapes, `run_length` for run-heavy low-cardinality series, `quantized_signal` +for signal-like values collapsed into a finite bucket set, `signal_with_sporadic_stalls` +for mostly continuous signals with periodic flat spots, and `mixed_signal_repeated` +for signal-plus-periodic-default mixtures. `unique` or high-cardinality buckets +still work for broad sample-value distributions. This is a generic sample-value +control for query/ingestion cases; it does not inspect or assert storage +encoding, Parquet footers, or storage policy choices. For chunked remote-write +ingestion, the runner passes the sample offset and total sample count to the +helper so non-linear value patterns use a stable global/per-series ordinal across +chunks. + +Case schema, value distribution defaults, storage defaults, and read-bench +defaults are owned by Rust. The Python runner calls +`query_perf_fixture plan --case ` and orchestrates the normalized JSON. +The same helper exposes `direct-sst`, `prom-remote-write`, and `inspect-footer` +subcommands; the old direct invocation (`query_perf_fixture --case ... --out-dir +...`) remains compatible. + +Remote-write cases that need to validate storage output can add +`[scenario.remote_write.storage]`. When present, the scenario body becomes: +remote-write → flush → visibility check → query measurement → stop and await +datanode → Parquet footer/size/encoding inspection → optional read-bench. Footer +inspection is handled by `query_perf_fixture inspect-footer`, which reads local +Parquet footers directly; it does not require Python `pyarrow`. + +```toml +[scenario.remote_write.storage] +inspect = true +column = "greptime_value" +include_metadata_files = false +# Optional: inspect below datanode data home instead of the whole fresh data dir. +# root_suffix = "greptime/public/" +min_files = 1 +min_files_with_column = 1 +require_encodings = ["BYTE_STREAM_SPLIT"] +forbid_encodings = ["PLAIN_DICTIONARY"] +max_total_file_size_bytes = 104857600 +max_column_compressed_size_bytes = 52428800 +max_column_uncompressed_size_bytes = 209715200 +max_candidate_total_file_size_regression_pct = 10.0 +max_candidate_column_compressed_size_regression_pct = 10.0 +max_candidate_column_uncompressed_size_regression_pct = 10.0 +``` + +The `*_pct` storage thresholds are percent regression limits comparing candidate +against base; `pct` means percent, not percentile. Per-target storage checks +(`min_files`, `min_files_with_column`, required/forbidden encodings, and absolute +`max_*_bytes`) run for both base and candidate. Comparative +`max_candidate_*_regression_pct` checks remain base-vs-candidate. These checks +are generic footer and byte-size assertions and do not encode product-specific +heuristics. When storage inspection is enabled, `min_files` and +`min_files_with_column` default to `1` even if omitted, so dry-runs show these +planned checks and an empty flush or missing target column fails the scenario. +By default the inspector root is the target datanode data home, which is suitable +for a fresh single-case data directory. For reused or more complex data homes, set +`root_suffix` to a path relative to the datanode data home to narrow inspection. + +When storage inspection is enabled, `[scenario.remote_write.read_bench]` defaults +to enabled with both datanode `parquetbench` and `scanbench`. Disable it with +`enabled = false`. `parquetbench` measures per-SST reader cost, `scanbench` +measures region scan cost, and query measurements still exercise the SQL/TQL +frontend path. Treat all performance conclusions as release-only; debug builds +are suitable only for command wiring and correctness checks. + The runner creates the configured database if needed, writes a per-target frontend config enabling `[prom_store]` with metric engine storage and a non-zero `pending_rows_flush_interval`, and validates that the logical metric table reaches `series_count * samples_per_series` rows before trusting the query measurements. -Use `--remote-write-generator /path/to/prom_remote_write_fixture` to provide the -Rust payload helper. `--fixture-only` is rejected for remote-write cases; use -`--dry-run` for planning. +Use `--fixture-generator /path/to/query_perf_fixture` to provide the Rust helper. +Deprecated `--remote-write-generator` and `--storage-inspector` options are kept +only for CLI compatibility. `--fixture-only` is rejected for remote-write cases; +use `--dry-run` for planning. Large manual remote-write cases can set `sample_chunk_size` to split ingestion by -time. For each chunk, the runner invokes `prom_remote_write_fixture` with the +time. For each chunk, the runner invokes `query_perf_fixture prom-remote-write` with the same series cardinality but a shorter `--samples-per-series` and an advanced `--start-unix-millis`. `flush_every_sample_chunks` controls periodic `ADMIN FLUSH_TABLE('')` calls; with `flush_every_sample_chunks = 1`, @@ -89,6 +183,25 @@ need multi-window file distribution should span multiple compaction windows. If `sample_chunk_size` is omitted, the runner keeps the older single-helper-invocation behavior and flushes once at the end. +The default remote-write coverage set contains four cases with the same 2048 +series × 14,400 samples (29,491,200 rows per target) shape. Each writes ten +one-day chunks, flushes every chunk, requires at least two visible SSTs with the +value column, runs seven read-bench iterations with `parquetbench` capped at +four SSTs while `scanbench` covers the landed region, and measures one ten-day +TQL selector with two warmups and 15 iterations: + +- `prom_remote_write_seeded_random`: high-distinct seeded-random values with + cardinality 29,491,200 and seed 8444. +- `prom_remote_write_run_heavy`: low-cardinality values in exact 16-sample runs. +- `prom_remote_write_mixed_every`: continuous signal values with every fifth + sample repeated at the base value. +- `prom_remote_write_integer_counter`: strictly increasing integer counter + values, with disjoint ranges for each series. + +These cases record normal storage/footer and read-bench results but do not gate +specific value encodings or storage-size outcomes. Use the controlled encoding +experiment matrix for those policy comparisons. + `tests/perf/query_cases/prom_remote_write_7913/case.toml` is a larger manual case for issue #7913. It writes 8192 series × 20160 samples through remote-write in 1440-sample daily time chunks, flushing after each chunk before running 1d/7d/14d @@ -246,7 +359,8 @@ as a fallback. Manifest files are reflinked or copied, not hardlinked. Fixture generator smoke test: ```bash -cargo run -p cmd --bin query_perf_fixture -- \ +cargo run -p cmd --bin query_perf_fixture --features dev-tools -- \ + direct-sst \ --case tests/perf/query_cases/smoke_direct_sst/case.toml \ --out-dir /tmp/query-perf-smoke ``` @@ -271,10 +385,10 @@ Remote-write runner dry-run: ```bash uv run --no-project python tests/perf/query_regression_runner.py \ - --case tests/perf/query_cases/prom_remote_write_smoke/case.toml \ + --case tests/perf/query_cases/prom_remote_write_seeded_random/case.toml \ --base-bin /path/to/base/greptime \ --candidate-bin /path/to/candidate/greptime \ - --remote-write-generator /path/to/prom_remote_write_fixture \ + --fixture-generator /path/to/query_perf_fixture \ --work-dir /tmp/query-perf-remote-write \ --dry-run ``` @@ -289,14 +403,22 @@ query regression runs. It builds its own binaries for now: candidate checkout - runner and summary formatter from the candidate checkout -The workflow runs automatically only for non-draft PRs labeled -`query-regression` (on label/ready-for-review/reopen events, not every push). +The workflow builds base and candidate `greptime` as normal release-equivalent +binaries. Candidate `query_perf_fixture` is the extra head-side helper binary; +the runner uses candidate `greptime datanode parquetbench/scanbench` as the +read-bench tool against each target's data directory. + +The workflow runs automatically only when `query-regression` is added to a +non-draft PR; it does not rerun on pushes, ready-for-review, or reopen events. PR runs build base/candidate once and then run the default case set with `--allow-large-fixture`. Manual `workflow_dispatch` runs can pass `all`, one case path, or a comma/whitespace-separated list of case paths, and can override refs. -It always uploads `query-regression-work/**` and `query-regression-summary.md`, -writes the Markdown summary to the workflow step summary, and updates a sticky PR -comment through the trusted follow-up workflow. +The main report artifact uploads only aggregate/per-target JSON reports, +component logs, and `query-regression-summary.md` with seven-day retention; +fixture data, SSTs, and cluster state are excluded. PR runs also upload a +separate trusted-comment artifact containing PR metadata and aggregate reports. +The workflow writes the Markdown summary to the workflow step summary and +updates a sticky PR comment through the trusted follow-up workflow. ## Built-in cases diff --git a/tests/perf/query_cases/prom_remote_write_integer_counter/case.toml b/tests/perf/query_cases/prom_remote_write_integer_counter/case.toml new file mode 100644 index 00000000000..3ac64a94b9d --- /dev/null +++ b/tests/perf/query_cases/prom_remote_write_integer_counter/case.toml @@ -0,0 +1,48 @@ +[case] +name = "prom_remote_write_integer_counter" +description = "End-to-end Prometheus remote-write regression with strictly increasing integer counter values" + +[scenario] +kind = "prom_remote_write_then_query" + +[scenario.remote_write] +database = "public" +metric = "prom_remote_write_integer_counter" +physical_table = "greptime_physical_table" +series_count = 2048 +samples_per_series = 14400 +sample_chunk_size = 1440 +flush_every_sample_chunks = 1 +start_unix_millis = 1704067200000 +step_millis = 60000 +chunk_series_count = 256 +timeout_seconds = 600 +visibility_timeout_seconds = 300 + +[scenario.remote_write.value] +pattern = "unique" +base = 0 +step = 1 + +[scenario.remote_write.prom_store] +pending_rows_flush_interval = "1s" +max_batch_rows = 100000 + +[scenario.remote_write.storage] +column = "greptime_value" +min_files = 2 +min_files_with_column = 2 + +[scenario.remote_write.read_bench] +iterations = 7 +max_files = 4 + +[[scenario.queries]] +name = "selector_10d_integer_counter" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704067200, 1704931200, '1h') prom_remote_write_integer_counter{host=~'host.*'}" +warmup = 2 +iterations = 15 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 75 diff --git a/tests/perf/query_cases/prom_remote_write_mixed_every/case.toml b/tests/perf/query_cases/prom_remote_write_mixed_every/case.toml new file mode 100644 index 00000000000..9ba39f2fbed --- /dev/null +++ b/tests/perf/query_cases/prom_remote_write_mixed_every/case.toml @@ -0,0 +1,49 @@ +[case] +name = "prom_remote_write_mixed_every" +description = "End-to-end Prometheus remote-write regression with periodic repeated signal values" + +[scenario] +kind = "prom_remote_write_then_query" + +[scenario.remote_write] +database = "public" +metric = "prom_remote_write_mixed_every" +physical_table = "greptime_physical_table" +series_count = 2048 +samples_per_series = 14400 +sample_chunk_size = 1440 +flush_every_sample_chunks = 1 +start_unix_millis = 1704067200000 +step_millis = 60000 +chunk_series_count = 256 +timeout_seconds = 600 +visibility_timeout_seconds = 300 + +[scenario.remote_write.value] +pattern = "mixed_signal_repeated" +base = 0 +step = 0.125 +mixed_every = 5 + +[scenario.remote_write.prom_store] +pending_rows_flush_interval = "1s" +max_batch_rows = 100000 + +[scenario.remote_write.storage] +column = "greptime_value" +min_files = 2 +min_files_with_column = 2 + +[scenario.remote_write.read_bench] +iterations = 7 +max_files = 4 + +[[scenario.queries]] +name = "selector_10d_mixed_every" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704067200, 1704931200, '1h') prom_remote_write_mixed_every{host=~'host.*'}" +warmup = 2 +iterations = 15 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 75 diff --git a/tests/perf/query_cases/prom_remote_write_run_heavy/case.toml b/tests/perf/query_cases/prom_remote_write_run_heavy/case.toml new file mode 100644 index 00000000000..243892e42b1 --- /dev/null +++ b/tests/perf/query_cases/prom_remote_write_run_heavy/case.toml @@ -0,0 +1,50 @@ +[case] +name = "prom_remote_write_run_heavy" +description = "End-to-end Prometheus remote-write regression with repeated 16-sample value runs" + +[scenario] +kind = "prom_remote_write_then_query" + +[scenario.remote_write] +database = "public" +metric = "prom_remote_write_run_heavy" +physical_table = "greptime_physical_table" +series_count = 2048 +samples_per_series = 14400 +sample_chunk_size = 1440 +flush_every_sample_chunks = 1 +start_unix_millis = 1704067200000 +step_millis = 60000 +chunk_series_count = 256 +timeout_seconds = 600 +visibility_timeout_seconds = 300 + +[scenario.remote_write.value] +pattern = "run_length" +base = 0 +step = 1 +cardinality = 1024 +run_length = 16 + +[scenario.remote_write.prom_store] +pending_rows_flush_interval = "1s" +max_batch_rows = 100000 + +[scenario.remote_write.storage] +column = "greptime_value" +min_files = 2 +min_files_with_column = 2 + +[scenario.remote_write.read_bench] +iterations = 7 +max_files = 4 + +[[scenario.queries]] +name = "selector_10d_run_heavy" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704067200, 1704931200, '1h') prom_remote_write_run_heavy{host=~'host.*'}" +warmup = 2 +iterations = 15 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 75 diff --git a/tests/perf/query_cases/prom_remote_write_seeded_random/case.toml b/tests/perf/query_cases/prom_remote_write_seeded_random/case.toml new file mode 100644 index 00000000000..6a63edbcbbf --- /dev/null +++ b/tests/perf/query_cases/prom_remote_write_seeded_random/case.toml @@ -0,0 +1,48 @@ +[case] +name = "prom_remote_write_seeded_random" +description = "End-to-end Prometheus remote-write regression with high-distinct seeded-random values" + +[scenario] +kind = "prom_remote_write_then_query" + +[scenario.remote_write] +database = "public" +metric = "prom_remote_write_seeded_random" +physical_table = "greptime_physical_table" +series_count = 2048 +samples_per_series = 14400 +sample_chunk_size = 1440 +flush_every_sample_chunks = 1 +start_unix_millis = 1704067200000 +step_millis = 60000 +chunk_series_count = 256 +timeout_seconds = 600 +visibility_timeout_seconds = 300 + +[scenario.remote_write.value] +pattern = "seeded_random" +cardinality = 29_491_200 +seed = 8444 + +[scenario.remote_write.prom_store] +pending_rows_flush_interval = "1s" +max_batch_rows = 100000 + +[scenario.remote_write.storage] +column = "greptime_value" +min_files = 2 +min_files_with_column = 2 + +[scenario.remote_write.read_bench] +iterations = 7 +max_files = 4 + +[[scenario.queries]] +name = "selector_10d_seeded_random" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704067200, 1704931200, '1h') prom_remote_write_seeded_random{host=~'host.*'}" +warmup = 2 +iterations = 15 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 75 diff --git a/tests/perf/query_cases/prom_remote_write_smoke/case.toml b/tests/perf/query_cases/prom_remote_write_smoke/case.toml deleted file mode 100644 index b946fbc8f51..00000000000 --- a/tests/perf/query_cases/prom_remote_write_smoke/case.toml +++ /dev/null @@ -1,31 +0,0 @@ -[case] -name = "prom_remote_write_smoke" -description = "Tiny Prometheus remote-write ingestion then TQL query smoke case" - -[scenario] -kind = "prom_remote_write_then_query" - -[scenario.remote_write] -database = "public" -metric = "prom_remote_write_smoke" -physical_table = "greptime_physical_table" -series_count = 8 -samples_per_series = 30 -start_unix_millis = 1_704_067_200_000 -step_millis = 15_000 -chunk_series_count = 4 -timeout_seconds = 30 - -[scenario.remote_write.prom_store] -pending_rows_flush_interval = "1s" -max_batch_rows = 10000 - -[[scenario.queries]] -name = "selector_5m_smoke" -kind = "tql" -query = "TQL ANALYZE VERBOSE (1704067200, 1704067650, '15s') prom_remote_write_smoke{host=~'host.*'}" -warmup = 0 -iterations = 2 - -[scenario.queries.thresholds] -max_candidate_latency_regression_pct = 100 diff --git a/tests/perf/query_regression_runner.py b/tests/perf/query_regression_runner.py index a7f2646ea8b..f5c9628b83d 100644 --- a/tests/perf/query_regression_runner.py +++ b/tests/perf/query_regression_runner.py @@ -66,7 +66,8 @@ def parse_args() -> argparse.Namespace: p.add_argument("--base-bin", required=True, type=Path) p.add_argument("--candidate-bin", required=True, type=Path) p.add_argument("--fixture-generator", type=Path) - p.add_argument("--remote-write-generator", type=Path, help="prom_remote_write_fixture helper binary") + p.add_argument("--remote-write-generator", type=Path, help="deprecated: use --fixture-generator query_perf_fixture") + p.add_argument("--storage-inspector", type=Path, help="deprecated: use --fixture-generator query_perf_fixture") p.add_argument("--work-dir", required=True, type=Path) p.add_argument("--fixture-cache-dir", type=Path, help="persistent directory for generated fixtures, keyed by case content") p.add_argument("--reuse-fixture", action="store_true") @@ -83,6 +84,16 @@ def load_case(path: Path) -> dict[str, Any]: return tomllib.load(f) +def load_normalized_case(case_path: Path, fixture_generator: Path | None) -> dict[str, Any]: + if fixture_generator is None: + raise ValueError("--fixture-generator query_perf_fixture is required for Rust-owned case planning") + result = run_command([str(fixture_generator), "plan", "--case", str(case_path)]) + if result["returncode"] != 0: + raise RuntimeError(f"query_perf_fixture plan failed: {result['stderr'][:2000]}") + plan = json.loads(result["stdout"]) + return {"case": load_case(case_path).get("case", {}), "scenario": plan["scenario"], "schema_version": plan.get("schema_version")} + + def require_binary(path: Path, name: str, *, dry_run: bool) -> None: if dry_run: return @@ -157,11 +168,9 @@ def scenario(case: dict[str, Any]) -> dict[str, Any]: def case_tables(case: dict[str, Any]) -> list[dict[str, Any]]: value = scenario(case) if value.get("kind") == "prom_remote_write_then_query": - remote = value.get("remote_write") or {} - metric = remote.get("metric") or remote.get("metric_name") - database = remote.get("database", "public") - if not metric: - raise ValueError("remote-write scenario requires scenario.remote_write.metric") + remote = value["remote_write"] + metric = remote["metric"] + database = remote["database"] return [{"database": database, "name": metric, "engine": "metric", "validate_show_create_engine": False}] tables = value.get("tables") or [] if not tables or value.get("layout", {}).get("regions") != 1: @@ -548,7 +557,7 @@ def generate_fixture(generator: Path | None, case_path: Path, fixture_dir: Path, if fixture_dir.exists() and not reuse_fixture: shutil.rmtree(fixture_dir) fixture_dir.mkdir(parents=True, exist_ok=True) - cmd = [str(generator), "--case", str(case_path), "--out-dir", str(fixture_dir)] + cmd = [str(generator), "direct-sst", "--case", str(case_path), "--out-dir", str(fixture_dir)] if table_name is not None: cmd += ["--table", table_name] if region_id is not None: @@ -777,6 +786,243 @@ def enforce_thresholds(case: dict[str, Any], base: dict[str, Any], candidate: di return results +def storage_config(remote: dict[str, Any]) -> dict[str, Any] | None: + value = remote["storage"] + if not isinstance(value, dict): + return None + if value["inspect"] is False: + return None + return value + + +def run_storage_inspection(helper: Path | None, target: RunTarget, storage: dict[str, Any], *, dry_run: bool) -> dict[str, Any]: + if helper is None: + if not dry_run: + raise ValueError("--fixture-generator query_perf_fixture is required when storage inspection is enabled") + helper = Path("query_perf_fixture") + root = target.datanode_data_dir + if storage.get("root_suffix"): + root = root / str(storage["root_suffix"]) + cmd = [ + str(helper), + "inspect-footer", + "--root", str(root), + "--column", str(storage["column"]), + ] + if storage["include_metadata_files"]: + cmd.append("--include-metadata-files") + if dry_run: + return {"status": "dry-run", "cmd": cmd, "root": str(root)} + result = run_command(cmd) + result["status"] = "ok" if result["returncode"] == 0 else "failed" + if result["returncode"] != 0: + raise RuntimeError(f"storage inspector failed for {target.name}: {result['stderr'][:2000]}") + try: + result["summary"] = json.loads(result["stdout"]) + except json.JSONDecodeError: + result["summary_parse_error"] = result["stdout"] + result["status"] = "failed" + return result + + +def storage_summary(inspection: dict[str, Any]) -> dict[str, Any]: + summary = inspection.get("summary") + if isinstance(summary, dict): + nested = summary.get("summary") + if isinstance(nested, dict): + return nested + return {} + + +def enforce_storage_thresholds(storage: dict[str, Any] | None, base: dict[str, Any] | None, candidate: dict[str, Any] | None) -> list[dict[str, Any]]: + if not storage: + return [] + results: list[dict[str, Any]] = [] + target_summaries = [("base", storage_summary(base or {})), ("candidate", storage_summary(candidate or {}))] + cand = target_summaries[1][1] + base_summary = storage_summary(base or {}) + + def add_abs(name: str, field: str) -> None: + limit = storage.get(name) + if limit is None: + return + for target_name, summary in target_summaries: + actual = summary.get(field) + ok = actual is not None and float(actual) <= float(limit) + results.append({"target": target_name, "threshold": name, "status": "passed" if ok else "failed", "actual": actual, "limit": limit}) + + def add_min(name: str, field: str) -> None: + limit = storage[name] + if limit is None: + return + for target_name, summary in target_summaries: + actual = summary.get(field) + ok = actual is not None and int(actual) >= int(limit) + results.append({"target": target_name, "threshold": name, "status": "passed" if ok else "failed", "actual": actual, "limit": limit}) + + def add_cmp(name: str, field: str) -> None: + limit = storage.get(name) + if limit is None: + return + base_value = base_summary.get(field) + candidate_value = cand.get(field) + if base_value in (None, 0) or candidate_value is None: + results.append({"threshold": name, "status": "failed", "reason": "missing or zero base/candidate value", "base": base_value, "candidate": candidate_value}) + return + actual = (float(candidate_value) - float(base_value)) / float(base_value) * 100.0 + results.append({"threshold": name, "status": "passed" if actual <= float(limit) else "failed", "actual_pct": actual, "limit_pct": limit, "base": base_value, "candidate": candidate_value}) + + add_min("min_files", "file_count") + add_min("min_files_with_column", "files_with_column") + add_abs("max_total_file_size_bytes", "total_file_size") + add_abs("max_column_compressed_size_bytes", "column_compressed_size") + add_abs("max_column_uncompressed_size_bytes", "column_uncompressed_size") + for target_name, summary in target_summaries: + encodings = set(summary.get("unique_encodings") or []) + for required in storage["require_encodings"]: + results.append({"target": target_name, "threshold": "require_encodings", "encoding": required, "status": "passed" if required in encodings else "failed"}) + for forbidden in storage["forbid_encodings"]: + results.append({"target": target_name, "threshold": "forbid_encodings", "encoding": forbidden, "status": "passed" if forbidden not in encodings else "failed"}) + add_cmp("max_candidate_total_file_size_regression_pct", "total_file_size") + add_cmp("max_candidate_column_compressed_size_regression_pct", "column_compressed_size") + add_cmp("max_candidate_column_uncompressed_size_regression_pct", "column_uncompressed_size") + return results + + +def planned_storage_thresholds(storage: dict[str, Any] | None) -> list[dict[str, Any]]: + if not storage: + return [] + return list(storage["planned_thresholds"]) + + +REGION_DIR_RE = re.compile(r"^(\d+)_(\d{10})$") + + +def inspection_report(inspection: dict[str, Any] | None) -> dict[str, Any]: + if not inspection: + return {} + report = inspection.get("summary") + return report if isinstance(report, dict) else {} + + +def inspected_relative_path(target: RunTarget, inspection: dict[str, Any], relative_path: str) -> Path: + report = inspection_report(inspection) + root_text = report.get("root") or inspection.get("root") + if not root_text: + return Path(relative_path) + root = Path(root_text).resolve(strict=False) + data_home = target.datanode_data_dir.resolve(strict=False) + try: + root_suffix = root.relative_to(data_home) + except ValueError as e: + raise RuntimeError(f"storage inspection root {root} is not under datanode data home {data_home}") from e + return root_suffix / relative_path + + +def parse_sst_bench_target(target: RunTarget, inspection: dict[str, Any], file: dict[str, Any], *, dry_run: bool) -> dict[str, str] | None: + relative_path = str(file.get("relative_path") or "") + if dry_run and (not relative_path or "<" in relative_path): + return { + "relative_path": relative_path or "/_/data/.parquet", + "table_dir": "/", + "region_id": ":", + "path_type": "data", + "file_id": "", + } + if not relative_path: + return None + full_relative = inspected_relative_path(target, inspection, relative_path) + parts = full_relative.parts + if not parts or not parts[-1].endswith(".parquet"): + return None + region_index = next((idx for idx, part in enumerate(parts) if REGION_DIR_RE.match(part)), None) + if region_index is None or region_index == 0: + if dry_run: + return { + "relative_path": relative_path, + "table_dir": "/", + "region_id": ":", + "path_type": "data", + "file_id": Path(parts[-1]).stem, + } + return None + region_match = REGION_DIR_RE.match(parts[region_index]) + assert region_match is not None + file_index = len(parts) - 1 + path_type = "bare" + if region_index + 1 < file_index and parts[region_index + 1] in ("data", "metadata"): + path_type = parts[region_index + 1] + if path_type == "metadata": + return None + table_dir = "/".join(parts[:region_index]).rstrip("/") + "/" + return { + "relative_path": str(full_relative), + "table_dir": table_dir, + "region_id": f"{int(region_match.group(1))}:{int(region_match.group(2))}", + "path_type": path_type, + "file_id": Path(parts[-1]).stem, + } + + +def run_read_bench(bench_binary: Path, target: RunTarget, read_bench: dict[str, Any] | None, storage_inspection: dict[str, Any] | None, *, dry_run: bool) -> dict[str, Any]: + if not read_bench or read_bench["enabled"] is False: + return {"status": "skipped", "reason": "read_bench disabled"} + run_parquetbench = read_bench["parquetbench"] + run_scanbench = read_bench["scanbench"] + if not run_parquetbench and not run_scanbench: + return {"status": "skipped", "reason": "parquetbench and scanbench disabled"} + if storage_inspection is None: + raise ValueError("read_bench requires storage inspection") + bench_dir = target.work_dir / "read_bench" + config_toml = bench_dir / "bench.toml" + scan_json = bench_dir / "scan.json" + files = inspection_report(storage_inspection).get("files") or [] + ssts = [f for f in files if f.get("relative_path", "").endswith(".parquet") and f.get("columns")] + if dry_run and not ssts: + ssts = [{"relative_path": ".parquet"}] + if read_bench["max_files"] is not None: + ssts = ssts[: int(read_bench["max_files"])] + bench_targets = [parsed for f in ssts if (parsed := parse_sst_bench_target(target, storage_inspection, f, dry_run=dry_run)) is not None] + if not bench_targets: + return {"status": "failed", "reason": "no inspected data SST files available for read_bench"} + config_text = f'[storage]\ndata_home = "{target.datanode_data_dir}"\ntype = "File"\n\n[[region_engine]]\n[region_engine.mito]\n' + commands = [] + parquet_runs = [] + scan_runs = [] + iterations = str(int(read_bench["iterations"])) + if run_parquetbench: + for bench_target in bench_targets: + cmd = [str(bench_binary), "datanode", "parquetbench", "--config", str(config_toml), "--region-id", bench_target["region_id"], "--table-dir", bench_target["table_dir"], "--file-id", bench_target["file_id"], "--scan-config", str(scan_json), "--path-type", bench_target["path_type"], "--iterations", iterations, "--reader", str(read_bench["parquet_reader"])] + commands.append(cmd) + parquet_runs.append({**bench_target, "cmd": cmd, "status": "dry-run" if dry_run else "planned"}) + regions: dict[tuple[str, str, str], list[str]] = {} + if run_scanbench: + for bench_target in bench_targets: + regions.setdefault((bench_target["table_dir"], bench_target["region_id"], bench_target["path_type"]), []).append(bench_target["relative_path"]) + for (table_dir, region_id, path_type), region_files in regions.items(): + cmd = [str(bench_binary), "datanode", "scanbench", "--config", str(config_toml), "--region-id", region_id, "--table-dir", table_dir, "--scan-config", str(scan_json), "--path-type", path_type, "--scanner", str(read_bench["scan_scanner"]), "--parallelism", str(int(read_bench["parallelism"])), "--iterations", iterations] + commands.append(cmd) + scan_runs.append({"table_dir": table_dir, "region_id": region_id, "path_type": path_type, "files": region_files, "cmd": cmd, "status": "dry-run" if dry_run else "planned"}) + if dry_run: + return {"status": "dry-run", "config_path": str(config_toml), "scan_config_path": str(scan_json), "commands": commands, "parquetbench": parquet_runs, "scanbench": scan_runs} + bench_dir.mkdir(parents=True, exist_ok=True) + config_toml.write_text(config_text) + scan_json.write_text(json.dumps({"projection_names": read_bench["projection"]}, indent=2) + "\n") + avg_re = re.compile(r"Average duration[^0-9]*([0-9.]+)\s*ms", re.I) + for runs in (parquet_runs, scan_runs): + for run in runs: + result = run_command(run["cmd"]) + run.update(result) + run["status"] = "ok" if result["returncode"] == 0 else "failed" + m = avg_re.search(result.get("stdout", "")) + if m: + run["average_ms"] = float(m.group(1)) + def median(runs: list[dict[str, Any]]) -> float | None: + vals = [r["average_ms"] for r in runs if "average_ms" in r] + return statistics.median(vals) if vals else None + return {"status": "ok" if all(r.get("status") == "ok" for r in parquet_runs + scan_runs) else "failed", "config_path": str(config_toml), "scan_config_path": str(scan_json), "parquetbench": parquet_runs, "scanbench": scan_runs, "aggregate": {"parquetbench_median_average_ms": median(parquet_runs), "scanbench_median_average_ms": median(scan_runs)}} + + def write_json(path: Path, payload: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") @@ -791,18 +1037,18 @@ def require_fresh_work_dirs(targets: list[RunTarget], *, reuse_work_dir: bool, d def write_frontend_prom_config(target: RunTarget, remote: dict[str, Any]) -> Path: - prom = remote.get("prom_store") or {} + prom = remote["prom_store"] path = target.work_dir / "frontend-prom-store.toml" path.parent.mkdir(parents=True, exist_ok=True) path.write_text( "[prom_store]\n" "enable = true\n" "with_metric_engine = true\n" - f"pending_rows_flush_interval = {json.dumps(str(prom.get('pending_rows_flush_interval', '1s')))}\n" - f"max_batch_rows = {int(prom.get('max_batch_rows', 100000))}\n" - f"max_concurrent_flushes = {int(prom.get('max_concurrent_flushes', 256))}\n" - f"worker_channel_capacity = {int(prom.get('worker_channel_capacity', 65526))}\n" - f"max_inflight_requests = {int(prom.get('max_inflight_requests', 3000))}\n" + f"pending_rows_flush_interval = {json.dumps(str(prom['pending_rows_flush_interval']))}\n" + f"max_batch_rows = {int(prom['max_batch_rows'])}\n" + f"max_concurrent_flushes = {int(prom['max_concurrent_flushes'])}\n" + f"worker_channel_capacity = {int(prom['worker_channel_capacity'])}\n" + f"max_inflight_requests = {int(prom['max_inflight_requests'])}\n" ) return path @@ -810,23 +1056,38 @@ def write_frontend_prom_config(target: RunTarget, remote: dict[str, Any]) -> Pat def run_remote_write(generator: Path | None, target: RunTarget, remote: dict[str, Any], *, dry_run: bool) -> dict[str, Any]: if generator is None: if not dry_run: - raise ValueError("--remote-write-generator is required for prom_remote_write_then_query unless --dry-run is set") - generator = Path("prom_remote_write_fixture") - metric = remote.get("metric") or remote.get("metric_name") - physical_table = remote.get("physical_table", "greptime_physical_table") + raise ValueError("--fixture-generator query_perf_fixture is required for prom_remote_write_then_query") + generator = Path("query_perf_fixture") + metric = remote["metric"] + physical_table = remote["physical_table"] cmd = [ str(generator), + "prom-remote-write", "--endpoint", f"http://127.0.0.1:{target.http_port}/v1/prometheus/write", - "--database", remote.get("database", "public"), + "--database", remote["database"], "--metric", metric, "--physical-table", physical_table, - "--series-count", str(int(remote.get("series_count", 8))), - "--samples-per-series", str(int(remote.get("samples_per_series", 30))), - "--start-unix-millis", str(int(remote.get("start_unix_millis", 1_704_067_200_000))), - "--step-millis", str(int(remote.get("step_millis", 15_000))), - "--chunk-series-count", str(int(remote.get("chunk_series_count", remote.get("batch_size", 8)))), - "--timeout-seconds", str(int(remote.get("timeout_seconds", 60))), + "--series-count", str(int(remote["series_count"])), + "--samples-per-series", str(int(remote["samples_per_series"])), + "--start-unix-millis", str(int(remote["start_unix_millis"])), + "--step-millis", str(int(remote["step_millis"])), + "--chunk-series-count", str(int(remote["chunk_series_count"])), + "--timeout-seconds", str(int(remote["timeout_seconds"])), ] + value = remote["value"] + cmd.extend(["--value-pattern", str(value["pattern"])]) + cmd.extend(["--value-base", str(value["base"])]) + cmd.extend(["--value-step", str(value["step"])]) + cmd.extend(["--value-cardinality", str(int(value["cardinality"]))]) + cmd.extend(["--value-seed", str(int(value["seed"]))]) + cmd.extend(["--value-run-length", str(int(value["run_length"]))]) + cmd.extend(["--value-stall-every", str(int(value["stall_every"]))]) + cmd.extend(["--value-stall-length", str(int(value["stall_length"]))]) + cmd.extend(["--value-mixed-every", str(int(value["mixed_every"]))]) + if "sample_offset" in remote: + cmd.extend(["--value-sample-offset", str(int(remote["sample_offset"]))]) + if "total_samples_per_series" in remote: + cmd.extend(["--value-total-samples-per-series", str(int(remote["total_samples_per_series"]))]) if dry_run: return {"status": "dry-run", "cmd": cmd} result = run_command(cmd) @@ -862,23 +1123,23 @@ def flush_remote_physical_table(target: RunTarget, db: str, physical_table: str, def run_remote_write_ingestion(generator: Path | None, target: RunTarget, remote: dict[str, Any], args: argparse.Namespace, *, dry_run: bool) -> tuple[dict[str, Any], list[dict[str, Any]]]: - sample_chunk_size = remote.get("sample_chunk_size") - db = remote.get("database", "public") - physical_table = remote.get("physical_table", "greptime_physical_table") + sample_chunk_size = remote["sample_chunk_size"] + db = remote["database"] + physical_table = remote["physical_table"] if sample_chunk_size is None: rw = run_remote_write(generator, target, remote, dry_run=dry_run) flush = flush_remote_physical_table(target, db, physical_table, args, dry_run=dry_run, reason="final") return rw, [flush] - total_samples = int(remote.get("samples_per_series", 30)) + total_samples = int(remote["samples_per_series"]) chunk_samples = int(sample_chunk_size) if chunk_samples <= 0: raise ValueError("scenario.remote_write.sample_chunk_size must be positive") - flush_every = int(remote.get("flush_every_sample_chunks", 1)) + flush_every = int(remote["flush_every_sample_chunks"]) if flush_every <= 0: raise ValueError("scenario.remote_write.flush_every_sample_chunks must be positive") - start = int(remote.get("start_unix_millis", 1_704_067_200_000)) - step = int(remote.get("step_millis", 15_000)) + start = int(remote["start_unix_millis"]) + step = int(remote["step_millis"]) chunks: list[dict[str, Any]] = [] flushes: list[dict[str, Any]] = [] chunk_index = 0 @@ -889,6 +1150,8 @@ def run_remote_write_ingestion(generator: Path | None, target: RunTarget, remote chunk_remote = dict(remote) chunk_remote["samples_per_series"] = current chunk_remote["start_unix_millis"] = start + offset * step + chunk_remote["sample_offset"] = offset + chunk_remote["total_samples_per_series"] = total_samples result = run_remote_write(generator, target, chunk_remote, dry_run=dry_run) result["sample_offset"] = offset result["samples_per_series"] = current @@ -901,8 +1164,8 @@ def run_remote_write_ingestion(generator: Path | None, target: RunTarget, remote flushes.append(flush_remote_physical_table(target, db, physical_table, args, dry_run=dry_run, reason="final", chunk_index=chunk_index)) aggregate = summarize_remote_write_chunks(chunks) if dry_run: - series_count = int(remote.get("series_count", 8)) - chunk_series_count = int(remote.get("chunk_series_count", remote.get("batch_size", 8))) + series_count = int(remote["series_count"]) + chunk_series_count = int(remote["chunk_series_count"]) aggregate["rows"] = series_count * total_samples aggregate["samples_written"] = aggregate["rows"] aggregate["batches"] = chunk_index * ((series_count + chunk_series_count - 1) // chunk_series_count) @@ -917,7 +1180,7 @@ def run_remote_write_ingestion(generator: Path | None, target: RunTarget, remote def expected_remote_write_rows(remote: dict[str, Any]) -> int: - return int(remote.get("series_count", 8)) * int(remote.get("samples_per_series", 30)) + return int(remote["series_count"]) * int(remote["samples_per_series"]) def extract_count_value(result: dict[str, Any]) -> int | None: @@ -966,16 +1229,21 @@ def poll_expected_count(target: RunTarget, table_name: str, db: str, expected_ro def run_remote_write_scenario(args: argparse.Namespace, case: dict[str, Any], case_path: Path, targets: list[RunTarget], report: dict[str, Any]) -> None: - remote = scenario(case).get("remote_write") or {} + remote = scenario(case)["remote_write"] tables = case_tables(case) if args.fixture_only: raise ValueError("--fixture-only is not supported for prom_remote_write_then_query; use --dry-run for planning") - if args.remote_write_generator is not None: - require_binary(args.remote_write_generator, "remote-write generator", dry_run=args.dry_run) + helper = args.fixture_generator or args.remote_write_generator + if helper is not None: + require_binary(helper, "query_perf_fixture", dry_run=args.dry_run) elif not args.dry_run: - raise ValueError("--remote-write-generator is required for prom_remote_write_then_query") + raise ValueError("--fixture-generator is required for prom_remote_write_then_query") + storage = storage_config(remote) + if storage and helper is None and not args.dry_run: + raise ValueError("--fixture-generator is required when scenario.remote_write.storage.inspect is enabled") clusters: list[DistributedCluster] = [] target_results = [] + storage_results = [] try: for target in targets: target.work_dir.mkdir(parents=True, exist_ok=True) @@ -984,28 +1252,40 @@ def run_remote_write_scenario(args: argparse.Namespace, case: dict[str, Any], ca clusters.append(cluster) if not args.dry_run: cluster.start_all(config_path) - db = remote.get("database", "public") + db = remote["database"] create_database = {"status": "dry-run", "database": db} if not args.dry_run: create_database = http_post_sql(target.http_port, f"CREATE DATABASE IF NOT EXISTS {sql_ident(db)}", "public", args.http_timeout) - rw, flushes = run_remote_write_ingestion(args.remote_write_generator, target, remote, args, dry_run=args.dry_run) - physical = remote.get("physical_table", "greptime_physical_table") + rw, flushes = run_remote_write_ingestion(helper, target, remote, args, dry_run=args.dry_run) visibility = {"status": "dry-run"} query_result = {"validation": [], "validation_errors": [], "measurements": [], "status": "planned"} if not args.dry_run: - visibility = poll_expected_count(target, tables[0]["name"], db, expected_remote_write_rows(remote), float(remote.get("visibility_timeout_seconds", 30)), args.http_timeout) + visibility = poll_expected_count(target, tables[0]["name"], db, expected_remote_write_rows(remote), float(remote["visibility_timeout_seconds"]), args.http_timeout) query_result = run_queries(target, case, tables, args.http_timeout) + cluster.stop_component("datanode") + storage_inspection = None + if storage: + storage_inspection = run_storage_inspection(helper or args.storage_inspector, target, storage, dry_run=args.dry_run) + storage_results.append(storage_inspection) + read_bench_result = run_read_bench(args.candidate_bin, target, remote["read_bench"], storage_inspection, dry_run=args.dry_run) tr = {"name": target.name, "binary": str(target.binary), "work_dir": str(target.work_dir), "components": cluster.component_report(), "frontend_config": str(config_path), "create_database": create_database, "remote_write": rw, "flushes": flushes, "flush": flushes[-1] if flushes else None, "visibility": visibility, **query_result} + if storage_inspection is not None: + tr["storage_inspection"] = storage_inspection + tr["read_bench"] = read_bench_result flushes_ok = all(flush.get("ok") for flush in flushes) - remote_checks_ok = args.dry_run or (create_database.get("ok") and flushes_ok and visibility.get("ok") and visibility.get("row_count_ok")) + storage_ok = storage_inspection is None or args.dry_run or storage_inspection.get("status") == "ok" + read_bench_ok = args.dry_run or read_bench_result.get("status") in ("ok", "skipped") + remote_checks_ok = args.dry_run or (create_database.get("ok") and flushes_ok and visibility.get("ok") and visibility.get("row_count_ok") and storage_ok and read_bench_ok) if not remote_checks_ok: - tr.setdefault("validation_errors", []).append({"phase": "remote_write_visibility", "create_database_ok": create_database.get("ok"), "flushes_ok": flushes_ok, "visibility_ok": visibility.get("ok"), "row_count_ok": visibility.get("row_count_ok"), "create_database": create_database, "flushes": flushes, "visibility": visibility}) + tr.setdefault("validation_errors", []).append({"phase": "remote_write_visibility_storage", "create_database_ok": create_database.get("ok"), "flushes_ok": flushes_ok, "visibility_ok": visibility.get("ok"), "row_count_ok": visibility.get("row_count_ok"), "storage_ok": storage_ok, "read_bench_ok": read_bench_ok, "create_database": create_database, "flushes": flushes, "visibility": visibility, "read_bench": read_bench_result}) tr["status"] = "planned" if args.dry_run else ("measured" if remote_checks_ok and query_result["status"] == "ok" else "failed") write_json(target.report_path, tr) report["targets"].append(tr) target_results.append(query_result) if not args.dry_run: - report["thresholds"] = enforce_thresholds(case, target_results[0], target_results[1]) + report["thresholds"] = enforce_thresholds(case, target_results[0], target_results[1]) + enforce_storage_thresholds(storage, storage_results[0] if storage_results else None, storage_results[1] if len(storage_results) > 1 else None) + elif storage: + report["thresholds"] = planned_storage_thresholds(storage) report["status"] = "planned" if args.dry_run else ("failed" if any(t["status"] == "failed" for t in report["thresholds"]) or any(t.get("status") == "failed" for t in report["targets"]) else "ok") finally: for cluster in reversed(clusters): @@ -1015,7 +1295,9 @@ def run_remote_write_scenario(args: argparse.Namespace, case: dict[str, Any], ca def main() -> int: args = parse_args() case_path = args.case.resolve() - case = load_case(case_path) + if args.fixture_generator is not None: + require_binary(args.fixture_generator, "fixture generator", dry_run=False) + case = load_normalized_case(case_path, args.fixture_generator) scenario_config = scenario(case) scenario_kind = scenario_config.get("kind") tables = case_tables(case) @@ -1025,9 +1307,6 @@ def main() -> int: require_binary(args.candidate_bin, "candidate", dry_run=args.dry_run or args.fixture_only) if scenario_kind == "direct_readable_sst" and not args.dry_run and args.fixture_generator is None: raise ValueError("--fixture-generator is required unless --dry-run is set") - if args.fixture_generator is not None: - require_binary(args.fixture_generator, "fixture generator", dry_run=args.dry_run) - ports = allocate_ports(16) targets = [make_target("base", args.base_bin.resolve(), work_root, ports[:8]), make_target("candidate", args.candidate_bin.resolve(), work_root, ports[8:])] require_fresh_work_dirs(targets, reuse_work_dir=args.reuse_work_dir, dry_run=args.dry_run, fixture_only=args.fixture_only) diff --git a/tests/perf/test_query_regression_runner_compaction_toctou.py b/tests/perf/test_query_regression_runner_compaction_toctou.py new file mode 100644 index 00000000000..675f451a10e --- /dev/null +++ b/tests/perf/test_query_regression_runner_compaction_toctou.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Regression coverage for quiescing remote-write storage before read benches.""" + +import argparse +import importlib.util +import sys +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + + +RUNNER_PATH = Path(__file__).with_name("query_regression_runner.py") +SPEC = importlib.util.spec_from_file_location("query_regression_runner_under_test", RUNNER_PATH) +assert SPEC is not None and SPEC.loader is not None +runner = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = runner +SPEC.loader.exec_module(runner) + + +class RemoteWriteCompactionToctouTest(unittest.TestCase): + def test_quiesces_each_target_before_storage_inspection_and_read_bench(self) -> None: + lifecycle: dict[str, list[str]] = {"base": [], "candidate": []} + clusters = [] + + class RecordingCluster: + def __init__(self, target): + self.target = target + self.active = set() + self.stop_all_calls = 0 + clusters.append(self) + + def start_all(self, _config_path): + self.active.update(("metasrv", "datanode", "frontend")) + + def stop_component(self, name): + if name in self.active: + lifecycle[self.target.name].append(f'cluster.stop_component("{name}")') + self.active.remove(name) + + def stop_all(self): + self.stop_all_calls += 1 + for name in ("frontend", "datanode", "metasrv"): + self.stop_component(name) + + def component_report(self): + return {} + + def run_queries(target, _case, _tables, _http_timeout): + lifecycle[target.name].append("query measurement completes") + return {"validation": [], "validation_errors": [], "measurements": [], "status": "ok"} + + def inspect_storage(_helper, target, _storage, *, dry_run): + self.assertFalse(dry_run) + lifecycle[target.name].append("storage inspection") + return {"status": "ok", "summary": {"summary": {}}} + + def read_bench(_binary, target, _read_bench, _inspection, *, dry_run): + self.assertFalse(dry_run) + lifecycle[target.name].append("read bench") + return {"status": "ok"} + + remote_write = { + "database": "public", + "metric": "cpu_usage", + "physical_table": "cpu_usage_physical", + "series_count": 1, + "samples_per_series": 1, + "visibility_timeout_seconds": 1, + "storage": { + "inspect": True, + "min_files": None, + "min_files_with_column": None, + "max_total_file_size_bytes": None, + "max_column_compressed_size_bytes": None, + "max_column_uncompressed_size_bytes": None, + "require_encodings": [], + "forbid_encodings": [], + }, + "read_bench": {"enabled": True}, + "prom_store": { + "pending_rows_flush_interval": "1s", + "max_batch_rows": 1, + "max_concurrent_flushes": 1, + "worker_channel_capacity": 1, + "max_inflight_requests": 1, + }, + } + case = { + "scenario": { + "kind": "prom_remote_write_then_query", + "remote_write": remote_write, + "queries": [], + } + } + args = argparse.Namespace( + fixture_only=False, + fixture_generator=Path("/bin/true"), + remote_write_generator=None, + storage_inspector=None, + candidate_bin=Path("/bin/true"), + dry_run=False, + http_timeout=1.0, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + root = Path(tmpdir) + targets = [ + runner.make_target("base", Path("/bin/true"), root, list(range(10000, 10008))), + runner.make_target("candidate", Path("/bin/true"), root, list(range(10008, 10016))), + ] + report = {"targets": [], "thresholds": []} + with ( + patch.object(runner, "DistributedCluster", RecordingCluster), + patch.object(runner, "run_remote_write_ingestion", return_value=({"status": "ok"}, [{"ok": True}])), + patch.object(runner, "http_post_sql", return_value={"ok": True}), + patch.object(runner, "poll_expected_count", return_value={"ok": True, "row_count_ok": True}), + patch.object(runner, "run_queries", side_effect=run_queries), + patch.object(runner, "run_storage_inspection", side_effect=inspect_storage), + patch.object(runner, "run_read_bench", side_effect=read_bench), + ): + runner.run_remote_write_scenario(args, case, root / "case.toml", targets, report) + + for cluster in clusters: + self.assertEqual(cluster.stop_all_calls, 1) + self.assertEqual(cluster.active, set()) + cluster.stop_all() + self.assertEqual(cluster.stop_all_calls, 2) + self.assertEqual(cluster.active, set()) + + expected = [ + "query measurement completes", + 'cluster.stop_component("datanode")', + "storage inspection", + "read bench", + ] + relevant = {name: [event for event in events if event in expected] for name, events in lifecycle.items()} + for target_name in ("base", "candidate"): + with self.subTest(target=target_name): + self.assertEqual(relevant[target_name], expected) + + +if __name__ == "__main__": + unittest.main()