rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-26 13:46:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into prom-integer

Browse Source 
Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

# Conflicts:
#	config/config.md
#	config/datanode.example.toml
#	config/standalone.example.toml
#	src/cmd/tests/load_config_test.rs
#	src/metric-engine/src/config.rs
This commit is contained in:
Ruihang Xia
2026-07-17 11:50:34 +08:00
parent 5d5574d411 c969fb0a5e
commit 078a77c8df
279 changed files with 24742 additions and 4109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.agents/skills/greptimedb-development-docker-image/SKILL.md
+368
View File
@@ -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 <skill-dir>/scripts/collect_context.py \
--source <source-checkout> \
--bin greptime \
--profile nightly \
[--target <rust-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 <skill-dir>/scripts/collect_context.py \
--source <source-checkout> \
--bin <binary> \
--profile <profile> \
--engine <docker|podman> \
--registry <registry> \
--repository <repository> \
--tag <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 <existing-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 <skill-dir>/scripts/update_image_env.py \
--env <source-checkout>/.env \
--registry <registry> \
--repository <repository> \
--tag <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
<skill-dir>/scripts/build_binary.sh \
--source <source-checkout> --package cmd --bin greptime --profile nightly
```
For open-source arm64:
```bash
<skill-dir>/scripts/build_binary.sh \
--source <source-checkout> \
--package cmd \
--bin greptime \
--profile nightly \
--target aarch64-unknown-linux-gnu
```
For Enterprise, use the user-provided executable target, for example:
```bash
<skill-dir>/scripts/build_binary.sh \
--source <source-checkout> --bin greptime-ent-cloud --profile nightly
```
For a user-selected profile, replace `nightly` in the helper call and resolve
the binary under `target/<profile>/<binary>` (or
`target/<rust-target>/<profile>/<binary>` 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 <source-target-binary> <build-context>/greptime
```
Verify it before building the image with `file <build-context>/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
<skill-dir>/scripts/prepare_context.sh \
--context "$(mktemp -d)" \
--binary <source-target-binary> \
--dockerfile <skill-dir>/assets/Dockerfile \
--platform <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
<skill-dir>/scripts/build_image.sh \
--context <build-context> \
--image <image-reference> \
--platform <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 <image-reference> --format '{{.Os}}/{{.Architecture}}'
docker run --rm <image-reference> --help
```
For Podman, use `podman image inspect <image-reference>` and
`podman run --rm <image-reference> --help`. For a pushed Podman image, use
`podman manifest inspect <image-reference>`.
For a pushed image, inspect its remote manifest:
```bash
docker buildx imagetools inspect <image-reference>
```
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.
.agents/skills/greptimedb-development-docker-image/assets/Dockerfile
+15
View File
@@ -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"]
.agents/skills/greptimedb-development-docker-image/scripts/binary_platform.py
+78
View File
@@ -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("<HHI", header, 16)
header_size = struct.unpack_from("<H", header, 52)[0]
if elf_type not in (2, 3) or version != 1 or header_size != 64:
raise ValueError("binary has an unsupported ELF executable header")
return machine
def validate_binary(path: Path, platform: str) -> 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())
.agents/skills/greptimedb-development-docker-image/scripts/build_binary.sh
Executable
+65
View File
@@ -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[@]}"
)
.agents/skills/greptimedb-development-docker-image/scripts/build_image.sh
Executable
+136
View File
@@ -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
.agents/skills/greptimedb-development-docker-image/scripts/collect_context.py
Executable
+200
View File
@@ -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())
.agents/skills/greptimedb-development-docker-image/scripts/image_config.py
+155
View File
@@ -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
.agents/skills/greptimedb-development-docker-image/scripts/next_image_tag.py
Executable
+48
View File
@@ -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())
.agents/skills/greptimedb-development-docker-image/scripts/prepare_context.sh
Executable
+65
View File
@@ -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"
.agents/skills/greptimedb-development-docker-image/scripts/update_image_env.py
Executable
+53
View File
@@ -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())
.agents/skills/greptimedb-development-docker-image/tests/test_binary_platform.py
+100
View File
@@ -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, 16, elf_type)
struct.pack_into("<H" if endian == 1 else ">H", header, 18, machine)
struct.pack_into("<I", header, 20, 1)
struct.pack_into("<H", header, 52, 64)
return bytes(header)
class BinaryPlatformTest(unittest.TestCase):
def make_binary(self, content: bytes) -> 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()
.agents/skills/greptimedb-development-docker-image/tests/test_image_config.py
+100
View File
@@ -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()
.github/actions/setup-greptimedb-cluster/with-remote-wal-logical-prune.yaml
+1
View File
@@ -3,6 +3,7 @@ logging:
format: "json"
filters:
- log_store=debug
- rskafka=debug
meta:
configData: |-
[runtime]
.github/actions/setup-greptimedb-cluster/with-remote-wal.yaml
+1
View File
@@ -3,6 +3,7 @@ logging:
format: "json"
filters:
- log_store=debug
- rskafka=debug
meta:
configData: |-
[runtime]
.github/actions/setup-kafka-cluster/action.yml
+2 -2
View File
@@ -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 \
.github/runner-scale-sets/query-regression/Dockerfile
+68
View File
@@ -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"
.github/runner-scale-sets/query-regression/README.md
+258 -100
View File
@@ -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=<private-key.pem>
```
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.
.github/runner-scale-sets/query-regression/cache-pvc.yaml
+15
View File
@@ -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
.github/runner-scale-sets/query-regression/values-8-cores.yaml
+121 -8
View File
@@ -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
.github/runner-scale-sets/query-regression/values-paused.yaml
+2
View File
@@ -0,0 +1,2 @@
minRunners: 0
maxRunners: 0
.github/scripts/query-regression-comment.cjs
+3 -1
View File
@@ -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');
.github/scripts/query-regression-pr-metadata.py
+4 -1
View File
@@ -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"]),
.github/scripts/query-regression-run.py
+4 -10
View File
@@ -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:
.github/scripts/wait-kafka-ready.sh
Executable
+68
View File
@@ -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
.github/workflows/develop.yml
+10
View File
@@ -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:
.github/workflows/nightly-jsonbench.yaml
+82 -12
View File
@@ -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
.github/workflows/query-regression.yml
+457 -28
View File
@@ -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
AGENTS.md
+4 -2
View File
@@ -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
Cargo.lock
Generated
+10 -5
View File
@@ -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",
Cargo.toml
+3 -2
View File
@@ -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"
README.md
+7 -1
View File
@@ -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)
<a href="https://www.star-history.com/?repos=GreptimeTeam%2FGreptimeDB&type=date&legend=top-left">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://api.star-history.com/chart?repos=GreptimeTeam/GreptimeDB&type=date&theme=dark&legend=top-left&sealed_token=1kD9cyZVPnWvxbpbPPTWo_r7W351gDAhGx9_55zfD1vSIfHsdTRfH6NAtdPVX_M2ac8sjNxiuw7-0Vk4WvbyXmOuufxlAoU3QF7z4onyltr0o0c9xVfUnPKltz7llUc3nEDW7l_r2fHBBMtJGf1-jsitNozdAUpb8eb7OJp9N1LqEnMQK78ESp2zgjyr" />
<source media="(prefers-color-scheme: light)" srcset="https://api.star-history.com/chart?repos=GreptimeTeam/GreptimeDB&type=date&legend=top-left&sealed_token=1kD9cyZVPnWvxbpbPPTWo_r7W351gDAhGx9_55zfD1vSIfHsdTRfH6NAtdPVX_M2ac8sjNxiuw7-0Vk4WvbyXmOuufxlAoU3QF7z4onyltr0o0c9xVfUnPKltz7llUc3nEDW7l_r2fHBBMtJGf1-jsitNozdAUpb8eb7OJp9N1LqEnMQK78ESp2zgjyr" />
<img alt="Star History Chart" src="https://api.star-history.com/chart?repos=GreptimeTeam/GreptimeDB&type=date&legend=top-left&sealed_token=1kD9cyZVPnWvxbpbPPTWo_r7W351gDAhGx9_55zfD1vSIfHsdTRfH6NAtdPVX_M2ac8sjNxiuw7-0Vk4WvbyXmOuufxlAoU3QF7z4onyltr0o0c9xVfUnPKltz7llUc3nEDW7l_r2fHBBMtJGf1-jsitNozdAUpb8eb7OJp9N1LqEnMQK78ESp2zgjyr" />
</picture>
</a>
<img alt="Known Users" src="https://greptime.com/logo/img/users.png"/>
config/config.md
+10 -7
View File
@@ -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.<br/>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.<br/>Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"<br/>Password verifier formats: "plain:<password>", "pbkdf2_sha256:<iterations>:<hex_salt>:<hex_hash>",<br/>"mysql_native_password:<hex_sha1_sha1_password>"<br/>"pbkdf2_sha256" protects passwords at rest but is not compatible with mysql_native_password. |
| `user_provider` | String | Unset | The user provider for authentication.<br/>Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"<br/>Password verifier formats: "plain:<password>", "pbkdf2_sha256:<iterations>:<hex_salt>:<hex_hash>",<br/>"mysql_native_password:<hex_sha1_sha1_password>",<br/>"pg_scram_sha256:<iterations>:<hex_salt>:<hex_stored_key>:<hex_server_key>"<br/>"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's<br/>native password handshake; a MySQL client must send the password in cleartext for such users.<br/>"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.<br/>PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or<br/>"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.<br/>For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are<br/>observable in the SCRAM server-first message, and non-default values weaken resistance to username<br/>enumeration. |
| `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).<br/>Set to 0 to disable the limit. Default: "0" (unlimited) |
| `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.<br/>Options: "wait" (default, 10s timeout), "wait(<duration>)" (e.g., "wait(30s)"), "fail" |
| `init_regions_in_background` | Bool | `false` | Initialize all regions in the background during the startup.<br/>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.<br/>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.<br/>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.<br/>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.<br/>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.<br/>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.<br/>- `auto`: automatically determine the threshold based on the system memory size (default)<br/>- `unlimited`: no memory limit<br/>- `[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.<br/>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.<br/>Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"<br/>Password verifier formats: "plain:<password>", "pbkdf2_sha256:<iterations>:<hex_salt>:<hex_hash>",<br/>"mysql_native_password:<hex_sha1_sha1_password>"<br/>"pbkdf2_sha256" protects passwords at rest but is not compatible with mysql_native_password. |
| `user_provider` | String | Unset | The user provider for authentication.<br/>Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"<br/>Password verifier formats: "plain:<password>", "pbkdf2_sha256:<iterations>:<hex_salt>:<hex_hash>",<br/>"mysql_native_password:<hex_sha1_sha1_password>",<br/>"pg_scram_sha256:<iterations>:<hex_salt>:<hex_stored_key>:<hex_server_key>"<br/>"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's<br/>native password handshake; a MySQL client must send the password in cleartext for such users.<br/>"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.<br/>PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or<br/>"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.<br/>For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are<br/>observable in the SCRAM server-first message, and non-default values weaken resistance to username<br/>enumeration. |
| `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).<br/>Set to 0 to disable the limit. Default: "0" (unlimited) |
| `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.<br/>Options: "wait" (default, 10s timeout), "wait(<duration>)" (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`<br/>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.<br/>**It's only used when the provider is `kafka`**. |
| `wal.max_batch_bytes` | String | `1MB` | The max size of a single producer batch.<br/>Warning: Kafka has a default limit of 1MB per message in a topic.<br/>**It's only used when the provider is `kafka`**. |
| `wal.consumer_wait_timeout` | String | `100ms` | The consumer wait timeout.<br/>**It's only used when the provider is `kafka`**. |
| `wal.create_index` | Bool | `true` | Whether to enable WAL index creation.<br/>**It's only used when the provider is `kafka`**. |
| `wal.create_index` | Bool | `false` | Whether to enable WAL index creation.<br/>**It's only used when the provider is `kafka`**. |
| `wal.dump_index_interval` | String | `60s` | The interval for dumping WAL indexes.<br/>**It's only used when the provider is `kafka`**. |
| `wal.overwrite_entry_start_id` | Bool | `false` | Ignore missing entries during read WAL.<br/>**It's only used when the provider is `kafka`**.<br/><br/>This option ensures that when Kafka messages are deleted, the system<br/>can still successfully replay memtable data without throwing an<br/>out-of-range error.<br/>However, enabling this option might lead to unexpected data loss,<br/>as the system will skip over missing entries instead of treating<br/>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.<br/>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.<br/>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.<br/>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.<br/>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.<br/>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).<br/>Only applies during full file listing GC.<br/>This uses the object's last-modified timestamp as a heuristic (strict less-than comparison);<br/>do not configure this value too small in production to avoid deleting pre-manifest files<br/>from in-progress compaction or flush.<br/>For active/open regions, an unknown file is deleted only if its object last-modified time exceeds this TTL.<br/>If the object store does not provide a last-modified timestamp, the file is conservatively kept.<br/>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. |
config/datanode.example.toml
+6 -5
View File
@@ -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
config/frontend.example.toml
+10 -2
View File
@@ -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:<password>", "pbkdf2_sha256:<iterations>:<hex_salt>:<hex_hash>",
## "mysql_native_password:<hex_sha1_sha1_password>"
## "pbkdf2_sha256" protects passwords at rest but is not compatible with mysql_native_password.
## "mysql_native_password:<hex_sha1_sha1_password>",
## "pg_scram_sha256:<iterations>:<hex_salt>:<hex_stored_key>:<hex_server_key>"
## "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"
config/metasrv.example.toml
+6
View File
@@ -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.
config/standalone.example.toml
+15 -6
View File
@@ -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:<password>", "pbkdf2_sha256:<iterations>:<hex_salt>:<hex_hash>",
## "mysql_native_password:<hex_sha1_sha1_password>"
## "pbkdf2_sha256" protects passwords at rest but is not compatible with mysql_native_password.
## "mysql_native_password:<hex_sha1_sha1_password>",
## "pg_scram_sha256:<iterations>:<hex_salt>:<hex_stored_key>:<hex_server_key>"
## "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
src/api/src/helper.rs
+85 -25
View File
@@ -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::<Vec<_>>()
.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::<BTreeMap<_, _>>()
.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::<BTreeMap<_, _>>();
let fields = fields.into_iter().collect::<BTreeMap<_, _>>();
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);
src/auth/Cargo.toml
+3
View File
@@ -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
src/auth/src/common.rs
+275 -1
View File
@@ -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<Sha256>;
/// 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<u8>,
stored_key: Vec<u8>,
server_key: Vec<u8>,
}
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", &"<REDACTED>")
.field("stored_key", &"<REDACTED>")
.field("server_key", &"<REDACTED>")
.finish()
}
}
impl PgScramSha256Verifier {
pub fn new(
iterations: u32,
salt: Vec<u8>,
stored_key: Vec<u8>,
server_key: Vec<u8>,
) -> Result<Self> {
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<Self> {
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<u8>,
salt: Vec<u8>,
iterations: u32,
) -> Result<Self> {
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<bool> {
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<Option<Vec<u8>>> {
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<String> {
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<Vec<u8>> {
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::<Sha256>(password, salt, iterations, &mut salted_password);
Ok(salted_password.to_vec())
}
fn hmac_sha256(key: &[u8], msg: &[u8]) -> Vec<u8> {
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<u8> {
let mut hasher = Sha256::new();
hasher.update(data);
hasher.finalize().to_vec()
}
fn xor(lhs: &[u8], rhs: &[u8]) -> Vec<u8> {
lhs.iter().zip(rhs).map(|(l, r)| l ^ r).collect()
}
fn sha1_two(input_1: &[u8], input_2: &[u8]) -> Vec<u8> {
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()
);
}
}
src/auth/src/lib.rs
+6 -5
View File
@@ -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<dyn UserInfo>;
src/auth/src/permission.rs
+42 -7
View File
@@ -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<PermissionResp>;
fn check_permission_with_context(
&self,
user_info: UserInfoRef,
req: PermissionReq,
_current_schema: Option<&str>,
) -> Result<PermissionResp> {
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<PermissionResp> {
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<PermissionResp> {
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());
}
}
src/auth/src/user_provider.rs
+336 -15
View File
@@ -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<PgAuthInfo> {
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<UserInfoRef>,
},
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<u8>,
@@ -81,12 +102,46 @@ pub(crate) enum PasswordVerifier {
MysqlNativePassword {
hash_stage_2: Vec<u8>,
},
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(&"<REDACTED>").finish()
}
PasswordVerifier::Pbkdf2Sha256 { iterations, .. } => f
@@ -99,14 +154,33 @@ impl fmt::Debug for PasswordVerifier {
.debug_struct("MysqlNativePassword")
.field("hash_stage_2", &"<REDACTED>")
.finish(),
PasswordVerifier::PgScramSha256(_) => f
.debug_struct("PgScramSha256")
.field("verifier", &"<REDACTED>")
.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<Self> {
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<Self> {
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::<u32>().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<PgScramSha256Verifier> {
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<UserInfoMap> {
/// - `plain:<password>`
/// - `pbkdf2_sha256:<iterations>:<hex-encoded-salt>:<hex-encoded-hash>`
/// - `mysql_native_password:<hex-encoded-sha1-sha1-password>`
/// - `pg_scram_sha256:<iterations>:<hex-encoded-salt>:<hex-encoded-stored-key>:<hex-encoded-server-key>`
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<PgAuthInfo> {
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<u8> {
@@ -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::<Sha256>(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("<REDACTED>"));
assert!(!debug.contains("secret"));
src/auth/src/user_provider/static_user_provider.rs
+6 -1
View File
@@ -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<PgAuthInfo> {
postgres_auth_info_with_credential(&self.users, id)
}
async fn authorize(
&self,
_catalog: &str,
src/auth/src/user_provider/watch_file_user_provider.rs
+9 -1
View File
@@ -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<PgAuthInfo> {
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(())
src/catalog/src/information_extension.rs
+191 -36
View File
@@ -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<SendableRecordBatchStream, crate::error::Error> {
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<PlanProperties>,
}
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<PlanProperties> {
&self.properties
}
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
vec![]
}
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,
) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
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<TaskContext>,
) -> datafusion::error::Result<DfSendableRecordBatchStream> {
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<SendableRecordBatchStream, Self::Error> {
// 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<Option<Arc<dyn ExecutionPlan>>, Self::Error> {
Ok(Some(Arc::new(DistributedInspectExec::new(
request,
schema,
self.meta_client.clone(),
self.datanode_manager.clone(),
))))
}
}
src/catalog/src/system_schema.rs
+16
View File
@@ -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<SendableRecordBatchStream>;
fn scan_plan(&self, _request: ScanRequest) -> Result<Option<Arc<dyn ExecutionPlan>>> {
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<Option<Arc<dyn ExecutionPlan>>, BoxedError> {
self.table
.scan_plan(request)
.map_err(BoxedError::new)
.context(TablesRecordBatchSnafu)
.map_err(BoxedError::new)
}
}
src/catalog/src/system_schema/information_schema.rs
+19
View File
@@ -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<SendableRecordBatchStream>;
fn scan_plan(&self, _request: ScanRequest) -> Result<Option<Arc<dyn ExecutionPlan>>> {
Ok(None)
}
fn table_type(&self) -> TableType {
TableType::Temporary
}
@@ -449,6 +454,10 @@ where
fn to_stream(&self, request: ScanRequest) -> Result<SendableRecordBatchStream> {
InformationTable::to_stream(self, request)
}
fn scan_plan(&self, request: ScanRequest) -> Result<Option<Arc<dyn ExecutionPlan>>> {
InformationTable::scan_plan(self, request)
}
}
pub type InformationExtensionRef = Arc<dyn InformationExtension<Error = Error> + Send + Sync>;
@@ -475,6 +484,16 @@ pub trait InformationExtension {
&self,
request: DatanodeInspectRequest,
) -> std::result::Result<SendableRecordBatchStream, Self::Error>;
/// 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<Option<Arc<dyn ExecutionPlan>>, Self::Error> {
Ok(None)
}
}
/// The request to inspect the datanode.
src/catalog/src/system_schema/information_schema/region_info.rs
+17
View File
@@ -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<Option<Arc<dyn ExecutionPlan>>> {
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)
}
}
src/catalog/src/system_schema/information_schema/ssts.rs
+158 -15
View File
@@ -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<SchemaRef> {
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<dyn CatalogManager>,
kind: DatanodeInspectKind,
schema: &SchemaRef,
request: ScanRequest,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
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<SendableRecordBatchStream> {
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<Option<Arc<dyn ExecutionPlan>>> {
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<SendableRecordBatchStream> {
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<Option<Arc<dyn ExecutionPlan>>> {
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<SendableRecordBatchStream> {
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<Option<Arc<dyn ExecutionPlan>>> {
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<Vec<NodeInfo>, Self::Error> {
Ok(vec![])
}
async fn procedures(
&self,
) -> std::result::Result<Vec<(String, ProcedureInfo)>, Self::Error> {
Ok(vec![])
}
async fn region_stats(&self) -> std::result::Result<Vec<RegionStat>, Self::Error> {
Ok(vec![])
}
async fn flow_stats(&self) -> std::result::Result<Option<FlowStat>, Self::Error> {
Ok(None)
}
async fn inspect_datanode(
&self,
_request: DatanodeInspectRequest,
) -> std::result::Result<SendableRecordBatchStream, Self::Error> {
Ok(RecordBatches::empty().as_stream())
}
fn inspect_datanode_plan(
&self,
_request: DatanodeInspectRequest,
schema: SchemaRef,
) -> std::result::Result<Option<Arc<dyn ExecutionPlan>>, 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::<EmptyExec>());
assert_eq!(1, plan.schema().fields().len());
}
}
src/cli/src/bench/metadata.rs
+1
View File
@@ -119,6 +119,7 @@ impl TableMetadataBencher {
&table_info.table_name(),
&table_route,
&region_wal_options,
None,
)
.await;
start.elapsed()
src/cli/src/metadata/control/del/table.rs
+7 -1
View File
@@ -118,7 +118,13 @@ impl TableMetadataDeleter {
.map_err(BoxedError::new)?;
self.table_metadata_manager
.delete_table_metadata(table_id, &table_name, &table_route, &region_wal_options)
.delete_table_metadata(
table_id,
&table_name,
&table_route,
&region_wal_options,
None,
)
.await
.map_err(BoxedError::new)?;
Ok(())
src/cmd/Cargo.toml
+2 -5
View File
@@ -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"]
src/cmd/src/bin/greptime.rs
-1
View File
@@ -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) => {
src/cmd/src/bin/prom_remote_write_fixture.rs
-144
View File
@@ -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<dyn std::error::Error>> {
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::<Vec<_>>();
rows += timeseries
.iter()
.map(|ts| ts.samples.len() as u64)
.sum::<u64>();
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)
}
src/cmd/src/bin/query_perf_fixture.rs
+3 -649
View File
@@ -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<String>,
/// Table name to generate when the case contains multiple [[tables]].
#[arg(long, value_name = "NAME")]
table: Option<String>,
/// 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<u64>,
tables: Vec<TableConfig>,
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<bool>,
#[serde(default)]
sst_format: Option<String>,
primary_key: Vec<String>,
time_index: String,
columns: Vec<ColumnConfig>,
}
#[derive(Debug, Deserialize)]
struct ColumnConfig {
name: String,
#[serde(rename = "type")]
ty: String,
semantic: String,
distribution: Option<Distribution>,
}
#[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<String, String>) -> Vec<u8> {
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::<UInt32Type>::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::<Vec<_>>(),
)) as ArrayRef);
}
("field", "UINT64") => columns.push(Arc::new(UInt64Array::from(
(0..rows).map(|r| (base_row + r) as u64).collect::<Vec<_>>(),
)) as ArrayRef),
("timestamp", "TIMESTAMP(9)") => columns.push(Arc::new(TimestampNanosecondArray::from(
(0..rows)
.map(|r| timestamp_for_row(layout, base_row, r))
.collect::<Vec<_>>(),
)) 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::<Vec<_>>(),
)) as ArrayRef)
}
other => panic!("unsupported column combination {other:?}"),
}
}
let mut pk_builder = BinaryDictionaryBuilder::<UInt32Type>::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 <name> 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;
}
src/cmd/src/bin/query_perf_fixture/case.rs
+470
View File
@@ -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<u64>,
pub(super) tables: Vec<TableConfig>,
pub(super) layout: LayoutConfig,
}
#[derive(Debug, Deserialize, Serialize)]
pub(super) struct PromRemoteWriteThenQueryScenario {
#[serde(default)]
pub(super) queries: Vec<serde_json::Value>,
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<u64>,
#[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<StorageConfig>,
#[serde(default)]
pub(super) read_bench: Option<ReadBenchConfig>,
}
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<String>,
#[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<String>,
#[serde(default)]
pub(super) forbid_encodings: Vec<String>,
#[serde(default)]
pub(super) max_total_file_size_bytes: Option<u64>,
#[serde(default)]
pub(super) max_column_compressed_size_bytes: Option<u64>,
#[serde(default)]
pub(super) max_column_uncompressed_size_bytes: Option<u64>,
#[serde(default)]
pub(super) max_candidate_total_file_size_regression_pct: Option<f64>,
#[serde(default)]
pub(super) max_candidate_column_compressed_size_regression_pct: Option<f64>,
#[serde(default)]
pub(super) max_candidate_column_uncompressed_size_regression_pct: Option<f64>,
#[serde(skip_deserializing, default)]
pub(super) planned_thresholds: Vec<StorageThresholdPlan>,
}
#[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<T: Serialize>(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<String>,
#[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<usize>,
#[serde(flatten)]
pub(super) thresholds: HashMap<String, serde_json::Value>,
}
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<bool>,
#[serde(default)]
pub(super) sst_format: Option<String>,
pub(super) primary_key: Vec<String>,
pub(super) time_index: String,
pub(super) columns: Vec<ColumnConfig>,
}
#[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<Distribution>,
}
#[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,
}
src/cmd/src/bin/query_perf_fixture/direct_sst.rs
+529
View File
@@ -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<String, String>) -> Vec<u8> {
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::<UInt32Type>::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::<Vec<_>>(),
)) as ArrayRef);
}
("field", "UINT64") => columns.push(Arc::new(UInt64Array::from(
(0..rows).map(|r| (base_row + r) as u64).collect::<Vec<_>>(),
)) as ArrayRef),
("timestamp", "TIMESTAMP(9)") => columns.push(Arc::new(TimestampNanosecondArray::from(
(0..rows)
.map(|r| timestamp_for_row(layout, base_row, r))
.collect::<Vec<_>>(),
)) 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::<Vec<_>>(),
)) as ArrayRef)
}
other => panic!("unsupported column combination {other:?}"),
}
}
let mut pk_builder = BinaryDictionaryBuilder::<UInt32Type>::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 <name> 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());
}
src/cmd/src/bin/query_perf_fixture/inspect_footer.rs
+181
View File
@@ -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<FooterFileReport>,
}
#[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<String>,
}
#[derive(Debug, Serialize)]
struct FooterFileReport {
path: String,
relative_path: String,
file_size: u64,
num_rows: i64,
num_row_groups: usize,
columns: Vec<FooterColumnChunkReport>,
}
#[derive(Debug, Serialize)]
struct FooterColumnChunkReport {
row_group_index: usize,
column_index: usize,
column_path: String,
encodings: Vec<String>,
compression: String,
compressed_size: i64,
uncompressed_size: i64,
num_values: i64,
}
pub(super) fn run_inspect_footer(
args: InspectFooterArgs,
) -> Result<(), Box<dyn std::error::Error>> {
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::<Result<Vec<_>, _>>()?;
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<Vec<PathBuf>, Box<dyn std::error::Error>> {
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<PathBuf>,
) -> Result<(), Box<dyn std::error::Error>> {
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<FooterFileReport, Box<dyn std::error::Error>> {
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
}
src/cmd/src/bin/query_perf_fixture/mod.rs
+133
View File
@@ -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>,
#[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<PathBuf>,
/// Output directory (creates object-store/, manifest/, files.jsonl, summary.json).
#[arg(long, value_name = "DIR")]
out_dir: Option<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<String>,
/// Table name to generate when the case contains multiple [[tables]].
#[arg(long, value_name = "NAME")]
table: Option<String>,
/// 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<dyn std::error::Error>> {
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(())
}
src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs
+333
View File
@@ -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<u64>,
}
pub(super) async fn run_prom_remote_write(
args: PromRemoteWriteArgs,
) -> Result<(), Box<dyn std::error::Error>> {
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::<Vec<_>>();
rows += timeseries
.iter()
.map(|ts| ts.samples.len() as u64)
.sum::<u64>();
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::<Vec<_>>();
let repeated = (0..args.samples_per_series)
.map(|sample_idx| deterministic_prom_value(&args, 0, sample_idx))
.collect::<Vec<_>>();
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::<HashSet<_>>()
.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::<Vec<_>>();
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);
}
}
src/cmd/src/bin/query_perf_fixture/util.rs
+23
View File
@@ -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()
}
src/cmd/src/datanode.rs
+2 -7
View File
@@ -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);
src/cmd/src/standalone.rs
+1
View File
@@ -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(
src/cmd/src/user.rs
+40 -5
View File
@@ -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<String>,
}
@@ -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());
src/cmd/tests/load_config_test.rs
+53 -2
View File
@@ -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::<MetasrvOptions>::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::<MetasrvOptions> {
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::<DatanodeOptions>::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)
);
}
src/common/datasource/src/file_format.rs
+1
View File
@@ -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";
src/common/datasource/src/file_format/csv.rs
+38 -1
View File
@@ -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<usize>,
pub compression_type: CompressionType,
@@ -111,6 +112,16 @@ impl TryFrom<&HashMap<String, String>> 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
src/common/meta/Cargo.toml
+1 -1
View File
@@ -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
src/common/meta/src/ddl.rs
+3
View File
@@ -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<Duration>,
}
impl DdlContext {
src/common/meta/src/ddl/drop_table.rs
+68 -6
View File
@@ -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<Self> {
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<Status> {
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<RegionNumber, WalOptions>,
#[serde(default)]
pub allow_rollback: bool,
#[serde(default)]
pub soft_drop_enabled: bool,
#[serde(default)]
pub dropped_at: Option<i64>,
#[serde(default)]
pub retention_expires_at: Option<i64>,
#[serde(default)]
pub soft_drop_retention_millis: Option<i64>,
#[serde(default)]
pub drop_generation: Option<String>,
}
impl DropTableData {
pub fn new(task: DropTableTask) -> Self {
pub fn new(
task: DropTableTask,
soft_drop_enabled: bool,
soft_drop_retention: Option<std::time::Duration>,
) -> 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,
}
}
src/common/meta/src/ddl/drop_table/executor.rs
+138 -34
View File
@@ -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<Control<()>> {
pub async fn on_prepare(
&self,
ctx: &DdlContext,
soft_drop_enabled: bool,
) -> Result<Control<()>> {
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<RegionNumber, WalOptions>,
dropped_at: Option<i64>,
retention_expires_at: Option<i64>,
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<RegionNumber, WalOptions>,
) -> 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::<Vec<_>>();
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::<Result<Vec<_>>>()?;
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::<Vec<_>>();
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::<Result<Vec<_>>>()?;
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());
}
}
src/common/meta/src/ddl/drop_table/metadata.rs
+8 -1
View File
@@ -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;
src/common/meta/src/ddl/purge_dropped_table.rs
+118 -28
View File
@@ -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<Self> {
let data: PurgeDroppedTableData = serde_json::from_str(json).context(FromJsonSnafu)?;
Ok(Self { context, data })
}
async fn on_prepare(&mut self) -> Result<Status> {
async fn eligible_dropped_table(&self) -> Result<Option<DroppedTableMetadata>> {
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<Status> {
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<Status> {
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<Status> {
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<Status> {
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<Status> {
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<TableRouteValue>,
#[serde(default)]
region_wal_options: HashMap<RegionNumber, WalOptions>,
#[serde(default)]
check_expired: bool,
#[serde(default)]
drop_generation: Option<String>,
#[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,
}
src/common/meta/src/ddl/tests/drop_table.rs
+1248 -133
View File
File diff suppressed because it is too large Load Diff
src/common/meta/src/ddl/undrop_table.rs
+391 -51
View File
@@ -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<TableName>,
) -> Self {
let mut data = UndropTableData::new(task);
data.table_name = table_name;
Self { context, data }
}
pub fn from_json(json: &str, context: DdlContext) -> ProcedureResult<Self> {
@@ -70,6 +77,7 @@ impl UndropTableProcedure {
}
pub(crate) async fn on_prepare(&mut self) -> Result<Status> {
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<Status> {
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<bool> {
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,
&region_routes,
false,
)
.await
} else {
Ok(())
};
self.context
.deregister_failure_detectors(convert_region_routes_to_detecting_regions(
&region_routes,
))
.await;
close_result?;
Ok(())
}
async fn on_open_regions(&mut self) -> Result<Status> {
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<RegionNumber, WalOptions>,
) -> 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<RegionNumber, WalOptions>,
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<TableRouteValue>,
#[serde(default)]
region_wal_options: HashMap<RegionNumber, WalOptions>,
#[serde(default)]
dropped_at: Option<i64>,
#[serde(default)]
retention_expires_at: Option<i64>,
#[serde(default)]
drop_generation: Option<String>,
#[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<Vec<DetectingRegion>>,
}
#[async_trait]
impl RegionFailureDetectorController for RecordingRegionFailureDetectorController {
async fn register_failure_detectors(&self, _detecting_regions: Vec<DetectingRegion>) {}
async fn reset_failure_detectors(&self, _detecting_regions: Vec<DetectingRegion>) {}
async fn deregister_failure_detectors(&self, detecting_regions: Vec<DetectingRegion>) {
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)]
);
}
}
src/common/meta/src/ddl_manager.rs
+95 -2
View File
@@ -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<Output>)> {
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<Output>)> {
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 { .. }));
}
}
src/common/meta/src/key.rs
+336 -6
View File
@@ -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<i64>,
/// Fixed automatic-GC deadline in Unix milliseconds.
pub retention_expires_at: Option<i64>,
/// Unique identity of this soft-drop generation.
pub drop_generation: Option<String>,
}
#[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<RegionNumber, WalOptions>,
/// Unix timestamp in milliseconds when this table was soft-dropped.
pub dropped_at: Option<i64>,
/// Fixed automatic-GC deadline in Unix milliseconds.
pub retention_expires_at: Option<i64>,
/// Unique identity of this soft-drop generation.
pub drop_generation: Option<String>,
}
/// Lifecycle markers stored with a logically deleted table.
pub struct DroppedTableLifecycle<'a> {
pub dropped_at: Option<i64>,
pub retention_expires_at: Option<i64>,
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<u8> {
format!("{DROPPED_AT_KEY_PREFIX}/{table_id}").into_bytes()
}
fn retention_expires_at_key(table_id: TableId) -> Vec<u8> {
format!("{RETENTION_EXPIRES_AT_KEY_PREFIX}/{table_id}").into_bytes()
}
fn drop_generation_key(table_id: TableId) -> Vec<u8> {
format!("{DROP_GENERATION_KEY_PREFIX}/{table_id}").into_bytes()
}
fn purging_key(table_id: TableId) -> Vec<u8> {
format!("{PURGING_KEY_PREFIX}/{table_id}").into_bytes()
}
impl<T: DeserializeOwned + Serialize> Deref for DeserializedValueWithBytes<T> {
@@ -998,10 +1039,78 @@ impl TableMetadataManager {
table_name: &TableName,
table_route_value: &TableRouteValue,
region_wal_options: &HashMap<RegionNumber, WalOptions>,
dropped_at: Option<i64>,
) -> 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<RegionNumber, WalOptions>,
dropped_at: Option<i64>,
retention_expires_at: Option<i64>,
) -> 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<RegionNumber, WalOptions>,
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::<Vec<_>>();
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<bool> {
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<Option<String>> {
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<Option<i64>> {
Self::parse_timestamp_marker(table_id, "dropped timestamp", kv)
}
fn parse_retention_expires_at(table_id: TableId, kv: Option<&KeyValue>) -> Result<Option<i64>> {
Self::parse_timestamp_marker(table_id, "retention deadline", kv)
}
fn parse_drop_generation(kv: Option<&KeyValue>) -> Option<String> {
kv.map(|kv| String::from_utf8_lossy(&kv.value).into_owned())
}
fn parse_timestamp_marker(
table_id: TableId,
description: &str,
kv: Option<&KeyValue>,
) -> Result<Option<i64>> {
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,
&region_wal_options,
None,
)
.await
.unwrap();
@@ -2188,6 +2430,7 @@ mod tests {
&table_name,
table_route_value,
&region_wal_options,
None,
)
.await
.unwrap();
@@ -2198,6 +2441,7 @@ mod tests {
&table_name,
table_route_value,
&region_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<_>>(),
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;
src/common/meta/src/key/tombstone.rs
+352 -8
View File
@@ -190,6 +190,7 @@ impl TombstoneManager {
keys: &[Vec<u8>],
dest_keys: &[Vec<u8>],
require_dest_not_exists: bool,
extra_ops: Vec<TxnOp>,
) -> Result<usize> {
ensure!(
keys.len() == dest_keys.len(),
@@ -213,6 +214,9 @@ impl TombstoneManager {
.into_iter()
.map(|kv| (kv.key, kv.value))
.collect::<HashMap<_, _>>();
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<TxnOp>) -> 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<Vec<u8>>,
dest_keys: Vec<Vec<u8>>,
require_dest_not_exists: bool,
mut extra_ops: Vec<TxnOp>,
) -> Result<usize> {
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::<Vec<_>>();
let dest_keys_chunks = dest_keys.chunks(chunk_size).collect::<Vec<_>>();
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<Vec<u8>>,
dest_keys: Vec<Vec<u8>>,
require_dest_not_exists: bool,
) -> Result<usize> {
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<Vec<u8>>,
marker_key: Vec<u8>,
marker_value: Vec<u8>,
) -> Result<usize> {
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<Vec<u8>>,
markers: Vec<(Vec<u8>, Vec<u8>)>,
) -> Result<usize> {
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<Vec<u8>>,
marker_key: Vec<u8>,
) -> Result<usize> {
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<Vec<u8>>,
marker_keys: Vec<Vec<u8>>,
) -> Result<usize> {
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<Vec<u8>>,
marker_key: Vec<u8>,
) -> Result<usize> {
keys.push(marker_key);
self.delete(keys).await
}
/// Deletes tombstones and associated markers.
pub async fn delete_with_markers(
&self,
mut keys: Vec<Vec<u8>>,
marker_keys: Vec<Vec<u8>>,
) -> Result<usize> {
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<MemoryKvBackend<Error>>,
max_txn_ops: usize,
txn_op_counts: Mutex<Vec<usize>>,
}
#[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());
src/common/meta/src/test_util.rs
+1
View File
@@ -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)),
}
}
src/common/query/src/lib.rs
+1
View File
@@ -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"))]
src/common/query/src/prometheus.rs
+37
View File
@@ -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
)));
}
}
src/common/recordbatch/Cargo.toml
+4
View File
@@ -34,3 +34,7 @@ tokio.workspace = true
[[bench]]
name = "iter_record_batch_rows"
harness = false
[[bench]]
name = "memory_accounting"
harness = false
src/common/recordbatch/benches/memory_accounting.rs
+117
View File
@@ -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<VectorRef> = 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::<Vec<_>>();
batch(vec![
(
"i32",
ConcreteDataType::int32_datatype(),
false,
Arc::new(Int32Vector::from_slice(
(0..rows).map(|value| value as i32).collect::<Vec<_>>(),
)),
),
(
"u64",
ConcreteDataType::uint64_datatype(),
false,
Arc::new(UInt64Vector::from_slice(
(0..rows).map(|value| value as u64).collect::<Vec<_>>(),
)),
),
(
"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::<Vec<_>>();
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);
src/common/recordbatch/src/lib.rs
+31 -2
View File
@@ -769,7 +769,7 @@ impl MemoryTrackedStream {
batch: RecordBatch,
cx: &mut Context<'_>,
) -> Poll<Option<Result<RecordBatch>>> {
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()])
src/common/recordbatch/src/recordbatch.rs
+117 -3
View File
@@ -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::<Vec<_>>();
let binary = (0_u32..1024)
.map(|value| (value % 3 != 0).then(|| value.to_le_bytes().to_vec()))
.collect::<Vec<_>>();
let columns: Vec<VectorRef> = 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::<usize>()
);
}
#[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::<Vec<_>>();
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::<Vec<_>>();
let logical_total = slices
.iter()
.map(RecordBatch::logical_slice_memory_size)
.sum::<usize>();
let buffer_total = slices
.iter()
.map(RecordBatch::buffer_memory_size)
.sum::<usize>();
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<VectorRef> =
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![
src/common/telemetry/src/logging.rs
+24 -15
View File
@@ -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);
}
}
src/common/time/src/date.rs
+4
View File
@@ -198,6 +198,10 @@ impl Date {
pub fn negative(&self) -> Self {
Self(-self.0)
}
pub fn checked_negative(&self) -> Option<Self> {
self.0.checked_neg().map(Self)
}
}
#[cfg(test)]
src/common/time/src/duration.rs
+5
View File
@@ -97,6 +97,11 @@ impl Duration {
self.value = -self.value;
self
}
pub fn checked_negative(mut self) -> Option<Self> {
self.value = self.value.checked_neg()?;
Some(self)
}
}
/// Convert i64 to Duration Type.
src/common/time/src/interval.rs
+19
View File
@@ -80,6 +80,10 @@ impl IntervalYearMonth {
Self::new(-self.months)
}
pub fn checked_negative(&self) -> Option<Self> {
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<Self> {
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<Self> {
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()
}
src/common/time/src/time.rs
+5
View File
@@ -145,6 +145,11 @@ impl Time {
self.value = -self.value;
self
}
pub fn checked_negative(mut self) -> Option<Self> {
self.value = self.value.checked_neg()?;
Some(self)
}
}
impl From<i64> for Time {
src/common/time/src/timestamp.rs
+5
View File
@@ -484,6 +484,11 @@ impl Timestamp {
self.value = -self.value;
self
}
pub fn checked_negative(mut self) -> Option<Self> {
self.value = self.value.checked_neg()?;
Some(self)
}
}
impl Timestamp {
src/common/wal/src/config/kafka/datanode.rs
+1 -1
View File
@@ -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,
src/datanode/src/alive_keeper.rs
+1 -1
View File
@@ -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
{
src/datanode/src/error.rs
+13
View File
@@ -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,
src/datanode/src/heartbeat.rs
+17
View File
@@ -121,6 +121,7 @@ impl HeartbeatTask {
pub async fn create_streams(
meta_client: &MetaClient,
local_gc_enabled: bool,
running: Arc<AtomicBool>,
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(),
src/datanode/src/heartbeat/handler.rs
+4 -1
View File
@@ -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();
src/datanode/src/heartbeat/handler/close_region.rs
+4 -2
View File
@@ -40,8 +40,10 @@ impl InstructionHandler for CloseRegionsHandler {
.collect::<Vec<_>>();
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;
src/datanode/src/heartbeat/handler/open_region.rs
+8 -2
View File
@@ -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();
src/datanode/src/region_server.rs
+143 -6
View File
@@ -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(&region_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(&region_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(&region_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 {
src/datatypes/src/json/value.rs
+9 -11
View File
@@ -317,13 +317,6 @@ impl JsonValue {
}
}
pub(crate) fn new_with(json_variant: JsonVariant, json_type: Arc<JsonNativeType>) -> 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())
}

Some files were not shown because too many files have changed in this diff Show More