rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-14 06:46:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into feat/allow-widen-time-index-unit

Browse Source
This commit is contained in:
Ning Sun
2026-08-31 15:42:33 +08:00
committed by GitHub
parent 79dea6d56d 51b94bb73f
commit 734c0d5db8
141 changed files with 11276 additions and 2880 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/runner-scale-sets/query-regression/README.md
+158 -295
View File
@@ -1,321 +1,184 @@
# Query regression self-hosted runners
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.
The `Query Regression` workflow runs on **Aliyun ECS ephemeral runners**
(`aliyun-ecs`, the default and only automated path): a `provision` job on
`ubuntu-latest` creates one pay-as-you-go ECS instance per run, and the
instance registers itself as an ephemeral GitHub runner with a per-run
label. A `teardown` job (`if: always()`) deletes the instance;
`query-regression-janitor.yml` sweeps tagged leftovers older than 4 hours
daily.
## Prerequisites and trust admission
Build caches live on the instance's system disk, so **every run compiles
cold**; only the within-run reuse (base build warms the candidate build
through the shared target dir and sccache) applies. There is no retained
data disk.
Install the ARC scale set controller if it is not already installed:
Dispatching with any other `runner` value uses it as a literal self-hosted
runner label, which is how a manually prepared host (see
`ecs-image/bootstrap-runner-host.sh`) runs the workflow. For PR labels, set
the repository variable `QUERY_REGRESSION_PR_RUNNER` to such a label to
redirect PR runs away from ECS.
The office ARC scale set `perf-regression-8-cores` that previously ran this
workflow is retired; see git history for its values files and pause/deploy
procedures. Tearing down the office cluster (helm release, runner namespace,
and the `query-regression-build-cache` PVC) is a manual operator action
outside this repository. This directory keeps its historical
`runner-scale-sets` name for path stability.
## Nightly vs previous nightly
`query-regression-nightly.yml` runs after a successful `GreptimeDB Nightly
Build` (`workflow_run`). It resolves that run's `head_sha` as the candidate
and the previous successful nightly (same branch, typically Friday when
Monday's nightly fires) as the base, then calls `query-regression.yml` with
those immutable SHAs. Both binaries are still compiled on the ECS runner;
this is not an artifact-download path. `workflow_dispatch` can pass explicit
`base_ref` / `candidate_ref` or a nightly run id. If there is no previous
nightly, or both nightlies built the same commit, the comparison is skipped.
## Aliyun ECS path
Configuration lives in repository variables/secrets:
| Kind | Name | Purpose |
| --- | --- | --- |
| secret | `ALICLOUD_ECS_ACCESS_KEY_ID` / `ALICLOUD_ECS_ACCESS_KEY_SECRET` | RAM user scoped to ECS RunInstances/DeleteInstances/Describe*/CreateImage/RunCommand. Used only by provision/teardown jobs on `ubuntu-latest`; never reaches the ECS instance. |
| secret | `GH_PERSONAL_ACCESS_TOKEN` | Creates the short-lived runner registration token (shared with the jsonbench EC2 path). |
| vars | `ALIYUN_ECS_REGION_ID` / `ALIYUN_ECS_VSWITCH_ID` / `ALIYUN_ECS_SECURITY_GROUP_ID` | Network placement. The security group should allow egress only; no inbound rules are needed. The vSwitch pins the zone. |
| vars | `ALIYUN_ECS_INSTANCE_TYPE` | Dedicated (non-burstable, non-shared) instance family. Prefer 32 GiB (e.g. `ecs.g8i.2xlarge`); `ecs.c9i.2xlarge` is 8c16g and nightly thin-LTO of greptime peaks above that. Both base and candidate clusters run on the same machine, so noisy neighbors break thresholds. |
| vars | `QUERY_REGRESSION_ECS_IMAGE_ID` | Custom image built by `ecs-image/build-ecs-image.py`. |
The system disk is 50 GiB, which covers the image, a 16 GiB swapfile, the
checkout, and cold build caches (target dir, cargo registry, sccache). ENOSPC
stops the runner itself from writing logs, which GitHub reports as `The
operation was canceled` with no telemetry, indistinguishable from a
platform-side cancellation.
cloud-init masks `systemd-oomd`, disables `unattended-upgrades` /
`apt-daily-upgrade`, creates `/swapfile`, and sets `OOMPolicy=continue`
on the runner unit. Ubuntu 24.04 defaults to `DefaultOOMPolicy=stop`,
which SIGTERM-s the whole unit when rustc is OOM-killed and GitHub
reports `The operation was canceled` with no telemetry. Unattended
upgrades can do the same via `systemctl restart` of the runner after a
library update; GitHub then records `UserCancelled` even though the job
was still valid. Swap is a safety net for 16 GiB types, not a substitute
for 32 GiB; linking on swap is slow.
When a run dies with an unexplained `The operation was canceled` (no
"Canceled by" banner, healthy machine), re-dispatch with `keep_instance`
checked: the teardown job is skipped and the instance survives for
post-mortem inspection. The security group is egress-only, so inspect via
Cloud Assistant (`RunCommand`) or VNC: the runner's `_diag` logs under the
runner home record reconnects and worker crashes, `journalctl -u
ephemeral-github-runner.service` mirrors the console stream, `dmesg -T`
shows kernel OOM kills, and `machine-telemetry.log` in the job workspace
has the 30 s sampler history. The janitor sweep still deletes the instance
after its TTL, so finish the inspection within that window.
Trust model on the ECS path: the instance receives only the one-hour runner
registration token via user data and holds no cloud credentials; the Aliyun
AK/SK exist only in the control-plane jobs. Instance tags
(`managed-by=query-regression-ci`, `query-regression-run-id`) feed the janitor
sweep.
### Building and updating the ECS image
The runner `Dockerfile` in the parent directory stays the single source of the
tool contract. Build a new ECS image from it:
```bash
helm upgrade --install arc \
oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set-controller \
--namespace arc-systems \
--create-namespace \
--version 0.14.2
ALIBABA_CLOUD_ACCESS_KEY_ID=... ALIBABA_CLOUD_ACCESS_KEY_SECRET=... \
uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \
--region-id <region> --vswitch-id <vsw-...> --security-group-id <sg-...> \
--base-image-id <ubuntu-24.04-image-id>
```
Create the GitHub App secret in the runner namespace. Prefer an App limited to
`GreptimeTeam/greptimedb`:
The script boots a temporary builder instance, `docker build`s the runner
image, materializes `/opt/rustup`, `/opt/cargo`, `/usr/local/bin` tools, and
`/home/runner` (actions-runner) onto the host, installs the ephemeral-runner
systemd unit from `ecs-image/`, and snapshots a custom image. It prints the
image id; set it as `QUERY_REGRESSION_ECS_IMAGE_ID`, and bump
`RUNNER_IMAGE_EPOCH` in `query-regression.yml` at the same time so the target
cache invalidates. Builder sentinel polling requires the Cloud Assistant
agent, which Aliyun public Ubuntu images include.
```bash
kubectl -n arc-runners create secret generic greptimedb-arc-github-app \
--from-literal=github_app_id=<app-id> \
--from-literal=github_app_installation_id=<installation-id> \
--from-file=github_app_private_key=<private-key.pem>
```
The values files here reference that secret by name.
## Trust admission for PR runs
A maintainer applying the `query-regression` or `heavy-regression` label is
**trust admission for that exact PR revision**. `query-regression` runs the six
routine default cases; `heavy-regression` runs only the high-cardinality
`prom_remote_write_7913` remote-write case. 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 the desired regression label; cancel the old
run if it is no longer wanted. An existing label does not automatically rerun
the benchmark.
**trust admission for that exact PR revision**. `query-regression` runs the
six routine default cases; `heavy-regression` runs only the high-cardinality
`prom_remote_write_7913` remote-write case. `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 the desired
regression label; cancel the old run if it is no longer wanted. An existing
label does not automatically rerun the benchmark.
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.
Admission does not relax runner hardening or GitHub permissions. Keep the ECS
instance free of cloud credentials and long-lived tokens, keep the security
group egress-only, keep GitHub tokens least-privilege, and review workflow
changes before admission.
## 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. The
image builds `otelgen` from
The runner `Dockerfile` builds `otelgen` from
[`WenyXu/otelgen`](https://github.com/WenyXu/otelgen) commit
[`863a3f395d062c7322cc1de08a38774b7fdaa6c8`](https://github.com/WenyXu/otelgen/commit/863a3f395d062c7322cc1de08a38774b7fdaa6c8)
so trace cases do not download or compile tools during a benchmark run:
so trace cases do not download or compile tools during a benchmark run. It is
built only as a toolchain factory: `build-ecs-image.py` and
`bootstrap-runner-host.sh` both materialize its contents onto a host, and the
benchmark itself runs host-native.
```bash
docker build \
--platform linux/amd64 \
-f .github/runner-scale-sets/query-regression/Dockerfile \
-t greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner:latest \
.github/runner-scale-sets/query-regression
docker push greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner:latest
```
Deploy by digest, not mutable tag, by updating both image references in
`values-8-cores.yaml` after a rebuild. Update `RUNNER_IMAGE_DIGEST` and bump
`RUNNER_IMAGE_EPOCH` in `query-regression.yml` at the same time. 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 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.
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`, `otelgen` commit `863a3f395d062c7322cc1de08a38774b7fdaa6c8`,
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`.
Before builds, the workflow asserts the runner UID/GID (1001 in the ECS image;
overridable via `QUERY_REGRESSION_RUNNER_UID`/`QUERY_REGRESSION_RUNNER_GID`)
and exact tool versions: `libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.40.4`,
`Python 3.14.4`, `sccache 0.16.0`, `otelgen` commit
`863a3f395d062c7322cc1de08a38774b7fdaa6c8`, root-owned `rustup 1.29.0`, and
the image-baked `nightly-2026-03-21` Rust toolchain. `mold` and `python3`
come from apt at image-build time (not Ubuntu 24.04's default 3.12); bump
the Verify pins together with `QUERY_REGRESSION_ECS_IMAGE_ID` when the
image is rebuilt. 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.
workflow sets its warning-denying mold `RUSTFLAGS` directly, disables
automatic Rustup installation, and performs no runtime toolchain downloads.
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.
The workflow no longer uses GitHub `rust-cache`, `setup-protoc`, `setup-uv`,
or runtime Rust setup: the image establishes immutable executable state and
each instance's system disk holds only that run's Cargo data. Do not
reintroduce those actions unless the image contract changes.
## Capacity and persistent cache
## Capacity
`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.
Each run provisions its own ECS instance, so overlapping dispatches proceed
in parallel. There is no workflow `concurrency` group. All Cargo state
(`CARGO_HOME` including registry/git, `CARGO_TARGET_DIR`, sccache, cache
metadata) lives on the 50 GiB system disk and is discarded with the VM.
`RUSTUP_HOME=/opt/rustup` and `/opt/cargo/bin` are image-owned. The runner
sets `RUSTC_WRAPPER=/usr/local/bin/sccache`,
`SCCACHE_DIR=/home/runner/.cache/sccache`, `SCCACHE_CACHE_SIZE=10G`, and
`CARGO_INCREMENTAL=0`. sccache uses its local disk backend and self-evicts
at 10G. Base and candidate builds share the target on that disk; Cargo
fingerprints invalidate source and dependency changes.
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 \
--reset-values --wait \
-f .github/runner-scale-sets/query-regression/values-8-cores.yaml
```
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`.
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.
## Canary and rollback
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.
Accept the canary only when all of the following hold:
- 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.
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.
The workflow reports `du` / `df` in telemetry. It does not try to reclaim
space across runs: a cold 50 GiB disk that fills up fails the build.
## Future optional phases
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.
The current phase uses a materialized runner toolchain with sccache 0.16.0 and
no shared cache service. Optional follow-ups are an image additionally seeded
with `cargo fetch --locked` results, or an internal read/write sccache
backend or Cargo/Git mirror. Evaluate them only if cold compile time on the
system disk becomes the bottleneck.
.github/runner-scale-sets/query-regression/cache-pvc.yaml
-15
View File
@@ -1,15 +0,0 @@
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/ecs-image/bootstrap-runner-host.sh
+195
View File
@@ -0,0 +1,195 @@
#!/usr/bin/env bash
# Copyright 2023 Greptime Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Bootstrap an Ubuntu 24.04 host as a query-regression self-hosted runner.
#
# This is the host-native equivalent of the ECS image contract: it
# installs docker-ce from Docker's official repository, builds the runner
# image from the sibling Dockerfile (the single source of the tool contract),
# materializes the toolchain onto the host, and optionally registers the
# runner and installs its systemd service.
#
# Large state lives under --data-root (a data disk): docker's data-root, the
# runner home, and the rustup/cargo toolchain roots, symlinked back to the
# hard-coded contract paths (/home/runner, /opt/rustup, /opt/cargo).
#
# The script is idempotent: re-running refreshes the toolchain in place, and
# the runner registration (.runner/.credentials) and _work survive. A live
# runner service is stopped first and restarted at the end.
#
# Usage:
# sudo bash bootstrap-runner-host.sh --repo-dir /path/to/greptimedb
# sudo RUNNER_TOKEN=<token> bash bootstrap-runner-host.sh --repo-dir . \
# --register --repo-url https://github.com/<owner>/<repo>
#
# Options:
# --repo-dir PATH greptimedb checkout containing the runner Dockerfile (required)
# --data-root PATH data disk mount point (default: /data)
# --uid / --gid N runner user id (default: 3141; the ECS image contract is 1001)
# --register also register the runner and install the systemd service
# --repo-url URL repository URL for registration (required with --register)
# --runner-name N runner name (default: qreg-host)
# --labels L runner labels (default: perf-regression-8-cores)
#
# With --register, provide a fresh registration token via RUNNER_TOKEN.
set -euo pipefail
REPO_DIR=""
DATA_ROOT="/data"
RUNNER_UID="3141"
RUNNER_GID="3141"
REGISTER="false"
REPO_URL=""
RUNNER_NAME="qreg-host"
RUNNER_LABELS="perf-regression-8-cores"
IMAGE_TAG="qreg-runner:manual"
while [[ $# -gt 0 ]]; do
case "$1" in
--repo-dir) REPO_DIR="$2"; shift 2 ;;
--data-root) DATA_ROOT="$2"; shift 2 ;;
--uid) RUNNER_UID="$2"; shift 2 ;;
--gid) RUNNER_GID="$2"; shift 2 ;;
--register) REGISTER="true"; shift ;;
--repo-url) REPO_URL="$2"; shift 2 ;;
--runner-name) RUNNER_NAME="$2"; shift 2 ;;
--labels) RUNNER_LABELS="$2"; shift 2 ;;
-h | --help) sed -n '17,40p' "$0"; exit 0 ;;
*) echo "Unknown argument: $1" >&2; exit 2 ;;
esac
done
[[ $EUID -eq 0 ]] || { echo "Run as root, e.g. sudo bash $0 ..." >&2; exit 1; }
[[ -n "${REPO_DIR}" ]] || { echo "--repo-dir is required" >&2; exit 2; }
DOCKERFILE="${REPO_DIR}/.github/runner-scale-sets/query-regression/Dockerfile"
[[ -f "${DOCKERFILE}" ]] || { echo "Dockerfile not found at ${DOCKERFILE}" >&2; exit 1; }
if [[ "${REGISTER}" == "true" ]]; then
[[ -n "${REPO_URL}" ]] || { echo "--repo-url is required with --register" >&2; exit 2; }
[[ -n "${RUNNER_TOKEN:-}" ]] || { echo "RUNNER_TOKEN env is required with --register" >&2; exit 2; }
fi
step() { printf '\n==> %s\n' "$*"; }
# A re-run refreshes files under the (possibly live) runner service, so stop
# it first and restart it at the end (unless --register re-creates it).
RUNNER_SERVICE="$(systemctl list-units --type=service --all --no-legend 'actions.runner.*' 2>/dev/null | awk '{print $1}' | head -n 1 || true)"
RUNNER_WAS_ACTIVE="false"
if [[ -n "${RUNNER_SERVICE}" ]] && systemctl is-active --quiet "${RUNNER_SERVICE}"; then
step "Stop running runner service ${RUNNER_SERVICE}"
systemctl stop "${RUNNER_SERVICE}"
RUNNER_WAS_ACTIVE="true"
fi
step "Install docker-ce from Docker's official repository"
for pkg in docker.io docker-doc docker-compose docker-compose-v2 podman-docker containerd runc; do
apt-get remove -y "${pkg}" 2>/dev/null || true
done
apt-get update
apt-get install -y ca-certificates curl
install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc
chmod a+r /etc/apt/keyrings/docker.asc
# shellcheck disable=SC1091
. /etc/os-release
tee /etc/apt/sources.list.d/docker.sources > /dev/null <<EOF
Types: deb
URIs: https://download.docker.com/linux/ubuntu
Suites: ${UBUNTU_CODENAME:-$VERSION_CODENAME}
Components: stable
Architectures: $(dpkg --print-architecture)
Signed-By: /etc/apt/keyrings/docker.asc
EOF
apt-get update
apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin
step "Move docker data-root to ${DATA_ROOT}/docker"
systemctl stop docker containerd
mkdir -p "${DATA_ROOT}/docker" "${DATA_ROOT}/containerd"
if [[ -f /etc/docker/daemon.json ]] && ! grep -q '"data-root"' /etc/docker/daemon.json; then
echo "ERROR: /etc/docker/daemon.json exists without data-root; merge manually." >&2
exit 1
fi
echo "{ \"data-root\": \"${DATA_ROOT}/docker\" }" > /etc/docker/daemon.json
# containerd keeps its own content store (where build-time image data lands);
# move it off the system disk too.
mkdir -p /etc/containerd
containerd config default | sed "s|root = \"/var/lib/containerd\"|root = \"${DATA_ROOT}/containerd\"|" > /etc/containerd/config.toml
systemctl start containerd docker
docker info | grep "Docker Root Dir"
step "Build the runner image"
docker build --platform linux/amd64 -f "${DOCKERFILE}" -t "${IMAGE_TAG}" "${REPO_DIR}"
step "Materialize the toolchain onto the host"
mkdir -p "${DATA_ROOT}/opt" "${DATA_ROOT}/runner" /opt /home
container="$(docker create "${IMAGE_TAG}")"
trap 'docker rm -f "${container}" >/dev/null 2>&1 || true' EXIT
# The `/.` suffix copies directory *contents*, so re-running over an existing
# materialization refreshes it in place instead of nesting runner/runner.
# Files absent from the image (runner registration, _work) are preserved.
docker cp "${container}:/home/runner/." "${DATA_ROOT}/runner"
docker cp "${container}:/opt/rustup/." "${DATA_ROOT}/opt/rustup"
docker cp "${container}:/opt/cargo/." "${DATA_ROOT}/opt/cargo"
for tool in uv uvx otelgen sccache; do
docker cp "${container}:/usr/local/bin/${tool}" "/usr/local/bin/${tool}"
done
docker rm "${container}" > /dev/null
trap - EXIT
ln -sfn "${DATA_ROOT}/runner" /home/runner
ln -sfn "${DATA_ROOT}/opt/rustup" /opt/rustup
ln -sfn "${DATA_ROOT}/opt/cargo" /opt/cargo
docker image rm "${IMAGE_TAG}" > /dev/null
step "Install system packages"
apt-get install -y --no-install-recommends \
build-essential clang cmake git gzip jq libprotobuf-dev libssl-dev mold \
openssh-client pkg-config protobuf-compiler python3 sudo tar unzip wget xz-utils zip zstd
step "Create runner user (${RUNNER_UID}:${RUNNER_GID}) and environment"
getent group "${RUNNER_GID}" > /dev/null || groupadd -g "${RUNNER_GID}" runner
id -u runner > /dev/null 2>&1 || useradd -u "${RUNNER_UID}" -g "${RUNNER_GID}" -d /home/runner -s /bin/bash runner
chown -R "${RUNNER_UID}:${RUNNER_GID}" "${DATA_ROOT}/runner"
echo 'PATH=/opt/cargo/bin:/usr/local/bin:/usr/bin:/bin' > "${DATA_ROOT}/runner/.env"
chown "${RUNNER_UID}:${RUNNER_GID}" "${DATA_ROOT}/runner/.env"
if [[ "${REGISTER}" == "true" ]]; then
step "Register runner ${RUNNER_NAME} (labels: ${RUNNER_LABELS})"
if [[ -n "${RUNNER_SERVICE}" ]]; then
(cd /home/runner && ./svc.sh uninstall) || true
fi
runuser -u runner -- bash -c "cd /home/runner && HOME=/home/runner ./config.sh \
--url '${REPO_URL}' --token '${RUNNER_TOKEN}' --name '${RUNNER_NAME}' \
--labels '${RUNNER_LABELS}' --unattended --replace --disableupdate"
step "Install and start the runner service"
cd /home/runner
./svc.sh install runner
./svc.sh start
./svc.sh status
else
step "Skipping registration (pass --register --repo-url ... with RUNNER_TOKEN to enable)"
if [[ "${RUNNER_WAS_ACTIVE}" == "true" ]]; then
step "Restart runner service ${RUNNER_SERVICE}"
systemctl start "${RUNNER_SERVICE}"
systemctl --no-pager status "${RUNNER_SERVICE}" || true
fi
fi
step "Done"
echo "Runner home: /home/runner -> ${DATA_ROOT}/runner"
echo "Before each workflow run on this persistent host, clean transient cargo state:"
echo " sudo rm -f /home/runner/.cargo/.package-cache"
echo " sudo rm -rf /home/runner/.cargo/.global-cache"
.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py
+392
View File
@@ -0,0 +1,392 @@
#!/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.
# PEP 723 inline metadata (see .github/scripts/aliyun-ecs-runner-provision.py
# for the convention).
# /// script
# requires-python = ">=3.10"
# dependencies = [
# "alibabacloud_ecs20140526>=4.1.0,<6",
# "alibabacloud_tea_openapi>=0.3.12,<1",
# ]
# ///
"""Build the query-regression ECS custom image (manual ops tool).
Boots a temporary pay-as-you-go ECS instance from a public Ubuntu 24.04 image,
builds the existing runner container image (the Dockerfile in the parent
directory remains the single source of the tool contract), materializes the
tool directories onto the host filesystem so the workflow's "Verify runner
image tools" step holds unchanged, installs the ephemeral-runner systemd unit,
and snapshots the result as a custom image. The temporary instance is deleted
afterwards.
Sentinel polling reads the instance's serial console output
(GetInstanceConsoleOutput) and looks for marker lines the user data writes to
/dev/console. This has no in-guest agent dependency.
Usage:
uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \
--region-id cn-hangzhou --vswitch-id vsw-... --security-group-id sg-... \
--base-image-id ubuntu_24_04_x64_20G_alibase_*.vhd
"""
from __future__ import annotations
import argparse
import base64
import os
import time
from pathlib import Path
ASSETS_DIR = Path(__file__).resolve().parent
DONE_MARKER = "QREG_IMAGE_BUILD_DONE"
FAILED_MARKER = "QREG_IMAGE_BUILD_FAILED"
POLL_INTERVAL_SECONDS = 15
CONSOLE_POLL_INTERVAL_SECONDS = 30
BUILD_TIMEOUT_SECONDS = 60 * 60
# Same apt package contract as the runner Dockerfile; the base
# actions-runner image is Ubuntu 24.04, so an Ubuntu 24.04 host resolves the
# same tool versions (protoc 3.21.12, mold 2.40.4, Python 3.14.4).
# Docker itself comes from Docker's official repository (docker-ce), not the
# distribution-packaged docker.io.
APT_PACKAGES = [
"build-essential",
"ca-certificates",
"clang",
"cmake",
"curl",
"git",
"gpg",
"gzip",
"jq",
"libprotobuf-dev",
"libssl-dev",
"mold",
"openssh-client",
"pkg-config",
"protobuf-compiler",
"python3",
"sudo",
"tar",
"unzip",
"wget",
"xz-utils",
"zip",
"zstd",
]
DOCKER_CE_PACKAGES = "docker-ce docker-ce-cli containerd.io docker-buildx-plugin"
def render_user_data(dockerfile: str, start_runner: str, unit: str) -> str:
dockerfile_b64 = base64.b64encode(dockerfile.encode()).decode()
start_runner_b64 = base64.b64encode(start_runner.encode()).decode()
unit_b64 = base64.b64encode(unit.encode()).decode()
packages = " ".join(APT_PACKAGES)
return f"""#!/bin/bash
set -euo pipefail
trap 'echo "{FAILED_MARKER} at line $LINENO" > /dev/console' ERR
# Stream the full build log to the serial console so the poller (and anyone
# watching GetInstanceConsoleOutput) sees real progress, not a silent login
# prompt for the whole build.
exec > >(tee -a /dev/console) 2>&1
apt-get update
DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends {packages}
# Docker from the official repository, not the distribution-packaged docker.io.
install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /etc/apt/keyrings/docker.gpg
chmod a+r /etc/apt/keyrings/docker.gpg
. /etc/os-release
echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu ${{VERSION_CODENAME}} stable" \
> /etc/apt/sources.list.d/docker.list
apt-get update
DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends {DOCKER_CE_PACKAGES}
base64 -d > /tmp/Dockerfile <<'EOF'
{dockerfile_b64}
EOF
mkdir -p /tmp/image-context
docker build --platform linux/amd64 -f /tmp/Dockerfile -t qreg-runner:local /tmp/image-context
# Materialize the tool contract onto the host filesystem.
container="$(docker create qreg-runner:local)"
trap 'docker rm -f "${{container}}" >/dev/null 2>&1 || true' EXIT
docker cp "${{container}}:/opt/rustup" /opt/rustup
docker cp "${{container}}:/opt/cargo" /opt/cargo
for tool in uv uvx otelgen sccache; do
docker cp "${{container}}:/usr/local/bin/${{tool}}" "/usr/local/bin/${{tool}}"
done
docker cp "${{container}}:/home/runner" /home/runner
# Runner identity mirrors the container image (UID/GID 1001).
groupadd --gid 1001 runner 2>/dev/null || true
useradd --uid 1001 --gid 1001 --home-dir /home/runner --shell /bin/bash runner 2>/dev/null || true
chown -R 1001:1001 /home/runner
mkdir -p /opt/ephemeral-github-runner
base64 -d > /opt/ephemeral-github-runner/start-runner.sh <<'EOF'
{start_runner_b64}
EOF
chmod 0755 /opt/ephemeral-github-runner/start-runner.sh
base64 -d > /etc/systemd/system/ephemeral-github-runner.service <<'EOF'
{unit_b64}
EOF
systemctl daemon-reload
systemctl enable ephemeral-github-runner.service
# Keep the image free of the build-time docker state (also shrinks the
# snapshot: buildkit cache is several GB).
docker rm -f "${{container}}" >/dev/null
docker system prune -af >/dev/null
trap - EXIT
echo "{DONE_MARKER}" > /dev/console
"""
def make_ecs_client(region_id: str):
from alibabacloud_ecs20140526.client import Client as EcsClient
from alibabacloud_tea_openapi.models import Config as OpenApiConfig
return EcsClient(
OpenApiConfig(
access_key_id=os.environ["ALIBABA_CLOUD_ACCESS_KEY_ID"],
access_key_secret=os.environ["ALIBABA_CLOUD_ACCESS_KEY_SECRET"],
region_id=region_id,
endpoint=f"ecs.{region_id}.aliyuncs.com",
)
)
def wait_for_instance_status(client, region_id: str, instance_id: str, wanted: str, deadline: float) -> None:
from alibabacloud_ecs20140526 import models as ecs_models
import json
while time.monotonic() < deadline:
response = call_api_with_retry(
lambda: client.describe_instances(
ecs_models.DescribeInstancesRequest(
region_id=region_id, instance_ids=json.dumps([instance_id])
)
),
"DescribeInstances",
)
instances = response.body.instances.instance
if instances and instances[0].status == wanted:
return
time.sleep(POLL_INTERVAL_SECONDS)
raise TimeoutError(f"Instance {instance_id} did not reach status {wanted} in time")
# Aliyun error codes worth retrying: throttling and server-side faults. Client
# errors (4xx: permissions, bad parameters) are configuration problems and must
# fail fast instead of being retried.
TRANSIENT_ERROR_CODES = {"Throttling", "Throttling.User", "InternalError", "ServiceUnavailable"}
def is_transient(error: Exception) -> bool:
# UnretryableException comes from the darabonba network layer: connection
# reset, read timeout, DNS blip.
if type(error).__name__ == "UnretryableException":
return True
code = getattr(error, "code", "") or ""
status = getattr(error, "statusCode", None)
return code in TRANSIENT_ERROR_CODES or (isinstance(status, int) and status >= 500)
def call_api_with_retry(fn, description: str, attempts: int = 5):
"""Retry transient API/network failures; the build is too long to die on a blip.
Only call this with idempotent or safely-repeatable requests (reads,
StopInstance, CreateImage). Never with RunInstances: if the request
succeeded but the response was lost, a retry double-creates instances.
"""
for attempt in range(1, attempts + 1):
try:
return fn()
except Exception as error: # noqa: BLE001
if attempt == attempts or not is_transient(error):
raise
print(f"{description} failed (attempt {attempt}/{attempts}): {error}", flush=True)
time.sleep(POLL_INTERVAL_SECONDS)
raise RuntimeError("unreachable: retry loop exited without returning")
def read_console_output(client, region_id: str, instance_id: str) -> str:
"""Fetch the instance's serial console output; no in-guest agent needed."""
from alibabacloud_ecs20140526 import models as ecs_models
response = call_api_with_retry(
lambda: client.get_instance_console_output(
ecs_models.GetInstanceConsoleOutputRequest(region_id=region_id, instance_id=instance_id)
),
"GetInstanceConsoleOutput",
)
return base64.b64decode(response.body.console_output or "").decode("utf-8", "replace")
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID"))
parser.add_argument("--vswitch-id", default=os.environ.get("ALIYUN_ECS_VSWITCH_ID"))
parser.add_argument("--security-group-id", default=os.environ.get("ALIYUN_ECS_SECURITY_GROUP_ID"))
parser.add_argument("--base-image-id", default=os.environ.get("ALIYUN_ECS_BASE_IMAGE_ID"))
parser.add_argument("--resource-group-id", default=os.environ.get("ALIYUN_ECS_RESOURCE_GROUP_ID"))
parser.add_argument("--instance-type", default="ecs.g7.xlarge")
parser.add_argument("--image-name", default=None, help="Defaults to a timestamped name.")
args = parser.parse_args()
for name in ("region_id", "vswitch_id", "security_group_id", "base_image_id"):
if not getattr(args, name):
raise SystemExit(f"Missing required configuration: --{name.replace('_', '-')}")
from alibabacloud_ecs20140526 import models as ecs_models
client = make_ecs_client(args.region_id)
image_name = args.image_name or time.strftime(
"greptimedb-query-regression-runner-%Y%m%d%H%M%S", time.gmtime()
)
user_data = base64.b64encode(
render_user_data(
(ASSETS_DIR.parent / "Dockerfile").read_text(),
(ASSETS_DIR / "start-runner.sh").read_text(),
(ASSETS_DIR / "ephemeral-github-runner.service").read_text(),
).encode()
).decode()
instance_id = None
try:
response = client.run_instances(
ecs_models.RunInstancesRequest(
region_id=args.region_id,
image_id=args.base_image_id,
resource_group_id=args.resource_group_id,
instance_type=args.instance_type,
v_switch_id=args.vswitch_id,
security_group_id=args.security_group_id,
instance_name=f"build-{image_name}",
description="Temporary builder for the query-regression ECS image",
amount=1,
instance_charge_type="PostPaid",
spot_strategy="NoSpot",
internet_charge_type="PayByTraffic",
internet_max_bandwidth_out=100,
# Peak usage is ~18G (OS+apt, docker image, and the materialized
# toolchain coexist briefly): deliberately tight, and a smaller
# disk makes the image snapshot faster. Instances created from
# the image get a larger system disk from the provision side.
system_disk=ecs_models.RunInstancesRequestSystemDisk(
category="cloud_essd", size="20"
),
user_data=user_data,
tag=[
ecs_models.RunInstancesRequestTag(key="managed-by", value="query-regression-ci"),
ecs_models.RunInstancesRequestTag(key="role", value="image-builder"),
],
)
)
instance_id = response.body.instance_id_sets.instance_id_set[0]
print(f"Builder instance: {instance_id}", flush=True)
wait_for_instance_status(client, args.region_id, instance_id, "Running", time.monotonic() + 10 * 60)
deadline = time.monotonic() + BUILD_TIMEOUT_SECONDS
while time.monotonic() < deadline:
console = read_console_output(client, args.region_id, instance_id)
if DONE_MARKER in console:
break
if FAILED_MARKER in console:
tail = "\n".join(console.splitlines()[-20:])
raise RuntimeError(f"Image build failed on the builder; console tail:\n{tail}")
lines = console.splitlines()
print(f"Build in progress; last console line: {lines[-1] if lines else '(none yet)'}", flush=True)
time.sleep(CONSOLE_POLL_INTERVAL_SECONDS)
else:
raise TimeoutError("Image build did not finish in time")
print("Stopping builder before image creation", flush=True)
try:
call_api_with_retry(
lambda: client.stop_instance(
ecs_models.StopInstanceRequest(
instance_id=instance_id,
# The instance is deleted right after the snapshot, but
# there is no reason to keep billing vCPU during the stop.
stopped_mode="StopCharging",
)
),
"StopInstance",
)
except Exception as error: # noqa: BLE001
# Instance families with local disks do not support StopCharging.
print(f"StopCharging unavailable ({error}); stopping with default mode", flush=True)
call_api_with_retry(
lambda: client.stop_instance(ecs_models.StopInstanceRequest(instance_id=instance_id)),
"StopInstance",
)
wait_for_instance_status(client, args.region_id, instance_id, "Stopped", time.monotonic() + 10 * 60)
image = call_api_with_retry(
lambda: client.create_image(
ecs_models.CreateImageRequest(
region_id=args.region_id, instance_id=instance_id, image_name=image_name
)
),
"CreateImage",
)
image_id = image.body.image_id
deadline = time.monotonic() + 60 * 60
while time.monotonic() < deadline:
description = call_api_with_retry(
lambda: client.describe_images(
ecs_models.DescribeImagesRequest(region_id=args.region_id, image_id=image_id)
),
"DescribeImages",
)
images = description.body.images.image
if images:
status = images[0].status
print(f"Image {image_id} status: {status} (progress {images[0].progress})", flush=True)
if status == "Available":
break
time.sleep(POLL_INTERVAL_SECONDS)
else:
raise TimeoutError(
f"Image {image_id} did not become Available in time. The snapshot "
f"continues server-side: check its status in the console and reuse it "
f"once Available instead of rebuilding."
)
print(f"Custom image ready: {image_id} ({image_name})", flush=True)
print(f"Set the repo variable QUERY_REGRESSION_ECS_IMAGE_ID={image_id}", flush=True)
return 0
finally:
if instance_id:
try:
client.delete_instance(
ecs_models.DeleteInstanceRequest(instance_id=instance_id, force=True)
)
print(f"Deleted builder instance {instance_id}", flush=True)
except Exception as error: # noqa: BLE001
print(f"Failed to delete builder instance {instance_id}: {error}", flush=True)
if __name__ == "__main__":
raise SystemExit(main())
.github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service
+45
View File
@@ -0,0 +1,45 @@
# 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.
[Unit]
Description=Ephemeral GitHub Actions runner
After=network-online.target
Wants=network-online.target
# cloud-final runs the provision user data, which writes
# /etc/ephemeral-github-runner.env and the drop-ins below. Ordering after it
# makes the runner come up configured on first boot, no restart needed.
# (cloud-init's `systemctl restart` in user data remains as the bridge for
# images built before this ordering existed.)
After=cloud-final.service
[Service]
Type=simple
# The toolchain lives in the image at /opt/cargo/bin (the Dockerfile's ENV
# PATH does not survive materialization onto the host). Publish it here so
# every runner job inherits it. /etc/ephemeral-github-runner.env currently
# carries the same PATH line as a bridge for images built before this
# directive existed; EnvironmentFile is applied after Environment=, and the
# two values are kept identical, so there is no conflict.
Environment=PATH=/opt/cargo/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
EnvironmentFile=/etc/ephemeral-github-runner.env
ExecStart=/opt/ephemeral-github-runner/start-runner.sh
Restart=no
# If the kernel OOM killer kills a job child (rustc/cargo), keep the runner
# service alive so the job is reported as a normal failure with logs instead
# of the runner vanishing and GitHub reporting a cancellation. cloud-init's
# oom.conf drop-in carries the same directive for images built before this.
OOMPolicy=continue
[Install]
WantedBy=multi-user.target
.github/runner-scale-sets/query-regression/ecs-image/start-runner.sh
+37
View File
@@ -0,0 +1,37 @@
#!/bin/bash
# Copyright 2023 Greptime Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Configure and start the ephemeral GitHub Actions runner on an ECS instance.
# Baked into the custom image at /opt/ephemeral-github-runner/ by
# build-ecs-image.py; the per-run environment is written by cloud-init to
# /etc/ephemeral-github-runner.env.
set -euo pipefail
# shellcheck disable=SC1091
source /etc/ephemeral-github-runner.env
export HOME=/home/runner
cd /home/runner
if [[ ! -f /home/runner/.runner ]]; then
runuser -u runner -- ./config.sh \
--url "${REPO_URL}" \
--token "${RUNNER_TOKEN}" \
--name "${RUNNER_NAME}" \
--labels "${RUNNER_LABELS}" \
--ephemeral --unattended --replace --disableupdate
fi
exec runuser -u runner -- ./run.sh
.github/runner-scale-sets/query-regression/values-8-cores.yaml
-146
View File
@@ -1,146 +0,0 @@
githubConfigUrl: "https://github.com/GreptimeTeam/greptimedb"
githubConfigSecret: greptimedb-arc-github-app
runnerScaleSetName: "perf-regression-8-cores"
minRunners: 0
maxRunners: 1
template:
spec:
automountServiceAccountToken: false
activeDeadlineSeconds: 12600
nodeSelector:
kubernetes.io/hostname: minipc-3
securityContext:
runAsNonRoot: true
runAsUser: 1001
runAsGroup: 1001
fsGroup: 1001
fsGroupChangePolicy: OnRootMismatch
seccompProfile:
type: RuntimeDefault
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:e713b294e23b7e15184e558866c90025e59930033e72c97650dbc7f1ca022d11
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: 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:e713b294e23b7e15184e558866c90025e59930033e72c97650dbc7f1ca022d11
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
@@ -1,2 +0,0 @@
minRunners: 0
maxRunners: 0
.github/scripts/aliyun-ecs-runner-provision.py
+488
View File
@@ -0,0 +1,488 @@
#!/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.
# This is the first `.github/scripts/` user of PEP 723 inline script metadata:
# `uv run` installs the pinned Aliyun SDK before executing. The SDK import stays
# lazy (inside `make_ecs_client`) so unit tests can import this module with a
# plain stdlib interpreter.
# /// script
# requires-python = ">=3.10"
# dependencies = [
# # Pin exactly after the first live run against the Aliyun API.
# "alibabacloud_ecs20140526>=4.1.0,<6",
# "alibabacloud_tea_openapi>=0.3.12,<1",
# ]
# ///
"""Provision an ephemeral Aliyun ECS query-regression runner.
Creates one pay-as-you-go ECS instance from the query-regression custom image
and waits until the instance registers itself as an ephemeral GitHub Actions
runner with a per-run label. Build caches live on the instance system disk
and disappear with the VM. The instance id is written to job outputs as soon
as the VM exists so the teardown job is the single DeleteInstance caller (a
timeout here used to delete and then teardown deleted again). On runner
online timeout the console is dumped and the job fails; teardown releases
the instance.
The Aliyun credentials are used only in this control-plane job (ubuntu-latest)
and are never passed to the ECS instance; the instance receives only a
short-lived GitHub runner registration token via user data.
"""
from __future__ import annotations
import argparse
import base64
import json
import os
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
RUNNER_NAME_PREFIX = "qreg-ecs"
RUNNER_LABEL_PREFIX = "query-regression-ecs"
MANAGED_BY_TAG_KEY = "managed-by"
MANAGED_BY_TAG_VALUE = "query-regression-ci"
RUN_TAG_KEY = "query-regression-run-id"
# Runner cache paths on the instance system disk. They are created empty
# every provision and deleted with the VM.
CACHE_PATHS = (
"/home/runner/.cargo/registry",
"/home/runner/.cargo/git",
"/home/runner/query-regression-target",
"/home/runner/query-regression-cache-meta",
"/home/runner/.cache/sccache",
)
RUNNER_ONLINE_TIMEOUT_SECONDS = 10 * 60
POLL_INTERVAL_SECONDS = 5
# Nightly thin-LTO of greptime peaks above ecs.c9i.2xlarge's 16 GiB. A swap
# file plus masking systemd-oomd lets the kernel reclaim rustc pages instead
# of SIGTERM-ing the runner service cgroup (GitHub then reports the step as
# "The operation was canceled"). Sized to match that 16 GiB instance; a
# 32 GiB type still benefits as a safety net. Lives on the system disk next
# to the cold build caches.
SWAP_FILE = "/swapfile"
SWAP_SIZE_GIB = 16
# Image (~12G), 16G swap, checkout, base+candidate data homes, and cold
# build caches. Deleted with the instance.
SYSTEM_DISK_GIB = 50
@dataclass(frozen=True)
class ProvisionConfig:
region_id: str
vswitch_id: str
security_group_id: str
image_id: str
instance_type: str
repo: str
run_id: str
github_token: str
# Optional resource group; required when the RAM grant is scoped to one.
resource_group_id: str | None = None
# Runner identity inside the image; the workflow asserts the same values.
runner_uid: str = "1001"
runner_gid: str = "1001"
def runner_name_for_run(run_id: str) -> str:
return f"{RUNNER_NAME_PREFIX}-{run_id}"
def runner_label_for_run(run_id: str) -> str:
return f"{RUNNER_LABEL_PREFIX}-{run_id}"
def render_user_data(
runner_name: str,
runner_label: str,
runner_token: str,
repo: str,
runner_uid: str = "1001",
runner_gid: str = "1001",
) -> str:
"""Render the cloud-init shell script for the runner instance."""
destinations = " ".join(f'"{dst}"' for dst in CACHE_PATHS)
cache_setup = f"""# Caches live on the system disk, are deleted with the instance, and every
# run compiles cold. Within-run reuse (base warming candidate via the shared
# target dir and sccache) still applies.
mkdir -p {destinations}
chown -R {runner_uid}:{runner_gid} /home/runner"""
swap_setup = f"""# Mask systemd-oomd before enabling swap: Ubuntu 24.04 kills the whole
# service cgroup on PSI pressure, and swap thrashing looks like pressure.
systemctl disable --now systemd-oomd.socket systemd-oomd.service || true
systemctl mask systemd-oomd.socket systemd-oomd.service || true
# apt-daily-upgrade / unattended-upgrades can `systemctl restart` services
# whose libraries were updated. That SIGTERM-s the runner mid-job; GitHub
# records UserCancelled (job still valid, child exit 143). These VMs live
# ~1h; security updates belong in the image, not at job time.
systemctl disable --now unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true
systemctl mask unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true
systemctl stop unattended-upgrades.service || true
killall -9 unattended-upgr || true
cat > /etc/apt/apt.conf.d/99disable-auto-updates <<'APTEOF'
APT::Periodic::Update-Package-Lists "0";
APT::Periodic::Unattended-Upgrade "0";
APT::Periodic::Download-Upgradeable-Packages "0";
APTEOF
if [[ ! -f "{SWAP_FILE}" ]]; then
fallocate --length {SWAP_SIZE_GIB}G "{SWAP_FILE}"
chmod 600 "{SWAP_FILE}"
mkswap "{SWAP_FILE}"
fi
swapon "{SWAP_FILE}"
sysctl --write vm.swappiness=10
swapon --show
free --human"""
return f"""#!/bin/bash
set -euo pipefail
{cache_setup}
{swap_setup}
cat > /etc/ephemeral-github-runner.env <<'ENVEOF'
RUNNER_NAME={runner_name}
RUNNER_LABELS={runner_label}
RUNNER_TOKEN={runner_token}
REPO_URL=https://github.com/{repo}
# The container image baked /opt/cargo/bin into PATH via Dockerfile ENV; on
# the host nothing inherits that, so publish it through the unit's
# EnvironmentFile: runner job processes inherit the runner's environment.
PATH=/opt/cargo/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
ENVEOF
# Stream the runner service's own logs (config output, job lifecycle, crash
# messages) to the serial console, next to the kernel's OOM-killer records.
# Teardown dumps the console tail before deleting the instance, so the
# witness of what killed the runner survives the machine.
mkdir -p /etc/systemd/system/ephemeral-github-runner.service.d
cat > /etc/systemd/system/ephemeral-github-runner.service.d/console.conf <<'CONFEOF'
[Service]
StandardOutput=journal+console
StandardError=journal+console
CONFEOF
# systemd's DefaultOOMPolicy=stop SIGTERM-s the whole unit when *any*
# process in the cgroup is OOM-killed (typically rustc during thin LTO).
# That is the "Session terminated, killing shell..." / job Canceled path:
# runuser dies, GitHub deletes the ephemeral registration, always() steps
# never run. continue keeps the runner up so cargo can report signal 9.
cat > /etc/systemd/system/ephemeral-github-runner.service.d/oom.conf <<'CONFEOF'
[Service]
OOMPolicy=continue
CONFEOF
systemctl daemon-reload
# restart (not start): the image enables this unit, so it may already
# have come up without the drop-ins if it raced cloud-init. restart
# applies OOMPolicy=continue to the running cgroup. --no-block matters:
# newer image units carry After=cloud-final.service, and this script IS
# cloud-final — a blocking restart would deadlock until TimeoutStartSec.
systemctl restart --no-block ephemeral-github-runner.service
"""
def encode_user_data(script: str) -> str:
return base64.b64encode(script.encode("utf-8")).decode("ascii")
def github_api(token: str, method: str, path: str, body: dict | None = None) -> dict:
data = json.dumps(body).encode("utf-8") if body is not None else None
request = urllib.request.Request(
f"https://api.github.com{path}",
data=data,
method=method,
headers={
"Accept": "application/vnd.github+json",
"Authorization": f"Bearer {token}",
"X-GitHub-Api-Version": "2022-11-28",
},
)
try:
with urllib.request.urlopen(request, timeout=30) as response:
return json.loads(response.read().decode("utf-8"))
except urllib.error.HTTPError as error:
# GitHub's error body says exactly why (e.g. "Must have admin rights to
# Repository" for a PAT without the required scope); surface it instead
# of a bare "HTTP Error 403".
body = error.read().decode("utf-8", "replace")
raise SystemExit(
f"GitHub API {method} {path} failed: HTTP {error.code}: {body}\n"
"The token comes from the GH_PERSONAL_ACCESS_TOKEN secret; it needs "
"'repo' scope (classic PAT) or 'Administration: write' on the "
"repository (fine-grained PAT)."
) from error
def create_registration_token(github_token: str, repo: str) -> str:
response = github_api(
github_token,
"POST",
f"/repos/{repo}/actions/runners/registration-token",
body={},
)
return response["token"]
def find_runner_by_name(github_token: str, repo: str, name: str) -> dict | None:
page = 1
while True:
response = github_api(
github_token,
"GET",
f"/repos/{repo}/actions/runners?per_page=100&page={page}",
)
runners = response.get("runners", [])
for runner in runners:
if runner.get("name") == name:
return runner
if len(runners) < 100:
return None
page += 1
def make_ecs_client(config: ProvisionConfig):
from alibabacloud_ecs20140526.client import Client as EcsClient
from alibabacloud_tea_openapi.models import Config as OpenApiConfig
access_key_id = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_ID", "")
access_key_secret = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_SECRET", "")
if not access_key_id or not access_key_secret:
raise SystemExit(
"ALIBABA_CLOUD_ACCESS_KEY_ID/SECRET are empty; in CI they come from the "
"ALICLOUD_ECS_ACCESS_KEY_ID/SECRET repository secrets (a missing or "
"misnamed secret expands to an empty string)."
)
return EcsClient(
OpenApiConfig(
access_key_id=access_key_id,
access_key_secret=access_key_secret,
region_id=config.region_id,
endpoint=f"ecs.{config.region_id}.aliyuncs.com",
)
)
def wait_for_instance_status(client, region_id: str, instance_id: str, wanted: str, deadline: float) -> None:
from alibabacloud_ecs20140526 import models as ecs_models
while time.monotonic() < deadline:
response = client.describe_instances(
ecs_models.DescribeInstancesRequest(
region_id=region_id, instance_ids=json.dumps([instance_id])
)
)
instances = response.body.instances.instance
if instances and instances[0].status == wanted:
return
time.sleep(POLL_INTERVAL_SECONDS)
raise TimeoutError(f"Instance {instance_id} did not reach status {wanted} in time")
def fetch_console_output(client, region_id: str, instance_id: str) -> str:
from alibabacloud_ecs20140526 import models as ecs_models
response = client.get_instance_console_output(
ecs_models.GetInstanceConsoleOutputRequest(region_id=region_id, instance_id=instance_id)
)
return base64.b64decode(response.body.console_output or "").decode("utf-8", "replace")
class ConsoleTailer:
"""Incrementally prints new serial-console lines so cloud-init progress is
visible in the CI log while waiting, not only after a failure dump."""
def __init__(self, client, region_id: str, instance_id: str) -> None:
self.client = client
self.region_id = region_id
self.instance_id = instance_id
self.printed_lines = 0
def poll(self) -> None:
try:
lines = fetch_console_output(self.client, self.region_id, self.instance_id).splitlines()
except Exception as error: # noqa: BLE001
print(f"[console] unable to fetch console output: {error}", flush=True)
return
# The API returns a bounded tail; if it ever shrinks, restart from the
# beginning of the new buffer rather than skipping lines.
if len(lines) < self.printed_lines:
self.printed_lines = 0
for line in lines[self.printed_lines :]:
print(f"[console] {line}", flush=True)
self.printed_lines = len(lines)
def dump_console_output(client, region_id: str, instance_id: str) -> None:
try:
output = fetch_console_output(client, region_id, instance_id)
except Exception as error: # noqa: BLE001
print(f"Unable to fetch console output for {instance_id}: {error}", flush=True)
return
tail = "\n".join(output.splitlines()[-80:])
print(f"::group::Console output tail for {instance_id}\n{tail}\n::endgroup::", flush=True)
summary_path = os.environ.get("GITHUB_STEP_SUMMARY")
if summary_path:
with open(summary_path, "a", encoding="utf-8") as summary:
summary.write(f"<details><summary>ECS console output ({instance_id})</summary>\n\n```\n")
summary.write(tail)
summary.write("\n```\n</details>\n")
def append_github_output(name: str, value: str) -> None:
output_path = os.environ.get("GITHUB_OUTPUT")
if output_path:
with open(output_path, "a", encoding="utf-8") as output:
output.write(f"{name}={value}\n")
def run_instance(client, config: ProvisionConfig, user_data: str) -> str:
from alibabacloud_ecs20140526 import models as ecs_models
request = ecs_models.RunInstancesRequest(
region_id=config.region_id,
image_id=config.image_id,
resource_group_id=config.resource_group_id,
instance_type=config.instance_type,
v_switch_id=config.vswitch_id,
security_group_id=config.security_group_id,
instance_name=runner_name_for_run(config.run_id),
host_name=runner_name_for_run(config.run_id),
description=f"Ephemeral query-regression runner for run {config.run_id}",
amount=1,
instance_charge_type="PostPaid",
spot_strategy="NoSpot",
internet_charge_type="PayByTraffic",
internet_max_bandwidth_out=100,
system_disk=ecs_models.RunInstancesRequestSystemDisk(
category="cloud_essd",
size=str(SYSTEM_DISK_GIB),
),
user_data=user_data,
tag=[
ecs_models.RunInstancesRequestTag(key=MANAGED_BY_TAG_KEY, value=MANAGED_BY_TAG_VALUE),
ecs_models.RunInstancesRequestTag(key=RUN_TAG_KEY, value=config.run_id),
],
)
response = client.run_instances(request)
instance_id = response.body.instance_id_sets.instance_id_set[0]
return instance_id
def provision(config: ProvisionConfig) -> int:
client = make_ecs_client(config)
runner_name = runner_name_for_run(config.run_id)
runner_label = runner_label_for_run(config.run_id)
registration_token = create_registration_token(config.github_token, config.repo)
user_data = encode_user_data(
render_user_data(
runner_name,
runner_label,
registration_token,
config.repo,
config.runner_uid,
config.runner_gid,
)
)
print(f"::group::Run instance {runner_name}", flush=True)
instance_id = run_instance(client, config, user_data)
print(f"Instance id: {instance_id}", flush=True)
# Teardown is the only DeleteInstance caller. Write outputs immediately
# so a later status/online failure still releases this VM.
append_github_output("label", runner_label)
append_github_output("instance_id", instance_id)
append_github_output("runner_name", runner_name)
wait_for_instance_status(client, config.region_id, instance_id, "Running", time.monotonic() + 5 * 60)
print("::endgroup::", flush=True)
print("::group::Wait for runner online", flush=True)
deadline = time.monotonic() + RUNNER_ONLINE_TIMEOUT_SECONDS
tailer = ConsoleTailer(client, config.region_id, instance_id)
last_console_poll = 0.0
while time.monotonic() < deadline:
runner = find_runner_by_name(config.github_token, config.repo, runner_name)
if runner is not None:
status = runner.get("status")
print(f"Runner {runner_name} status: {status}", flush=True)
if status == "online":
print("::endgroup::", flush=True)
print(f"Runner {runner_name} is online with label {runner_label}", flush=True)
return 0
# The serial console lags the guest by a minute or so; 30s polling is
# enough to follow cloud-init without tripping API throttling.
if time.monotonic() - last_console_poll >= 30:
tailer.poll()
last_console_poll = time.monotonic()
time.sleep(POLL_INTERVAL_SECONDS)
print("::endgroup::", flush=True)
print(
f"Runner {runner_name} did not come online in time; "
"leaving instance for the teardown job to delete",
flush=True,
)
dump_console_output(client, config.region_id, instance_id)
return 1
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID"))
parser.add_argument("--vswitch-id", default=os.environ.get("ALIYUN_ECS_VSWITCH_ID"))
parser.add_argument("--security-group-id", default=os.environ.get("ALIYUN_ECS_SECURITY_GROUP_ID"))
parser.add_argument("--image-id", default=os.environ.get("QUERY_REGRESSION_ECS_IMAGE_ID"))
parser.add_argument("--instance-type", default=os.environ.get("ALIYUN_ECS_INSTANCE_TYPE"))
parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY"))
parser.add_argument("--run-id", default=os.environ.get("GITHUB_RUN_ID"))
parser.add_argument("--github-token", default=os.environ.get("GH_PERSONAL_ACCESS_TOKEN"))
parser.add_argument("--runner-uid", default=os.environ.get("QUERY_REGRESSION_RUNNER_UID", "1001"))
parser.add_argument("--runner-gid", default=os.environ.get("QUERY_REGRESSION_RUNNER_GID", "1001"))
parser.add_argument("--resource-group-id", default=os.environ.get("ALIYUN_ECS_RESOURCE_GROUP_ID"))
args = parser.parse_args()
missing = [
name
for name, value in vars(args).items()
if name not in ("resource_group_id",)
and (value is None or (isinstance(value, str) and not value))
]
if missing:
flags = ", ".join(f"--{name.replace('_', '-')}" for name in missing)
raise SystemExit(f"Missing required configuration: {flags} (or their env defaults)")
config = ProvisionConfig(
region_id=args.region_id,
vswitch_id=args.vswitch_id,
security_group_id=args.security_group_id,
image_id=args.image_id,
instance_type=args.instance_type,
repo=args.repo,
run_id=args.run_id,
github_token=args.github_token,
resource_group_id=args.resource_group_id or None,
runner_uid=args.runner_uid,
runner_gid=args.runner_gid,
)
return provision(config)
if __name__ == "__main__":
raise SystemExit(main())
.github/scripts/aliyun-ecs-runner-teardown.py
+237
View File
@@ -0,0 +1,237 @@
#!/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.
# PEP 723 inline metadata (see aliyun-ecs-runner-provision.py for the
# convention). The SDK import stays lazy so unit tests run on a plain
# stdlib interpreter.
# /// script
# requires-python = ">=3.10"
# dependencies = [
# "alibabacloud_ecs20140526>=4.1.0,<6",
# "alibabacloud_tea_openapi>=0.3.12,<1",
# ]
# ///
"""Tear down ephemeral Aliyun ECS query-regression runners.
Two modes:
- Targeted (per workflow run): delete one instance by id and deregister its
runner by name. Both steps are idempotent and best-effort; a missing
instance or runner is not an error.
- Sweep (scheduled janitor): delete every instance tagged as managed by
query-regression CI whose creation time is older than the given TTL, and
deregister the matching runners. This is the safety net for runs whose
teardown job never executed.
"""
from __future__ import annotations
import argparse
import importlib.util
import os
import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path
# Reuse the GitHub API client, runner lookup, and tag constants from the
# provision script.
_PROVISION_SPEC = importlib.util.spec_from_file_location(
"aliyun_ecs_runner_provision",
Path(__file__).resolve().parent / "aliyun-ecs-runner-provision.py",
)
assert _PROVISION_SPEC is not None and _PROVISION_SPEC.loader is not None
provision = importlib.util.module_from_spec(_PROVISION_SPEC)
sys.modules[_PROVISION_SPEC.name] = provision
_PROVISION_SPEC.loader.exec_module(provision)
SWEEP_TTL = timedelta(hours=4)
# ECS DescribeInstances creation_time is ISO8601 UTC, e.g. "2026-08-17T02:13Z".
CREATION_TIME_FORMATS = ("%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%MZ")
def parse_creation_time(value: str) -> datetime:
for fmt in CREATION_TIME_FORMATS:
try:
return datetime.strptime(value, fmt).replace(tzinfo=timezone.utc)
except ValueError:
continue
raise ValueError(f"Unparsable ECS creation time: {value}")
def expired_instance_names(
instances: list[tuple[str, str, str]], now: datetime, ttl: timedelta
) -> list[tuple[str, str]]:
"""Pick (instance_id, instance_name) pairs whose creation is older than ttl.
`instances` items are (instance_id, instance_name, creation_time).
"""
expired = []
for instance_id, instance_name, creation_time in instances:
age = now - parse_creation_time(creation_time)
if age >= ttl:
expired.append((instance_id, instance_name))
return expired
def make_ecs_client(region_id: str):
from alibabacloud_ecs20140526.client import Client as EcsClient
from alibabacloud_tea_openapi.models import Config as OpenApiConfig
access_key_id = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_ID", "")
access_key_secret = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_SECRET", "")
if not access_key_id or not access_key_secret:
raise SystemExit(
"ALIBABA_CLOUD_ACCESS_KEY_ID/SECRET are empty; in CI they come from the "
"ALICLOUD_ECS_ACCESS_KEY_ID/SECRET repository secrets (a missing or "
"misnamed secret expands to an empty string)."
)
return EcsClient(
OpenApiConfig(
access_key_id=access_key_id,
access_key_secret=access_key_secret,
region_id=region_id,
endpoint=f"ecs.{region_id}.aliyuncs.com",
)
)
def delete_instance(client, instance_id: str, region_id: str | None = None) -> bool:
import time
from alibabacloud_ecs20140526 import models as ecs_models
# Export the serial console before destroying the evidence: a cancelled
# run still runs this teardown job, and the console tail (cloud-init,
# runner service logs, kernel OOM records) is the only witness of what
# happened on the machine.
if region_id is not None:
provision.dump_console_output(client, region_id, instance_id)
# DeleteInstance rejects instances that are still Initializing (or in
# another transitional status). Teardown runs right after a cancelled
# run, when the instance may be only seconds old, so retry the transient
# status errors for a few minutes instead of leaking the instance.
deadline = time.monotonic() + 5 * 60
while True:
try:
client.delete_instance(ecs_models.DeleteInstanceRequest(instance_id=instance_id, force=True))
print(f"Deleted instance {instance_id}", flush=True)
return True
except Exception as error: # noqa: BLE001
message = str(error)
if "InvalidInstanceId.NotFound" in message:
print(f"Instance {instance_id} already gone", flush=True)
return True
if "IncorrectInstanceStatus" in message and time.monotonic() < deadline:
print(f"Instance {instance_id} is in a transitional status; retrying delete", flush=True)
time.sleep(15)
continue
print(f"Failed to delete instance {instance_id}: {error}", flush=True)
return False
def deregister_runner(token: str, repo: str, runner_name: str) -> bool:
runner = provision.find_runner_by_name(token, repo, runner_name)
if runner is None:
print(f"Runner {runner_name} is not registered", flush=True)
return True
try:
provision.github_api(token, "DELETE", f"/repos/{repo}/actions/runners/{runner['id']}")
print(f"Deregistered runner {runner_name} (id {runner['id']})", flush=True)
return True
except Exception as error: # noqa: BLE001
print(f"Failed to deregister runner {runner_name}: {error}", flush=True)
return False
def list_managed_instances(client, region_id: str) -> list[tuple[str, str, str]]:
from alibabacloud_ecs20140526 import models as ecs_models
result: list[tuple[str, str, str]] = []
next_token = None
while True:
request = ecs_models.DescribeInstancesRequest(
region_id=region_id,
tag=[
ecs_models.DescribeInstancesRequestTag(
key=provision.MANAGED_BY_TAG_KEY, value=provision.MANAGED_BY_TAG_VALUE
)
],
max_results=100,
next_token=next_token,
)
response = client.describe_instances(request)
for instance in response.body.instances.instance:
result.append((instance.instance_id, instance.instance_name, instance.creation_time))
next_token = response.body.next_token
if not next_token:
return result
def sweep(client, region_id: str, repo: str, github_token: str, ttl: timedelta) -> int:
instances = list_managed_instances(client, region_id)
print(f"Found {len(instances)} managed instance(s) in {region_id}", flush=True)
expired = expired_instance_names(instances, datetime.now(timezone.utc), ttl)
ok = True
for instance_id, instance_name in expired:
print(f"Instance {instance_id} ({instance_name}) exceeds TTL {ttl}; deleting", flush=True)
# Runner names mirror instance names by construction in the provision
# script (both are qreg-ecs-<run_id>).
ok &= delete_instance(client, instance_id, region_id)
ok &= deregister_runner(github_token, repo, instance_name)
return 0 if ok else 1
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID"))
parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY"))
parser.add_argument("--github-token", default=os.environ.get("GH_PERSONAL_ACCESS_TOKEN"))
parser.add_argument("--instance-id", default=os.environ.get("QUERY_REGRESSION_ECS_INSTANCE_ID"))
parser.add_argument("--runner-name", default=os.environ.get("QUERY_REGRESSION_ECS_RUNNER_NAME"))
parser.add_argument(
"--sweep",
action="store_true",
help="Janitor mode: delete all managed instances older than --sweep-ttl-hours.",
)
parser.add_argument("--sweep-ttl-hours", type=float, default=SWEEP_TTL.total_seconds() / 3600)
args = parser.parse_args()
for name in ("region_id", "repo", "github_token"):
if not getattr(args, name):
raise SystemExit(f"Missing required configuration: --{name.replace('_', '-')}")
client = make_ecs_client(args.region_id)
if args.sweep:
return sweep(
client, args.region_id, args.repo, args.github_token, timedelta(hours=args.sweep_ttl_hours)
)
ok = True
if args.instance_id:
ok &= delete_instance(client, args.instance_id, args.region_id)
else:
print("No instance id given; skipping instance deletion", flush=True)
if args.runner_name:
ok &= deregister_runner(args.github_token, args.repo, args.runner_name)
else:
print("No runner name given; skipping runner deregistration", flush=True)
return 0 if ok else 1
if __name__ == "__main__":
raise SystemExit(main())
.github/scripts/query-regression-nightly-refs.py
+270
View File
@@ -0,0 +1,270 @@
#!/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.
"""Resolve consecutive successful Nightly Build SHAs for query-regression.
Query regression builds both binaries from git; this script only picks the
head SHAs of two Nightly Build workflow runs (today vs the previous success).
"""
from __future__ import annotations
import argparse
import json
import os
import sys
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass
from typing import Any
NIGHTLY_WORKFLOW = "nightly-build.yml"
NIGHTLY_EVENTS = frozenset({"schedule", "workflow_dispatch"})
@dataclass(frozen=True)
class WorkflowRun:
id: int
head_sha: str
head_branch: str
html_url: str
created_at: str
conclusion: str
event: str
@dataclass(frozen=True)
class RefPair:
base: WorkflowRun | None
candidate: WorkflowRun | None
skip: bool
reason: str
def parse_run(payload: dict[str, Any]) -> WorkflowRun:
return WorkflowRun(
id=int(payload["id"]),
head_sha=str(payload["head_sha"]),
head_branch=str(payload.get("head_branch") or ""),
html_url=str(payload.get("html_url") or ""),
created_at=str(payload.get("created_at") or ""),
conclusion=str(payload.get("conclusion") or ""),
event=str(payload.get("event") or ""),
)
def github_get(token: str, path: str, query: dict[str, str] | None = None) -> dict[str, Any]:
encoded = f"?{urllib.parse.urlencode(query)}" if query else ""
request = urllib.request.Request(
f"https://api.github.com{path}{encoded}",
method="GET",
headers={
"Accept": "application/vnd.github+json",
"Authorization": f"Bearer {token}",
"X-GitHub-Api-Version": "2022-11-28",
},
)
try:
with urllib.request.urlopen(request, timeout=30) as response:
return json.loads(response.read().decode("utf-8"))
except urllib.error.HTTPError as error:
body = error.read().decode("utf-8", "replace")
raise SystemExit(
f"GitHub API GET {path} failed: HTTP {error.code}: {body}"
) from error
def list_successful_nightly_runs(
token: str,
repo: str,
*,
branch: str | None = None,
per_page: int = 30,
) -> list[WorkflowRun]:
query = {
"status": "success",
"per_page": str(per_page),
}
if branch:
query["branch"] = branch
payload = github_get(
token,
f"/repos/{repo}/actions/workflows/{NIGHTLY_WORKFLOW}/runs",
query,
)
runs = [parse_run(item) for item in payload.get("workflow_runs") or []]
return [
run
for run in runs
if run.conclusion == "success" and run.event in NIGHTLY_EVENTS and run.head_sha
]
def fetch_run(token: str, repo: str, run_id: int) -> WorkflowRun:
payload = github_get(token, f"/repos/{repo}/actions/runs/{run_id}")
return parse_run(payload)
def select_base_and_candidate(
runs: list[WorkflowRun],
*,
candidate_run_id: int | None = None,
candidate: WorkflowRun | None = None,
) -> RefPair:
"""Pick candidate (newest / requested) and the previous successful run.
`runs` must be newest-first. Same-SHA consecutive nightlies are skipped
(no commits that day).
"""
if candidate is None:
if candidate_run_id is not None:
candidate = next((run for run in runs if run.id == candidate_run_id), None)
if candidate is None:
return RefPair(
None,
None,
True,
f"candidate nightly run {candidate_run_id} was not a successful "
f"{NIGHTLY_WORKFLOW} run",
)
elif runs:
candidate = runs[0]
else:
return RefPair(None, None, True, "no successful Nightly Build runs")
branch = candidate.head_branch
older = [
run
for run in runs
if run.id != candidate.id
and (not branch or not run.head_branch or run.head_branch == branch)
and (not candidate.created_at or run.created_at <= candidate.created_at)
]
if not older:
return RefPair(
None,
candidate,
True,
"no previous successful Nightly Build to compare against",
)
base = older[0]
if base.head_sha.lower() == candidate.head_sha.lower():
return RefPair(
base,
candidate,
True,
f"previous nightly SHA {base.head_sha} matches candidate; nothing to compare",
)
return RefPair(base, candidate, False, "")
def _override_run(sha: str) -> WorkflowRun:
return WorkflowRun(
id=0,
head_sha=sha,
head_branch="",
html_url="",
created_at="",
conclusion="success",
event="workflow_dispatch",
)
def override_pair(base_ref: str, candidate_ref: str) -> RefPair:
return RefPair(_override_run(base_ref), _override_run(candidate_ref), False, "explicit refs")
def write_outputs(pair: RefPair) -> None:
values = {
"skip": "true" if pair.skip else "false",
"reason": pair.reason,
"base_sha": pair.base.head_sha if pair.base else "",
"candidate_sha": pair.candidate.head_sha if pair.candidate else "",
"base_run_url": pair.base.html_url if pair.base else "",
"candidate_run_url": pair.candidate.html_url if pair.candidate else "",
"base_run_id": str(pair.base.id) if pair.base and pair.base.id else "",
"candidate_run_id": str(pair.candidate.id) if pair.candidate and pair.candidate.id else "",
}
output_path = os.environ.get("GITHUB_OUTPUT")
if output_path:
with open(output_path, "a", encoding="utf-8") as handle:
for key, value in values.items():
handle.write(f"{key}={value}\n")
for key, value in values.items():
print(f"{key}={value}")
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY"))
parser.add_argument(
"--token",
default=os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN"),
)
parser.add_argument("--branch", default=os.environ.get("NIGHTLY_BRANCH") or "")
parser.add_argument(
"--candidate-run-id",
default=os.environ.get("CANDIDATE_RUN_ID") or "",
)
parser.add_argument("--base-ref", default=os.environ.get("BASE_REF") or "")
parser.add_argument("--candidate-ref", default=os.environ.get("CANDIDATE_REF") or "")
return parser.parse_args(argv)
def main(argv: list[str] | None = None) -> int:
args = parse_args(argv)
base_ref = args.base_ref.strip()
candidate_ref = args.candidate_ref.strip()
if base_ref and candidate_ref:
write_outputs(override_pair(base_ref, candidate_ref))
return 0
if bool(base_ref) != bool(candidate_ref):
print("Provide both --base-ref and --candidate-ref, or neither.", file=sys.stderr)
return 2
if not args.repo or not args.token:
print("--repo and --token (or GITHUB_REPOSITORY / GITHUB_TOKEN) are required.", file=sys.stderr)
return 2
candidate_run_id = int(args.candidate_run_id) if str(args.candidate_run_id).strip() else None
candidate: WorkflowRun | None = None
branch = args.branch.strip() or None
if candidate_run_id is not None:
candidate = fetch_run(args.token, args.repo, candidate_run_id)
if candidate.conclusion != "success":
write_outputs(
RefPair(
None,
candidate,
True,
f"nightly run {candidate_run_id} conclusion is {candidate.conclusion}",
)
)
return 0
branch = branch or candidate.head_branch or None
runs = list_successful_nightly_runs(args.token, args.repo, branch=branch)
pair = select_base_and_candidate(
runs,
candidate_run_id=candidate_run_id if candidate is None else None,
candidate=candidate,
)
write_outputs(pair)
return 0
if __name__ == "__main__":
raise SystemExit(main())
.github/workflows/checks.yml
+8
View File
@@ -80,6 +80,14 @@ jobs:
run: python3 scripts/check-enterprise-license-test.py
- name: Run fuzz orchestration script tests
run: .github/scripts/run-fuzz-targets-test.sh
- name: Run query-regression Python tooling tests
run: |
python3 tests/perf/test_query_regression_runner_compaction_toctou.py
python3 tests/perf/test_query_regression_runner_otlp_trace_load.py
python3 tests/perf/test_query_regression_summary_otlp.py
python3 tests/perf/test_query_regression_case_selection.py
python3 tests/perf/test_query_regression_nightly_refs.py
python3 tests/perf/test_aliyun_ecs_runner_scripts.py
check:
if: ${{ github.repository == 'GreptimeTeam/greptimedb' }}
.github/workflows/query-regression-janitor.yml
+36
View File
@@ -0,0 +1,36 @@
name: Query Regression ECS Janitor
# Safety net for the aliyun-ecs query-regression path: deletes managed ECS
# instances (and their runner registrations) that outlived their workflow run,
# so a failed teardown never leaks pay-as-you-go spend.
on:
schedule:
- cron: "23 3 * * *"
workflow_dispatch:
permissions:
contents: read
jobs:
sweep:
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- name: Checkout teardown script
uses: actions/checkout@v4
with:
persist-credentials: false
- name: Install uv
uses: astral-sh/setup-uv@v6
- name: Sweep expired ECS runners
env:
ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }}
ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }}
GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }}
run: >-
uv run .github/scripts/aliyun-ecs-runner-teardown.py
--sweep --sweep-ttl-hours 4
.github/workflows/query-regression-nightly.yml
+104
View File
@@ -0,0 +1,104 @@
name: Query Regression Nightly
# After a successful GreptimeDB Nightly Build, compare that commit against
# the previous successful nightly. The reusable Query Regression workflow
# still compiles both SHAs; this wrapper only resolves which two SHAs.
on:
workflow_run:
workflows:
- GreptimeDB Nightly Build
types:
- completed
workflow_dispatch:
inputs:
base_ref:
description: Base ref/SHA (empty = previous successful nightly)
required: false
type: string
default: ""
candidate_ref:
description: Candidate ref/SHA (empty = latest successful nightly)
required: false
type: string
default: ""
candidate_run_id:
description: Nightly Build run id to treat as candidate (empty = latest)
required: false
type: string
default: ""
case:
description: Query perf case path(s), all, or heavy
required: false
type: string
default: all
permissions:
contents: read
actions: read
jobs:
resolve-refs:
name: Resolve previous vs current nightly SHAs
if: ${{ github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success' }}
runs-on: ubuntu-latest
timeout-minutes: 5
outputs:
skip: ${{ steps.refs.outputs.skip }}
reason: ${{ steps.refs.outputs.reason }}
base_sha: ${{ steps.refs.outputs.base_sha }}
candidate_sha: ${{ steps.refs.outputs.candidate_sha }}
steps:
- name: Checkout ref resolver
uses: actions/checkout@v4
with:
persist-credentials: false
- name: Resolve nightly SHAs
id: refs
env:
GITHUB_TOKEN: ${{ github.token }}
GITHUB_REPOSITORY: ${{ github.repository }}
CANDIDATE_RUN_ID: ${{ github.event_name == 'workflow_run' && github.event.workflow_run.id || inputs.candidate_run_id || '' }}
BASE_REF: ${{ github.event_name == 'workflow_dispatch' && inputs.base_ref || '' }}
CANDIDATE_REF: ${{ github.event_name == 'workflow_dispatch' && inputs.candidate_ref || '' }}
run: python3 .github/scripts/query-regression-nightly-refs.py
- name: Summarize comparison
if: always()
env:
SKIP: ${{ steps.refs.outputs.skip }}
REASON: ${{ steps.refs.outputs.reason }}
BASE_SHA: ${{ steps.refs.outputs.base_sha }}
CANDIDATE_SHA: ${{ steps.refs.outputs.candidate_sha }}
BASE_RUN_URL: ${{ steps.refs.outputs.base_run_url }}
CANDIDATE_RUN_URL: ${{ steps.refs.outputs.candidate_run_url }}
run: |
set -euo pipefail
{
if [[ "${SKIP}" == "true" ]]; then
printf 'Skipping query-regression nightly: %s\n' "${REASON}"
else
printf 'Comparing previous nightly `%s` -> current nightly `%s`\n' \
"${BASE_SHA}" "${CANDIDATE_SHA}"
if [[ -n "${BASE_RUN_URL}" ]]; then
printf -- '- Previous Nightly Build: %s\n' "${BASE_RUN_URL}"
fi
if [[ -n "${CANDIDATE_RUN_URL}" ]]; then
printf -- '- Current Nightly Build: %s\n' "${CANDIDATE_RUN_URL}"
fi
fi
} | tee -a "${GITHUB_STEP_SUMMARY}"
query-regression:
name: Query regression nightly
needs: [resolve-refs]
if: ${{ needs.resolve-refs.outputs.skip != 'true' && needs.resolve-refs.outputs.base_sha != '' && needs.resolve-refs.outputs.candidate_sha != '' }}
uses: ./.github/workflows/query-regression.yml
secrets: inherit
with:
case: ${{ github.event_name == 'workflow_dispatch' && inputs.case || 'all' }}
base_ref: ${{ needs.resolve-refs.outputs.base_sha }}
candidate_ref: ${{ needs.resolve-refs.outputs.candidate_sha }}
cargo_profile: nightly
runner: aliyun-ecs
.github/workflows/query-regression.yml
+306 -127
View File
@@ -32,10 +32,10 @@ on:
type: string
default: nightly
runner:
description: Self-hosted runner label or ARC runner scale set for this query regression run
description: Self-hosted runner label; aliyun-ecs provisions a fresh ECS instance per run
required: false
type: string
default: perf-regression-8-cores
default: aliyun-ecs
workflow_dispatch:
inputs:
case:
@@ -68,12 +68,20 @@ on:
- release
- dev
runner:
description: Self-hosted runner label or ARC runner scale set for this query regression run
description: >-
Self-hosted runner label; aliyun-ecs provisions a fresh ECS instance
per run, any other value is used as a literal runner label
required: true
type: choice
default: perf-regression-8-cores
options:
- perf-regression-8-cores
type: string
default: aliyun-ecs
keep_instance:
description: >-
Debug: keep the ECS instance after the run (skip teardown) so its
runner _diag logs, journal, and telemetry can be inspected; the
janitor still sweeps it after the TTL
required: false
type: boolean
default: false
pull_request:
types: [labeled]
@@ -81,19 +89,99 @@ permissions:
contents: read
jobs:
query-regression:
test-tooling:
# Stdlib unittests for the workflow Python (case selection, report
# helpers, nightly SHA picking, and rendered ECS user-data). They do
# not talk to Aliyun or the Actions runner process; running them on
# ubuntu-latest fails fast before any ECS spend. Not gated on the
# regression labels: those labels boot a VM, these tests should not.
# Ordinary PRs also run the same tests from checks.yml, because this
# workflow only starts on `labeled` (or dispatch / workflow_call).
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- name: Checkout
uses: actions/checkout@v4
with:
persist-credentials: false
- name: Test query regression tooling
run: |
python3 tests/perf/test_query_regression_runner_compaction_toctou.py
python3 tests/perf/test_query_regression_runner_otlp_trace_load.py
python3 tests/perf/test_query_regression_summary_otlp.py
python3 tests/perf/test_query_regression_case_selection.py
python3 tests/perf/test_query_regression_nightly_refs.py
python3 tests/perf/test_aliyun_ecs_runner_scripts.py
provision:
# Runs when the aliyun-ecs path is selected: explicitly via the runner
# input, or for PR labels by default (a QUERY_REGRESSION_PR_RUNNER
# repository variable set to another value redirects PRs to that literal
# runner label instead). Uses trusted scripts from the PR base (or the
# dispatched ref), never from candidate code. Waits for test-tooling so
# a broken user-data template does not still create a VM.
needs: [test-tooling]
if: >-
${{ github.event_name != 'pull_request' ||
(github.event_name == 'pull_request' &&
!github.event.pull_request.draft &&
(github.event.label.name == 'query-regression' ||
github.event.label.name == 'heavy-regression')) }}
runs-on: ${{ github.event_name != 'pull_request' && inputs.runner || 'perf-regression-8-cores' }}
${{ !failure() && !cancelled() &&
((github.event_name != 'pull_request' && inputs.runner == 'aliyun-ecs') ||
(github.event_name == 'pull_request' &&
!github.event.pull_request.draft &&
(github.event.label.name == 'query-regression' ||
github.event.label.name == 'heavy-regression') &&
(vars.QUERY_REGRESSION_PR_RUNNER || 'aliyun-ecs') == 'aliyun-ecs')) }}
runs-on: ubuntu-latest
timeout-minutes: 45
outputs:
label: ${{ steps.provision.outputs.label }}
instance_id: ${{ steps.provision.outputs.instance_id }}
runner_name: ${{ steps.provision.outputs.runner_name }}
steps:
- name: Checkout trusted provisioning scripts
uses: actions/checkout@v4
with:
repository: ${{ github.repository }}
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || github.sha }}
persist-credentials: false
- name: Install uv
uses: astral-sh/setup-uv@v6
- name: Provision ECS runner
id: provision
env:
ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }}
ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }}
GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }}
ALIYUN_ECS_VSWITCH_ID: ${{ vars.ALIYUN_ECS_VSWITCH_ID }}
ALIYUN_ECS_SECURITY_GROUP_ID: ${{ vars.ALIYUN_ECS_SECURITY_GROUP_ID }}
ALIYUN_ECS_INSTANCE_TYPE: ${{ vars.ALIYUN_ECS_INSTANCE_TYPE }}
ALIYUN_ECS_RESOURCE_GROUP_ID: ${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }}
QUERY_REGRESSION_ECS_IMAGE_ID: ${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }}
QUERY_REGRESSION_RUNNER_UID: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }}
QUERY_REGRESSION_RUNNER_GID: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }}
run: >-
uv run .github/scripts/aliyun-ecs-runner-provision.py
query-regression:
needs: [provision, test-tooling]
# `!failure() && !cancelled()` both suppresses the implicit success() and
# lets the job run when provision was intentionally skipped because a
# literal runner label was selected; a failed or cancelled provision still
# blocks the run because no ECS runner would be waiting.
if: >-
${{ !failure() && !cancelled() &&
(github.event_name != 'pull_request' ||
(github.event_name == 'pull_request' &&
!github.event.pull_request.draft &&
(github.event.label.name == 'query-regression' ||
github.event.label.name == 'heavy-regression'))) }}
runs-on: >-
${{ needs.provision.outputs.label ||
(github.event_name != 'pull_request' && inputs.runner ||
(vars.QUERY_REGRESSION_PR_RUNNER || 'aliyun-ecs')) }}
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_HOME: /home/runner/.cargo
@@ -105,7 +193,8 @@ jobs:
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
# Caps the local sccache on the system disk next to target and cargo.
SCCACHE_CACHE_SIZE: 10G
CARGO_INCREMENTAL: "0"
RUSTFLAGS: -D warnings -C link-arg=-fuse-ld=mold
QUERY_REGRESSION_CACHE_EPOCH: "1"
@@ -114,7 +203,55 @@ jobs:
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 || '' }}
# Runner identity contract. The ECS image uses 1001; a manually prepared
# host may override via repo variables when 1001 is already taken.
EXPECTED_RUNNER_UID: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }}
EXPECTED_RUNNER_GID: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }}
steps:
- name: Report provisioned ECS runner
if: ${{ needs.provision.outputs.instance_id != '' }}
shell: bash
run: |
set -euo pipefail
{
printf -- '- ECS instance: `%s` (type `%s`, image `%s`, region `%s`)\n' \
"${{ needs.provision.outputs.instance_id }}" \
"${{ vars.ALIYUN_ECS_INSTANCE_TYPE }}" \
"${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }}" \
"${{ vars.ALIYUN_ECS_REGION_ID }}"
} >> "${GITHUB_STEP_SUMMARY}"
- name: Start machine telemetry sampler
shell: bash
# The ECS instance is deleted when the run ends, so machine state must
# be captured while the job runs. A background sampler appends load,
# memory, disk, and top-memory processes to a log every 30s; the
# "Dump machine telemetry" step (if: always()) prints and uploads it.
# Note: when the job is *cancelled*, even always() steps are killed —
# telemetry covers failures, not cancellations.
run: |
set -euo pipefail
log="${GITHUB_WORKSPACE}/machine-telemetry.log"
{
echo "== baseline $(date -u +%Y-%m-%dT%H:%M:%SZ) =="
nproc
free -h
df -h / /home/runner
} >> "${log}"
nohup bash -c "
while true; do
{
date -u '+== %Y-%m-%dT%H:%M:%SZ =='
uptime
free -m
df -h / /home/runner
ps -eo pid,comm,%mem,%cpu --sort=-%mem | head -6
} >> '${log}' 2>&1
sleep 30
done
" >/dev/null 2>&1 &
echo "Telemetry sampler started (pid $!), logging to ${log}"
- name: Checkout base source
uses: actions/checkout@v4
with:
@@ -263,14 +400,64 @@ jobs:
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" ]]
errors=()
record() {
printf 'FAIL %s\n' "$*"
errors+=("$*")
}
require_eq() {
local name="$1" actual="$2" expected="$3"
printf 'check %s: %s\n' "${name}" "${actual}"
[[ "${actual}" == "${expected}" ]] || record "${name}: expected '${expected}', got '${actual}'"
}
require_match() {
local name="$1" actual="$2" pattern="$3"
printf 'check %s: %s\n' "${name}" "${actual}"
[[ "${actual}" =~ $pattern ]] || record "${name}: expected to match ${pattern}, got '${actual}'"
}
require() {
local name="$1"
shift
printf 'check %s\n' "${name}"
"$@" || record "${name}"
}
capture() {
local out
if out="$("$@" 2>&1)"; then
printf '%s' "${out}"
else
printf '<failed to run: %s>' "$*"
fi
}
printf 'uid=%s gid=%s PATH=%s CARGO_HOME=%s\n' "$(id -u)" "$(id -g)" "${PATH}" "${CARGO_HOME}"
ls -la "${CARGO_HOME}" 2>&1 || printf '(CARGO_HOME missing)\n'
# Image/host hygiene: CARGO_HOME must be empty of config and
# install-state before cargo/rustup run. Those tools create
# .package-cache / bin on first use, so this cannot come after.
cargo_home_dirty=()
for entry in config config.toml credentials credentials.toml bin .crates.toml .crates2.json .global-cache .package-cache; do
if [[ -e "${CARGO_HOME}/${entry}" ]]; then
cargo_home_dirty+=("${entry}")
record "CARGO_HOME must not contain ${entry}"
fi
done
if (( ${#cargo_home_dirty[@]} > 0 )); then
printf 'CARGO_HOME contents:\n' >&2
ls -la "${CARGO_HOME}" >&2 || true
fi
require_eq uid "$(id -u)" "${EXPECTED_RUNNER_UID}"
require_eq gid "$(id -g)" "${EXPECTED_RUNNER_GID}"
require_eq UV_CACHE_DIR "${UV_CACHE_DIR}" "/home/runner/.cargo/uv-cache"
require_eq protoc "$(capture 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
require "readable /usr/include/google/protobuf/any.proto" \
test -r /usr/include/google/protobuf/any.proto
require "readable /usr/include/google/protobuf/empty.proto" \
test -r /usr/include/google/protobuf/empty.proto
printf '%s\n' \
'syntax = "proto3";' \
'package smoke;' \
@@ -278,47 +465,51 @@ jobs:
'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 \
if 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" ]]
"${temporary_proto_dir}/smoke.proto"; then
require "protoc smoke descriptor is non-empty" test -s "${temporary_proto_dir}/smoke.pb"
else
record "protoc smoke compile failed"
fi
require_match uv "$(capture uv --version)" '^uv[[:space:]]0\.11\.26([[:space:]]|$)'
require_match mold "$(capture mold --version)" '^mold[[:space:]]2\.40\.4([[:space:]]|$)'
require_eq python3 "$(capture python3 --version)" "Python 3.14.4"
otelgen_path="$(command -v otelgen || true)"
otelgen_version="<unavailable>"
if [[ -n "${otelgen_path}" ]]; then
otelgen_version="$("${otelgen_path}" --version 2>&1 || true)"
otelgen_version="$(capture "${otelgen_path}" --version)"
fi
printf 'otelgen path: %s\n' "${otelgen_path:-<missing>}"
printf 'otelgen version: %s\n' "${otelgen_version}"
[[ "${otelgen_path}" == "/usr/local/bin/otelgen" ]]
[[ "${otelgen_version}" == *"863a3f395d062c7322cc1de08a38774b7fdaa6c8"* ]]
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
require_eq otelgen_path "${otelgen_path}" "/usr/local/bin/otelgen"
printf 'check otelgen_version: %s\n' "${otelgen_version}"
[[ "${otelgen_version}" == *"863a3f395d062c7322cc1de08a38774b7fdaa6c8"* ]] \
|| record "otelgen_version: expected commit 863a3f395d062c7322cc1de08a38774b7fdaa6c8, got '${otelgen_version}'"
require_match sccache "$(capture sccache --version)" '^sccache[[:space:]]0\.16\.0([[:space:]]|$)'
require_eq rustup_path "$(command -v rustup || true)" "/opt/cargo/bin/rustup"
require_eq cargo_path "$(command -v cargo || true)" "/opt/cargo/bin/cargo"
require_eq rustc_path "$(command -v rustc || true)" "/opt/cargo/bin/rustc"
require_match rustup "$(capture rustup --version)" '^rustup[[:space:]]1\.29\.0([[:space:]]|$)'
require_match cargo "$(capture cargo --version)" \
'^cargo[[:space:]]1\.96\.0-nightly[[:space:]]\(cbb9bb8bd[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$'
require_match rustc "$(capture rustc --version)" \
'^rustc[[:space:]]1\.96\.0-nightly[[:space:]]\(ac7f9ec7d[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$'
require_match active_toolchain "$(capture rustup show active-toolchain)" \
'^nightly-2026-03-21-x86_64-unknown-linux-gnu([[:space:]]|$)'
require_eq RUSTUP_HOME "${RUSTUP_HOME}" "/opt/rustup"
require_eq RUSTUP_TOOLCHAIN "${RUSTUP_TOOLCHAIN}" "nightly-2026-03-21"
require_eq RUSTUP_AUTO_INSTALL "${RUSTUP_AUTO_INSTALL}" "0"
require "readable /opt/rustup" test -r /opt/rustup
require "executable /opt/rustup" test -x /opt/rustup
require "runner cannot write /opt/rustup" test ! -w /opt/rustup
require "runner cannot write /opt/cargo/bin" test ! -w /opt/cargo/bin
mkdir -p "${CARGO_HOME}/registry" "${CARGO_HOME}/git"
if (( ${#errors[@]} > 0 )); then
printf '\nVerify runner image tools failed (%d checks):\n' "${#errors[@]}" >&2
printf ' - %s\n' "${errors[@]}" >&2
exit 1
fi
- name: Prepare persistent query regression cache
- name: Prepare query regression cache
shell: bash
working-directory: src
run: |
@@ -362,24 +553,6 @@ jobs:
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}"
@@ -400,7 +573,7 @@ jobs:
printf 'Refusing unexpected RUSTC_WRAPPER: %s\n' "${RUSTC_WRAPPER}" >&2
exit 1
}
[[ "${SCCACHE_CACHE_SIZE}" == "40G" ]] || {
[[ "${SCCACHE_CACHE_SIZE}" == "10G" ]] || {
printf 'Refusing unexpected SCCACHE_CACHE_SIZE: %s\n' "${SCCACHE_CACHE_SIZE}" >&2
exit 1
}
@@ -487,46 +660,10 @@ jobs:
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
# Every run starts from a fresh system disk, so size/free-space
# watermarks from the retired retained-disk era are not restored.
# A run that overflows the disk fails the build outright, which the
# telemetry step makes diagnosable.
sccache --start-server
sccache --zero-stats
@@ -567,14 +704,6 @@ jobs:
git reset --hard "${VERIFIED_CANDIDATE_SHA}"
git clean -ffdx
- name: Test query regression tooling
working-directory: src
run: |
uv run --no-project python tests/perf/test_query_regression_runner_compaction_toctou.py
uv run --no-project python tests/perf/test_query_regression_runner_otlp_trace_load.py
uv run --no-project python tests/perf/test_query_regression_summary_otlp.py
uv run --no-project python tests/perf/test_query_regression_case_selection.py
- name: Build candidate greptime and query regression helpers
working-directory: src
run: |
@@ -643,6 +772,7 @@ jobs:
query-regression-work/**/logs/**
query-regression-work/**/otelgen/**
query-regression-summary.md
machine-telemetry.log
if-no-files-found: warn
retention-days: 7
@@ -657,7 +787,7 @@ jobs:
if-no-files-found: warn
retention-days: 7
- name: Report persistent cache usage
- name: Report cache usage
if: ${{ always() }}
shell: bash
run: |
@@ -697,6 +827,55 @@ jobs:
printf 'sccache is unavailable (report only)\n' >&2
fi
- name: Dump machine telemetry
if: ${{ always() }}
shell: bash
run: |
log="${GITHUB_WORKSPACE}/machine-telemetry.log"
if [[ -f "${log}" ]]; then
echo "::group::Machine telemetry (last 200 lines)"
tail -n 200 "${log}"
echo "::endgroup::"
else
echo "No telemetry log found (sampler never started?)"
fi
echo "::group::dmesg tail (OOM killer records)"
sudo dmesg -T 2>/dev/null | tail -n 50 || dmesg -T 2>/dev/null | tail -n 50 || \
echo "dmesg unavailable without root"
echo "::endgroup::"
- name: Fail on regression failure
if: ${{ steps.run.outputs.status != '0' }}
run: exit 1
teardown:
# Releases the dynamically provisioned ECS runner. Runs even when the
# benchmark job fails or is cancelled; skipped when a literal runner label
# was selected because the provision outputs are empty, or when the
# dispatch set keep_instance to preserve the machine for post-mortem
# debugging (the janitor sweep still reclaims it after the TTL).
if: ${{ always() && needs.provision.outputs.instance_id != '' && !inputs.keep_instance }}
needs: [provision, query-regression]
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- name: Checkout trusted teardown scripts
uses: actions/checkout@v4
with:
repository: ${{ github.repository }}
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || github.sha }}
persist-credentials: false
- name: Install uv
uses: astral-sh/setup-uv@v6
- name: Teardown ECS runner
env:
ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }}
ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }}
GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }}
QUERY_REGRESSION_ECS_INSTANCE_ID: ${{ needs.provision.outputs.instance_id }}
QUERY_REGRESSION_ECS_RUNNER_NAME: ${{ needs.provision.outputs.runner_name }}
run: >-
uv run .github/scripts/aliyun-ecs-runner-teardown.py
.github/workflows/release.yml
+1 -1
View File
@@ -510,7 +510,7 @@ jobs:
base_ref: ${{ needs.prepare-release-validation.outputs.previous-release-tag }}
candidate_ref: ${{ needs.prepare-release-validation.outputs.candidate-ref }}
cargo_profile: nightly
runner: perf-regression-8-cores
runner: aliyun-ecs
release-images-to-dockerhub:
name: Build and push images to DockerHub
Cargo.lock
Generated
+144 -46
View File
@@ -2167,6 +2167,7 @@ dependencies = [
"common-time",
"common-version",
"common-wal",
"daemonize",
"datafusion",
"datafusion-common",
"datafusion-physical-plan",
@@ -3570,6 +3571,15 @@ dependencies = [
"syn 2.0.117",
]
[[package]]
name = "daemonize"
version = "0.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ab8bfdaacb3c887a54d41bdf48d3af8873b3f5566469f8ba21b92057509f116e"
dependencies = [
"libc",
]
[[package]]
name = "darling"
version = "0.14.4"
@@ -4628,10 +4638,23 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb"
dependencies = [
"const-oid 0.9.6",
"der_derive",
"flagset",
"pem-rfc7468",
"zeroize",
]
[[package]]
name = "der_derive"
version = "0.7.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8034092389675178f570469e6c3b0465d3d30b4505c294a6550db47f3c17ad18"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "deranged"
version = "0.5.8"
@@ -5334,6 +5357,12 @@ version = "0.5.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99"
[[package]]
name = "flagset"
version = "0.4.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b7ac824320a75a52197e8f2d787f6a38b6718bb6897a35142d749af3c0e8f4fe"
[[package]]
name = "flatbuffers"
version = "25.2.10"
@@ -5973,7 +6002,7 @@ dependencies = [
"cfg-if",
"js-sys",
"libc",
"wasi",
"wasi 0.11.1+wasi-snapshot-preview1",
"wasm-bindgen",
]
@@ -7866,13 +7895,14 @@ checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de"
[[package]]
name = "libredox"
version = "0.1.4"
version = "0.1.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1580801010e535496706ba011c15f8532df6b42297d2e471fec38ceadd8c0638"
checksum = "28d0a00925a9f930d679b6789b721e3a7f9ed110f41b86d2497caa780c3a070a"
dependencies = [
"bitflags 2.12.1",
"libc",
"redox_syscall 0.5.13",
"plain",
"redox_syscall 0.9.3",
]
[[package]]
@@ -8543,7 +8573,7 @@ checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda"
dependencies = [
"libc",
"log",
"wasi",
"wasi 0.11.1+wasi-snapshot-preview1",
"windows-sys 0.61.2",
]
@@ -9358,6 +9388,15 @@ dependencies = [
"objc2-encode",
]
[[package]]
name = "objc2-core-foundation"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536"
dependencies = [
"bitflags 2.12.1",
]
[[package]]
name = "objc2-encode"
version = "4.1.0"
@@ -9374,6 +9413,15 @@ dependencies = [
"objc2",
]
[[package]]
name = "objc2-system-configuration"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396"
dependencies = [
"objc2-core-foundation",
]
[[package]]
name = "object"
version = "0.36.7"
@@ -10552,16 +10600,7 @@ dependencies = [
"tokio",
"tokio-rustls",
"tokio-util",
"x509-certificate 0.25.0",
]
[[package]]
name = "phf"
version = "0.11.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078"
dependencies = [
"phf_shared 0.11.3",
"x509-certificate",
]
[[package]]
@@ -10814,6 +10853,12 @@ version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c"
[[package]]
name = "plain"
version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6"
[[package]]
name = "plist"
version = "1.7.2"
@@ -10924,16 +10969,16 @@ dependencies = [
[[package]]
name = "postgres-types"
version = "0.2.9"
version = "0.2.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "613283563cd90e1dfc3518d548caee47e0e725455ed619881f5cf21f36de4b48"
checksum = "851ca9db4932932d69f3ea811b1abe63087a0f740a47692619dd40d4899b68be"
dependencies = [
"array-init",
"bytes",
"chrono",
"fallible-iterator",
"postgres-protocol",
"serde",
"serde_core",
"serde_json",
]
@@ -11299,7 +11344,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4"
dependencies = [
"bytes",
"heck 0.4.1",
"heck 0.5.0",
"itertools 0.12.1",
"log",
"multimap",
@@ -11319,7 +11364,7 @@ version = "0.14.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1"
dependencies = [
"heck 0.4.1",
"heck 0.5.0",
"itertools 0.14.0",
"log",
"multimap",
@@ -12087,6 +12132,15 @@ dependencies = [
"bitflags 2.12.1",
]
[[package]]
name = "redox_syscall"
version = "0.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d678d17679829e73d371e96880897e98fee2ded7acc0a50bdf8af2affa4b2fe5"
dependencies = [
"bitflags 2.12.1",
]
[[package]]
name = "redox_users"
version = "0.4.6"
@@ -13785,7 +13839,7 @@ version = "0.8.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7"
dependencies = [
"heck 0.4.1",
"heck 0.5.0",
"proc-macro2",
"quote",
"syn 2.0.117",
@@ -14117,7 +14171,7 @@ dependencies = [
"stringprep",
"thiserror 2.0.17",
"tracing",
"whoami",
"whoami 1.6.0",
]
[[package]]
@@ -14156,7 +14210,7 @@ dependencies = [
"stringprep",
"thiserror 2.0.17",
"tracing",
"whoami",
"whoami 1.6.0",
]
[[package]]
@@ -15175,6 +15229,27 @@ version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
[[package]]
name = "tls_codec"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0de2e01245e2bb89d6f05801c564fa27624dbd7b1846859876c7dad82e90bf6b"
dependencies = [
"tls_codec_derive",
"zeroize",
]
[[package]]
name = "tls_codec_derive"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2d2e76690929402faae40aebdda620a2c0e25dd6d3b9afe48867dfd95991f4bd"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "tokio"
version = "1.52.3"
@@ -15241,9 +15316,9 @@ dependencies = [
[[package]]
name = "tokio-postgres"
version = "0.7.13"
version = "0.7.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6c95d533c83082bb6490e0189acaa0bbeef9084e60471b696ca6988cd0541fb0"
checksum = "a528f7d280f6d5b9cd149635c8705b0dd049754bc67d81d31fa25169a93809d3"
dependencies = [
"async-trait",
"byteorder",
@@ -15254,29 +15329,29 @@ dependencies = [
"log",
"parking_lot 0.12.4",
"percent-encoding",
"phf 0.11.3",
"phf 0.13.1",
"pin-project-lite",
"postgres-protocol",
"postgres-types",
"rand 0.9.4",
"socket2 0.5.10",
"rand 0.10.1",
"socket2 0.6.4",
"tokio",
"tokio-util",
"whoami",
"whoami 2.1.3",
]
[[package]]
name = "tokio-postgres-rustls"
version = "0.12.0"
version = "0.14.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "04fb792ccd6bbcd4bba408eb8a292f70fc4a3589e5d793626f45190e6454b6ab"
checksum = "4c2ad44aa0ae96db89c4742212ed41645b2f597311ff6e1945542a4d9fadc2fb"
dependencies = [
"ring",
"rustls",
"sha2 0.11.0",
"tokio",
"tokio-postgres",
"tokio-rustls",
"x509-certificate 0.23.1",
"x509-cert",
]
[[package]]
@@ -16284,6 +16359,15 @@ version = "0.11.1+wasi-snapshot-preview1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
[[package]]
name = "wasi"
version = "0.14.7+wasi-0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c"
dependencies = [
"wasip2",
]
[[package]]
name = "wasip2"
version = "1.0.2+wasi-0.2.9"
@@ -16308,6 +16392,15 @@ version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b8dad83b4f25e74f184f64c43b150b91efe7647395b42289f38e50566d82855b"
[[package]]
name = "wasite"
version = "1.0.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "66fe902b4a6b8028a753d5424909b764ccf79b7a209eac9bf97e59cda9f71a42"
dependencies = [
"wasi 0.14.7+wasi-0.2.4",
]
[[package]]
name = "wasm-bindgen"
version = "0.2.122"
@@ -16530,7 +16623,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6994d13118ab492c3c80c1f81928718159254c53c472bf9ce36f8dae4add02a7"
dependencies = [
"redox_syscall 0.5.13",
"wasite",
"wasite 0.1.0",
]
[[package]]
name = "whoami"
version = "2.1.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "626c4bac6755d76ffc12cb01b2eac751db1996b9e0041de9aa02c8c211ddc82c"
dependencies = [
"libc",
"libredox",
"objc2-system-configuration",
"wasite 1.0.2",
"web-sys",
]
@@ -17139,22 +17244,15 @@ dependencies = [
]
[[package]]
name = "x509-certificate"
version = "0.23.1"
name = "x509-cert"
version = "0.2.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "66534846dec7a11d7c50a74b7cdb208b9a581cad890b7866430d438455847c85"
checksum = "1301e935010a701ae5f8655edc0ad17c44bad3ac5ce8c39185f75453b720ae94"
dependencies = [
"bcder",
"bytes",
"chrono",
"const-oid 0.9.6",
"der",
"hex",
"pem",
"ring",
"signature",
"spki",
"thiserror 1.0.69",
"zeroize",
"tls_codec",
]
[[package]]
Cargo.toml
+1 -1
View File
@@ -255,7 +255,7 @@ strum = { version = "0.27", features = ["derive"] }
sysinfo = "0.33"
tempfile = "3"
tokio = { version = "1.47", features = ["full"] }
tokio-postgres = "0.7"
tokio-postgres = "0.7.18"
tokio-rustls = { version = "0.26.2", default-features = false }
tokio-stream = "0.1"
tokio-util = { version = "0.7", features = ["io-util", "compat"] }
config/config.md
+6 -4
View File
@@ -76,6 +76,7 @@
| `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. |
| `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. |
| `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. |
| `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric<br/>resource attributes, so metrics-only services reach the semantic graph. |
| `prom_store` | -- | -- | Prometheus remote storage options |
| `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. |
| `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. |
@@ -160,7 +161,7 @@
| `region_engine.mito.max_background_purges` | Integer | Auto | Max number of running background purge jobs (default: number of cpu cores). |
| `region_engine.mito.experimental_compaction_memory_limit` | String | 0 | Memory budget for compaction tasks.<br/>Supports absolute size (e.g., "2GiB", "512MB") or percentage of system memory (e.g., "50%").<br/>Setting it to 0 or "unlimited" disables the limit. |
| `region_engine.mito.experimental_compaction_on_exhausted` | String | wait | Behavior when compaction cannot acquire memory from the budget.<br/>Options: "wait" (default, 10s), "wait(<duration>)", "fail" |
| `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. |
| `region_engine.mito.auto_flush_interval` | String | `10m` | 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.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`. |
@@ -236,7 +237,7 @@
| `tracing.tokio_console_addr` | String | Unset | The tokio console address. |
| `event_recorder` | -- | -- | Configuration options for the event recorder. |
| `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. |
| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,<br/>`alter_database`, `drop_database`, `create_flow`, `drop_flow`,<br/>`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,<br/>`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,<br/>`create_view`, `drop_view`, `admin_function`.<br/>When omitted, all current and future event types are recorded.<br/>Set to an empty array to disable event recording. |
| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,<br/>`alter_database`, `drop_database`, `create_flow`, `drop_flow`,<br/>`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,<br/>`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,<br/>`create_view`, `drop_view`, `admin_function`, `reconcile_table`.<br/>When omitted, all current and future event types are recorded.<br/>Set to an empty array to disable event recording. |
| `memory` | -- | -- | The memory options. |
| `memory.enable_heap_profiling` | Bool | `true` | Whether to enable heap profiling activation during startup.<br/>When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable<br/>is set to "prof:true,prof_active:false". The official image adds this env variable.<br/>Default is true. |
@@ -319,6 +320,7 @@
| `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. |
| `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. |
| `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. |
| `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric<br/>resource attributes, so metrics-only services reach the semantic graph. |
| `prom_store` | -- | -- | Prometheus remote storage options |
| `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. |
| `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. |
@@ -454,7 +456,7 @@
| `wal.create_topic_timeout` | String | `30s` | The timeout for creating a Kafka topic.<br/>**It's only used when the provider is `kafka`**. |
| `event_recorder` | -- | -- | Configuration options for the event recorder. |
| `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. |
| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `region_migration`,<br/>`create_database`, `alter_database`, `drop_database`, `create_flow`,<br/>`drop_flow`, `create_table`, `create_logical_tables`, `alter_table`,<br/>`alter_logical_tables`, `drop_table`, `undrop_table`, `purge_dropped_table`,<br/>`truncate_table`, `create_view`, `drop_view`, `repartition`,<br/>`repartition_group`, `wal_prune`, `batch_gc`.<br/>When omitted, all current and future event types are recorded.<br/>Set to an empty array to disable event recording. |
| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `region_migration`,<br/>`create_database`, `alter_database`, `drop_database`, `create_flow`,<br/>`drop_flow`, `create_table`, `create_logical_tables`, `alter_table`,<br/>`alter_logical_tables`, `drop_table`, `undrop_table`, `purge_dropped_table`,<br/>`truncate_table`, `create_view`, `drop_view`, `repartition`,<br/>`repartition_group`, `wal_prune`, `batch_gc`, `reconcile_table`.<br/>When omitted, all current and future event types are recorded.<br/>Set to an empty array to disable event recording. |
| `stats_persistence` | -- | -- | Configuration options for the stats persistence. |
| `stats_persistence.ttl` | String | `0s` | TTL for the stats table that will be used to store the stats.<br/>Set to `0s` to disable stats persistence.<br/>Default is `0s`.<br/>If you want to enable stats persistence, set the TTL to a value greater than 0.<br/>It is recommended to set a small value, e.g., `3h`. |
| `stats_persistence.interval` | String | `10m` | The interval to persist the stats. Default is `10m`.<br/>The minimum value is `10m`, if the value is less than `10m`, it will be overridden to `10m`. |
@@ -583,7 +585,7 @@
| `region_engine.mito.max_background_purges` | Integer | Auto | Max number of running background purge jobs (default: number of cpu cores). |
| `region_engine.mito.experimental_compaction_memory_limit` | String | 0 | Memory budget for compaction tasks.<br/>Supports absolute size (e.g., "2GiB", "512MB") or percentage of system memory (e.g., "50%").<br/>Setting it to 0 or "unlimited" disables the limit. |
| `region_engine.mito.experimental_compaction_on_exhausted` | String | wait | Behavior when compaction cannot acquire memory from the budget.<br/>Options: "wait" (default, 10s), "wait(<duration>)", "fail" |
| `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. |
| `region_engine.mito.auto_flush_interval` | String | `10m` | 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.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`. |
config/datanode.example.toml
+1 -1
View File
@@ -462,7 +462,7 @@ compress_manifest = false
#+ experimental_compaction_on_exhausted = "wait"
## Interval to auto flush a region if it has not flushed yet.
auto_flush_interval = "1h"
auto_flush_interval = "10m"
## 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.
## @toml2docs:none-default="Auto"
config/frontend.example.toml
+3
View File
@@ -241,6 +241,9 @@ enable = true
experimental_enable_exponential_histogram = false
## Maximum spans per trace ingest chunk. Set to 0 to disable splitting.
trace_ingest_chunk_size = 512
## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
## resource attributes, so metrics-only services reach the semantic graph.
experimental_enable_resource_info = false
## Prometheus remote storage options
[prom_store]
config/metasrv.example.toml
+1 -1
View File
@@ -320,7 +320,7 @@ ttl = "90d"
## `drop_flow`, `create_table`, `create_logical_tables`, `alter_table`,
## `alter_logical_tables`, `drop_table`, `undrop_table`, `purge_dropped_table`,
## `truncate_table`, `create_view`, `drop_view`, `repartition`,
## `repartition_group`, `wal_prune`, `batch_gc`.
## `repartition_group`, `wal_prune`, `batch_gc`, `reconcile_table`.
## When omitted, all current and future event types are recorded.
## Set to an empty array to disable event recording.
#+ event_types = ["region_migration"]
config/standalone.example.toml
+5 -2
View File
@@ -208,6 +208,9 @@ enable = true
experimental_enable_exponential_histogram = false
## Maximum spans per trace ingest chunk. Set to 0 to disable splitting.
trace_ingest_chunk_size = 512
## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
## resource attributes, so metrics-only services reach the semantic graph.
experimental_enable_resource_info = false
## Prometheus remote storage options
[prom_store]
@@ -622,7 +625,7 @@ compress_manifest = false
#+ experimental_compaction_on_exhausted = "wait"
## Interval to auto flush a region if it has not flushed yet.
auto_flush_interval = "1h"
auto_flush_interval = "10m"
## 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.
## @toml2docs:none-default="Auto"
@@ -927,7 +930,7 @@ ttl = "90d"
## `alter_database`, `drop_database`, `create_flow`, `drop_flow`,
## `create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
## `drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
## `create_view`, `drop_view`, `admin_function`.
## `create_view`, `drop_view`, `admin_function`, `reconcile_table`.
## When omitted, all current and future event types are recorded.
## Set to an empty array to disable event recording.
#+ event_types = ["create_database"]
docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md
+4 -4
View File
@@ -149,7 +149,7 @@ The decisions behind the shape:
- **`provenance` is part of edge identity**, so a declared edge and a derived edge between the same pair coexist, and a user-declared edge survives when derivation is off.
- **`confidence` expresses derivation certainty, not statistical completeness**: `1.0` for a successfully paired or declared edge, `< 1.0` for virtual-node or agent-inferred edges. It does not correct for trace sampling.
- **RED metrics describe the observed span-pair population** — the `service_graph` connector semantics. Under sampling, counts understate true traffic, and ratios (error rate, mean duration) are representative only when sampling is unbiased with respect to status and latency — tail sampling that keeps errors and slow traces skews both.
- **Endpoint encoding (v1).** A single-attribute identity is the value verbatim; a composite identity is the attributes sorted by key, rendered `k1=v1,k2=v2` — a human-readable traversal key, not a collision-free canonical encoding. Nodes carry the structured `entity_id_attrs`; edges carry only the string key. Known v1 limitation: composite components containing unescaped `,` or `=` are not guaranteed to be distinguished.
- **Endpoint encoding.** An identity renders as its values in declared order (broad to narrow), joined by `,` with `,` and `\` escaped in each value. The identifying attribute *keys* are deliberately absent: the same identity reaches the graph under different column names per signal — a trace table's `service_name` against a metric table's `job` — and encoding the keys would split one entity into one entity per signal. Where a source carries a qualifier separately (`service.namespace` beside `service.name`) the convention folds it in as `<qualifier>/<value>`, reproducing the identity a source that pre-composed it (Prometheus `job`) already emits. Nodes carry the structured `entity_id_attrs`; edges carry only the string key.
- **Hand-declared edges** (`provenance = 'declared'`) live in a physical table of the same shape (plus a business validity window) in `greptime_private` and are unioned into the computed `semantic_relationships`.
The relationship vocabulary is small, typed, and inverse-paired; stored direction is `src -> dst`, the inverse is a query concern:
@@ -197,7 +197,7 @@ GROUP BY 1, src_id, dst_id;
Edge endpoints are built from each trace table's `service` entity declaration (the same id expression the registry uses), so edges land on exactly the entity ids the registry emits, self-calls are judged on the full endpoint id, and a table whose declaration is unusable contributes no edges rather than silently switching identity. The join is bounded by a time-proximity window and the results of **all** trace tables combine into one edge per window with merged RED metrics. One assumption to note: a deployment can route traces to multiple tables (`x-greptime-trace-table-name`), so one distributed trace is not guaranteed to live in a single table — pairing must account for that, and how is an implementation concern. Uninstrumented peers become **virtual nodes**: a client span with no matching server span yields an edge to a synthetic node named from `peer.service`/`db.name`/`server.address`.
**The entity registry.** Per declaring table: filter to the window, project (time bin, casts, JSON assembly), `DISTINCT`. The per-table results are unioned without any cross-source merge: no descriptive-conflict or JSON-merging policy lives in the derivation. Rows whose identity columns are NULL identify nothing and are filtered out; single-column ids are used verbatim, composite ids render the sorted `k=v` form plus the `entity_id_attrs` JSON. Consumers deduplicate with `SELECT DISTINCT entity_type, entity_id`; a unique node set is the snapshot relation's job (see *Querying the graph*).
**The entity registry.** Per declaring table: filter to the window, project (time bin, casts, JSON assembly), `DISTINCT`. The per-table results are unioned without any cross-source merge: no descriptive-conflict or JSON-merging policy lives in the derivation. Rows whose identity columns are NULL identify nothing and are filtered out; ids render as their escaped values joined in declared order, with the `entity_id_attrs` JSON alongside once an id is assembled from more than one column. Consumers deduplicate with `SELECT DISTINCT entity_type, entity_id`; a unique node set is the snapshot relation's job (see *Querying the graph*).
**Attribute edges.** Shared attributes provide join keys; co-declaration or a relationship template provides the relationship semantics. A shared value alone determines neither direction nor edge type, so the derivation is rule-based, not join-everything. A table declaring both a `service.instance` and a `host` identity on the same rows derives `runs_on` between them: the row itself witnesses the relationship, and the built-in vocabulary fixes the direction. Cross-table edges require a declared relationship template (endpoint mappings, `rel_type`, direction): a derivation *rule* declared on schema, not per-instance edge data.
@@ -284,7 +284,7 @@ Drawbacks:
- Read-time derivation costs a self-join over trace tables on every scan of the computed tables. The default one-hour window bounds it, but very large deployments will want materialisation (Future Work).
- `semantic_entities` exposes per-observation rows; deduplication falls on the consumer (`DISTINCT`) until the snapshot relation lands in M2.
- The v1 `k=v` id encoding can collide on hostile values (unescaped `,`/`=`); hardening is an open question.
- The id encoding is readable rather than canonical: escaping keeps components distinguishable, but a value containing the qualifier separator (`/`) is still ambiguous against a qualified name.
- The graph is only as connected as the identity columns tables actually share: the `job`/`service_name` misalignment above is the canonical example, and enrichment discipline gates connectivity.
# Implementation Plan
@@ -306,7 +306,7 @@ Drawbacks:
# Open Questions
1. **Endpoint encoding hardening.** The v1 `k=v` rendering (with `entity_id_attrs` JSON as the source of truth) is readable but collides when values contain `,`/`=`. Open: a fixed-width `entity_id_hash` join key (New-Relic-style); typed value encoding; the criteria for introducing the `semantic_key=column` mapping form; `schema_url`-gated entity merge. Related and unresolved: whether the identifying attribute *key* should participate in entity equality — a single-attribute id keeps only the value, so `job=api` and `service_name=api` currently name the same entity even when the keys differ semantically.
1. **Endpoint encoding hardening.** The rendering (with `entity_id_attrs` JSON as the source of truth) is readable rather than canonical. Open: a fixed-width `entity_id_hash` join key (New-Relic-style); typed value encoding; the criteria for introducing the `semantic_key=column` mapping form; `schema_url`-gated entity merge. Settled: identifying attribute *keys* do not participate in entity equality, so `job=api` and `service_name=api` name the same entity — the property that lets one service reach the graph from several signals.
2. **Identity reconciliation across signals.** Trace edges key on the trace table's service declaration; host edges key on `host.id`/`k8s.pod.uid`. Linking them relies on tables carrying both columns; how aggressively to enrich (a `k8sattributes`-equivalent) gates how connected the graph is.
3. **`GRAPH_TABLE` surface scope.** How much of SQL/PGQ to implement vs. a minimal single-hop subset. Start minimal and grow with demand.
4. **Visibility beyond table grants.** The derivation contract filters by table-read permission; whether topology needs finer policies (row-level, tenant-scoped visibility of cross-schema edges) is open.
src/client/src/database.rs
+336 -154
View File
@@ -12,20 +12,25 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashMap;
use std::pin::Pin;
use std::str::FromStr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, RwLock};
use std::task::{Context, Poll};
use std::time::Duration;
use api::v1::auth_header::AuthScheme;
#[cfg(feature = "testing")]
use api::v1::ddl_request::Expr as DdlExpr;
use api::v1::greptime_database_client::GreptimeDatabaseClient;
use api::v1::greptime_request::Request;
use api::v1::query_request::Query;
#[cfg(feature = "testing")]
use api::v1::{AlterTableExpr, CreateTableExpr, DdlRequest};
use api::v1::{
AlterTableExpr, AuthHeader, Basic, CreateTableExpr, DdlRequest, GreptimeRequest,
InsertRequests, QueryRequest, RequestHeader, RowInsertRequests,
AuthHeader, Basic, GreptimeRequest, InsertRequests, QueryRequest, RequestHeader,
RowInsertRequests,
};
use arc_swap::ArcSwapOption;
use arrow_flight::{FlightData, Ticket};
@@ -34,7 +39,7 @@ use base64::Engine;
use base64::prelude::BASE64_STANDARD;
use common_catalog::build_db_string;
use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
use common_error::ext::BoxedError;
use common_error::ext::{BoxedError, ErrorExt};
use common_grpc::flight::do_put::DoPutResponse;
use common_grpc::flight::{
FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage, SNAPSHOT_SEQS_METADATA_KEY,
@@ -49,21 +54,23 @@ use common_telemetry::{error, warn};
use futures::future;
use futures_util::{Stream, StreamExt, TryStreamExt};
use prost::Message;
use snafu::ResultExt;
use snafu::{IntoError, ResultExt};
use tonic::metadata::{AsciiMetadataKey, AsciiMetadataValue, MetadataMap, MetadataValue};
use tonic::transport::Channel;
use crate::error::{
ConvertFlightDataSnafu, Error, FlightGetSnafu, IllegalFlightMessagesSnafu,
ConvertFlightDataSnafu, Error, FlightGetSnafu, FlightStreamSnafu, IllegalFlightMessagesSnafu,
InvalidTonicMetadataValueSnafu,
};
use crate::flight::decode_flight_data;
use crate::flight::{FlightMessageReader, decode_flight_data};
use crate::{Client, Result, error, from_grpc_response};
type FlightDataStream = Pin<Box<dyn Stream<Item = FlightData> + Send>>;
type DoPutResponseStream = Pin<Box<dyn Stream<Item = Result<DoPutResponse>>>>;
const HINTS_METADATA_KEY: &str = "x-greptime-hints";
/// Terminal metrics associated with a query output.
///
/// For streaming outputs, metrics are only final after the stream is fully
@@ -257,19 +264,17 @@ fn attach_terminal_metrics(output: Output, terminal_metrics: &OutputMetrics) ->
}
async fn output_from_flight_message_stream<S>(
mut flight_message_stream: S,
remote_addr: String,
flight_message_stream: S,
) -> Result<OutputWithMetrics>
where
S: Stream<Item = Result<FlightMessage>> + Send + Unpin + 'static,
{
let Some(first_flight_message) = flight_message_stream.next().await else {
return IllegalFlightMessagesSnafu {
reason: "Expect the response not to be empty",
}
.fail();
};
let first_flight_message = first_flight_message?;
let mut reader = FlightMessageReader::new(remote_addr, flight_message_stream);
let first_flight_message = reader
.read_first()
.await
.map_err(|error| flight_stream_error(reader.remote_addr(), error))?;
match first_flight_message {
FlightMessage::AffectedRows { rows, metrics } => {
@@ -277,7 +282,10 @@ where
if let Some(metrics) = metrics {
terminal_metrics.update(Some(parse_terminal_metrics(&metrics)?));
}
let next_message = flight_message_stream.next().await.transpose()?;
let next_message = reader
.read_next()
.await
.map_err(|error| flight_stream_error(reader.remote_addr(), error))?;
match next_message {
None => terminal_metrics.mark_ready(),
Some(FlightMessage::Metrics(s)) if terminal_metrics.get().is_none() => {
@@ -316,15 +324,19 @@ where
);
let schema_cloned = schema.clone();
let stream = Box::pin(stream!({
while let Some(flight_message_item) = flight_message_stream.next().await {
let flight_message = match flight_message_item {
Ok(message) => message,
Err(e) => {
yield Err(BoxedError::new(e)).context(ExternalSnafu);
loop {
let flight_message = match reader.read_next().await {
Ok(Some(message)) => message,
Ok(None) => break,
Err(error) => {
yield Err(BoxedError::new(flight_stream_error(
reader.remote_addr(),
error,
)))
.context(ExternalSnafu);
break;
}
};
match flight_message {
FlightMessage::RecordBatch(arrow_batch) => {
yield Ok(RecordBatch::from_df_record_batch(
@@ -371,6 +383,25 @@ where
}
}
fn flight_stream_error(addr: &str, error: Error) -> Error {
let tonic_code = error.tonic_code().unwrap_or(tonic::Code::Unknown);
let message = error.to_string();
if error.status_code().should_log_error() {
error!(
error; "Failed to receive Flight data, addr: {}, code: {}",
addr,
tonic_code
);
}
FlightStreamSnafu {
addr: addr.to_string(),
tonic_code,
message,
}
.into_error(BoxedError::new(error))
}
#[derive(Clone, Debug, Default)]
pub struct Database {
// The "catalog" and "schema" to be used in processing the requests at the server side.
@@ -389,6 +420,42 @@ pub struct Database {
ctx: FlightContext,
}
#[derive(Default)]
struct FlightRequestOptions {
hints: Option<String>,
flow_extensions: Option<String>,
snapshot_seqs: Option<String>,
timeout: Option<Duration>,
}
impl FlightRequestOptions {
fn apply_to<T>(self, request: &mut tonic::Request<T>) -> Result<()> {
let metadata = request.metadata_mut();
if let Some(hints) = self.hints {
Database::put_metadata_value(metadata, HINTS_METADATA_KEY, hints)?;
}
if let Some(flow_extensions) = self.flow_extensions {
Database::put_metadata_value(metadata, FLOW_EXTENSIONS_METADATA_KEY, flow_extensions)?;
}
if let Some(snapshot_seqs) = self.snapshot_seqs {
Database::put_metadata_value(metadata, SNAPSHOT_SEQS_METADATA_KEY, snapshot_seqs)?;
}
if let Some(timeout) = self.timeout {
request.set_timeout(timeout);
}
Ok(())
}
}
/// A single Flight DoGet request to a [`Database`].
///
/// The builder carries request-scoped metadata and options. It does not modify
/// the underlying [`Database`], so its configuration cannot affect later RPCs.
pub struct DatabaseFlightRequest<'a> {
database: &'a Database,
options: FlightRequestOptions,
}
pub struct DatabaseClient {
pub addr: String,
pub inner: GreptimeDatabaseClient<Channel>,
@@ -483,6 +550,14 @@ impl Database {
});
}
/// Creates a builder for a single Flight DoGet request.
pub fn flight_request(&self) -> DatabaseFlightRequest<'_> {
DatabaseFlightRequest {
database: self,
options: FlightRequestOptions::default(),
}
}
/// Make an InsertRequests request to the database.
pub async fn insert(&self, requests: InsertRequests) -> Result<u32> {
self.handle(Request::Inserts(requests)).await
@@ -538,44 +613,31 @@ impl Database {
}
fn put_hints(metadata: &mut MetadataMap, hints: &[(&str, &str)]) -> Result<()> {
let Some(value) = hints
.iter()
.map(|(k, v)| format!("{}={}", k, v))
.reduce(|a, b| format!("{},{}", a, b))
else {
let Some(value) = Self::encode_hints(hints) else {
return Ok(());
};
let key = AsciiMetadataKey::from_static("x-greptime-hints");
let value = AsciiMetadataValue::from_str(&value).context(InvalidTonicMetadataValueSnafu)?;
metadata.insert(key, value);
Ok(())
Self::put_metadata_value(metadata, HINTS_METADATA_KEY, value)
}
fn put_flow_extensions(
metadata: &mut MetadataMap,
flow_extensions: &[(&str, &str)],
) -> Result<()> {
if flow_extensions.is_empty() {
return Ok(());
}
let value = serde_json::to_string(&flow_extensions.to_vec())
.expect("flow extension pairs should serialize");
Self::put_metadata_value(metadata, FLOW_EXTENSIONS_METADATA_KEY, value)
fn encode_hints(hints: &[(&str, &str)]) -> Option<String> {
hints
.iter()
.map(|(k, v)| format!("{}={}", k, v))
.reduce(|a, b| format!("{},{}", a, b))
}
fn put_snapshot_seqs(
metadata: &mut MetadataMap,
snapshot_seqs: &std::collections::HashMap<u64, u64>,
) -> Result<()> {
if snapshot_seqs.is_empty() {
return Ok(());
}
fn encode_flow_extensions(flow_extensions: &[(&str, &str)]) -> Option<String> {
(!flow_extensions.is_empty()).then(|| {
serde_json::to_string(&flow_extensions.to_vec())
.expect("flow extension pairs should serialize")
})
}
let value =
serde_json::to_string(snapshot_seqs).expect("snapshot sequence map should serialize");
Self::put_metadata_value(metadata, SNAPSHOT_SEQS_METADATA_KEY, value)
fn encode_snapshot_seqs(snapshot_seqs: &HashMap<u64, u64>) -> Option<String> {
(!snapshot_seqs.is_empty()).then(|| {
serde_json::to_string(snapshot_seqs).expect("snapshot sequence map should serialize")
})
}
fn put_metadata_value(
@@ -673,7 +735,7 @@ impl Database {
where
S: AsRef<str>,
{
self.sql_with_hint(sql, &[]).await
self.flight_request().sql(sql).await
}
/// Executes a SQL query with optional hints for query optimization.
@@ -681,12 +743,7 @@ impl Database {
where
S: AsRef<str>,
{
let request = Request::Query(QueryRequest {
query: Some(Query::Sql(sql.as_ref().to_string())),
});
self.do_get(request, hints, &[], &Default::default())
.await
.map(OutputWithMetrics::into_output)
self.flight_request().with_hints(hints).sql(sql).await
}
/// Executes a SQL query and returns the output with terminal metrics.
@@ -701,77 +758,33 @@ impl Database {
where
S: AsRef<str>,
{
self.query_with_terminal_metrics_and_flow_extensions(
QueryRequest {
query: Some(Query::Sql(sql.as_ref().to_string())),
},
hints,
&[],
&Default::default(),
)
.await
self.flight_request()
.with_hints(hints)
.sql_with_terminal_metrics(sql)
.await
}
/// Executes a logical plan directly without SQL parsing.
pub async fn logical_plan(&self, logical_plan: Vec<u8>) -> Result<Output> {
self.query_with_terminal_metrics_and_flow_extensions(
QueryRequest {
query: Some(Query::LogicalPlan(logical_plan)),
},
&[],
&[],
&Default::default(),
)
.await
.map(OutputWithMetrics::into_output)
}
/// Executes a query and carries flow extensions through Flight metadata.
///
/// This is the lower-level terminal-metrics API for Flow callers that need
/// to pass JSON-bearing flow extensions without going through hint metadata.
pub async fn query_with_terminal_metrics_and_flow_extensions(
&self,
request: QueryRequest,
hints: &[(&str, &str)],
flow_extensions: &[(&str, &str)],
snapshot_seqs: &std::collections::HashMap<u64, u64>,
) -> Result<OutputWithMetrics> {
self.do_get(
Request::Query(request),
hints,
flow_extensions,
snapshot_seqs,
)
.await
self.flight_request().logical_plan(logical_plan).await
}
/// Creates a new table using the provided table expression.
#[cfg(feature = "testing")]
pub async fn create(&self, expr: CreateTableExpr) -> Result<Output> {
let request = Request::Ddl(DdlRequest {
expr: Some(DdlExpr::CreateTable(expr)),
});
self.do_get(request, &[], &[], &Default::default())
.await
.map(OutputWithMetrics::into_output)
self.flight_request().create(expr).await
}
/// Alters an existing table using the provided alter expression.
#[cfg(feature = "testing")]
pub async fn alter(&self, expr: AlterTableExpr) -> Result<Output> {
let request = Request::Ddl(DdlRequest {
expr: Some(DdlExpr::AlterTable(expr)),
});
self.do_get(request, &[], &[], &Default::default())
.await
.map(OutputWithMetrics::into_output)
self.flight_request().alter(expr).await
}
async fn do_get(
&self,
request: Request,
hints: &[(&str, &str)],
flow_extensions: &[(&str, &str)],
snapshot_seqs: &std::collections::HashMap<u64, u64>,
options: FlightRequestOptions,
) -> Result<OutputWithMetrics> {
let request = self.to_rpc_request(request);
let request = Ticket {
@@ -779,12 +792,10 @@ impl Database {
};
let mut request = tonic::Request::new(request);
let metadata = request.metadata_mut();
Self::put_hints(metadata, hints)?;
Self::put_flow_extensions(metadata, flow_extensions)?;
Self::put_snapshot_seqs(metadata, snapshot_seqs)?;
options.apply_to(&mut request)?;
let mut client = self.client.make_flight_client(false, false)?;
let remote_addr = client.addr().to_string();
let response = client.mut_inner().do_get(request).await.or_else(|e| {
let tonic_code = e.code();
@@ -796,7 +807,7 @@ impl Database {
e
);
Err(BoxedError::new(e)).with_context(|_| FlightGetSnafu {
addr: client.addr().to_string(),
addr: remote_addr.clone(),
tonic_code,
})
})?;
@@ -808,7 +819,7 @@ impl Database {
future::ready(decode_flight_data(&mut decoder, flight_data))
});
output_from_flight_message_stream(flight_message_stream).await
output_from_flight_message_stream(remote_addr, flight_message_stream).await
}
/// Ingest a stream of [RecordBatch]es that belong to a table, using Arrow Flight's "`DoPut`"
@@ -847,6 +858,98 @@ impl Database {
}
}
impl<'a> DatabaseFlightRequest<'a> {
/// Adds query optimization hints to this Flight request.
pub fn with_hints(mut self, hints: &[(&str, &str)]) -> Self {
self.options.hints = Database::encode_hints(hints);
self
}
/// Adds Flow extensions to this Flight request.
pub fn with_flow_extensions(mut self, flow_extensions: &[(&str, &str)]) -> Self {
self.options.flow_extensions = Database::encode_flow_extensions(flow_extensions);
self
}
/// Adds snapshot sequence fences to this Flight request.
pub fn with_snapshot_seqs(mut self, snapshot_seqs: &HashMap<u64, u64>) -> Self {
self.options.snapshot_seqs = Database::encode_snapshot_seqs(snapshot_seqs);
self
}
/// Sets a timeout for this Flight request only.
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.options.timeout = Some(timeout);
self
}
/// Executes a SQL query.
pub async fn sql<S>(self, sql: S) -> Result<Output>
where
S: AsRef<str>,
{
let request = Request::Query(QueryRequest {
query: Some(Query::Sql(sql.as_ref().to_string())),
});
self.do_get(request)
.await
.map(OutputWithMetrics::into_output)
}
/// Executes a SQL query and returns terminal metrics.
pub async fn sql_with_terminal_metrics<S>(self, sql: S) -> Result<OutputWithMetrics>
where
S: AsRef<str>,
{
self.query_with_terminal_metrics(QueryRequest {
query: Some(Query::Sql(sql.as_ref().to_string())),
})
.await
}
/// Executes a logical plan directly without SQL parsing.
pub async fn logical_plan(self, logical_plan: Vec<u8>) -> Result<Output> {
self.query_with_terminal_metrics(QueryRequest {
query: Some(Query::LogicalPlan(logical_plan)),
})
.await
.map(OutputWithMetrics::into_output)
}
/// Executes a query and returns terminal metrics.
pub async fn query_with_terminal_metrics(
self,
request: QueryRequest,
) -> Result<OutputWithMetrics> {
self.do_get(Request::Query(request)).await
}
/// Creates a new table using the provided table expression.
#[cfg(feature = "testing")]
pub async fn create(self, expr: CreateTableExpr) -> Result<Output> {
self.do_get(Request::Ddl(DdlRequest {
expr: Some(DdlExpr::CreateTable(expr)),
}))
.await
.map(OutputWithMetrics::into_output)
}
/// Alters an existing table using the provided alter expression.
#[cfg(feature = "testing")]
pub async fn alter(self, expr: AlterTableExpr) -> Result<Output> {
self.do_get(Request::Ddl(DdlRequest {
expr: Some(DdlExpr::AlterTable(expr)),
}))
.await
.map(OutputWithMetrics::into_output)
}
async fn do_get(self, request: Request) -> Result<OutputWithMetrics> {
let Self { database, options } = self;
database.do_get(request, options).await
}
}
/// by grpc standard, only `Unavailable` is retryable, see: https://github.com/grpc/grpc/blob/master/doc/statuscodes.md#status-codes-and-their-use-in-grpc
pub fn is_grpc_retryable(err: &tonic::Status) -> bool {
matches!(err.code(), tonic::Code::Unavailable)
@@ -934,12 +1037,14 @@ mod tests {
#[test]
fn test_put_flow_extensions_preserves_comma_bearing_values() {
let mut metadata = MetadataMap::new();
Database::put_flow_extensions(
Database::put_metadata_value(
&mut metadata,
&[
FLOW_EXTENSIONS_METADATA_KEY,
Database::encode_flow_extensions(&[
("flow.return_region_seq", "true"),
("flow.incremental_after_seqs", r#"{"1":10,"2":20}"#),
],
])
.unwrap(),
)
.unwrap();
@@ -969,7 +1074,12 @@ mod tests {
(9_007_199_254_740_993_u64, 9_007_199_254_740_995_u64),
]);
Database::put_snapshot_seqs(&mut metadata, &snapshot_seqs).unwrap();
Database::put_metadata_value(
&mut metadata,
SNAPSHOT_SEQS_METADATA_KEY,
Database::encode_snapshot_seqs(&snapshot_seqs).unwrap(),
)
.unwrap();
let value = metadata
.get(SNAPSHOT_SEQS_METADATA_KEY)
@@ -980,6 +1090,50 @@ mod tests {
assert_eq!(decoded, snapshot_seqs);
}
#[test]
fn test_flight_request_builder_applies_request_options() {
let database = Database::new("greptime", "public", Client::default());
let snapshot_seqs = HashMap::from([(42, 99)]);
let request = database
.flight_request()
.with_hints(&[("query_parallelism", "1")])
.with_flow_extensions(&[("flow.return_region_seq", "true")])
.with_snapshot_seqs(&snapshot_seqs)
.with_timeout(Duration::from_millis(50));
let mut tonic_request = tonic::Request::new(());
request.options.apply_to(&mut tonic_request).unwrap();
let metadata = tonic_request.metadata();
assert_eq!(
metadata.get(HINTS_METADATA_KEY).unwrap(),
"query_parallelism=1"
);
assert_eq!(
serde_json::from_str::<Vec<(String, String)>>(
metadata
.get(FLOW_EXTENSIONS_METADATA_KEY)
.unwrap()
.to_str()
.unwrap(),
)
.unwrap(),
vec![("flow.return_region_seq".to_string(), "true".to_string())]
);
assert_eq!(
serde_json::from_str::<HashMap<u64, u64>>(
metadata
.get(SNAPSHOT_SEQS_METADATA_KEY)
.unwrap()
.to_str()
.unwrap(),
)
.unwrap(),
snapshot_seqs
);
assert!(metadata.get("grpc-timeout").is_some());
}
#[test]
fn test_flight_ctx() {
let mut ctx = FlightContext::default();
@@ -1020,6 +1174,28 @@ mod tests {
assert_eq!(expected.should_retry(), actual.should_retry());
}
#[test]
fn test_flight_stream_error_preserves_addr_and_message() {
let error = flight_stream_error(
"127.0.0.1:4001",
Status::out_of_range("message length too large").into(),
);
assert!(matches!(
&error,
Error::FlightStream {
addr,
tonic_code: Code::OutOfRange,
message,
..
} if addr == "127.0.0.1:4001" && message == "message length too large"
));
assert_eq!(
"Failed to receive Flight data from 127.0.0.1:4001, code: Operation was attempted past the valid range: message length too large",
error.to_string(),
);
}
#[test]
fn test_from_tonic_status_with_retry_hint() {
let mut headers = HeaderMap::new();
@@ -1117,13 +1293,13 @@ mod tests {
#[tokio::test]
async fn test_affected_rows_inline_metrics_are_parsed() {
let output = output_from_flight_message_stream(futures_util::stream::iter(vec![Ok(
FlightMessage::AffectedRows {
let output = output_from_flight_message_stream(
"test-peer".to_string(),
futures_util::stream::iter(vec![Ok(FlightMessage::AffectedRows {
rows: 3,
metrics: Some(terminal_metrics_json()),
},
)]
as Vec<Result<FlightMessage>>))
})] as Vec<Result<FlightMessage>>),
)
.await
.unwrap();
@@ -1138,14 +1314,16 @@ mod tests {
#[tokio::test]
async fn test_affected_rows_inline_metrics_rejects_trailing_metrics() {
let metrics_json = terminal_metrics_json();
let err = output_from_flight_message_stream(futures_util::stream::iter(vec![
Ok(FlightMessage::AffectedRows {
rows: 3,
metrics: Some(metrics_json.clone()),
}),
Ok(FlightMessage::Metrics(metrics_json)),
]
as Vec<Result<FlightMessage>>))
let err = output_from_flight_message_stream(
"test-peer".to_string(),
futures_util::stream::iter(vec![
Ok(FlightMessage::AffectedRows {
rows: 3,
metrics: Some(metrics_json.clone()),
}),
Ok(FlightMessage::Metrics(metrics_json)),
] as Vec<Result<FlightMessage>>),
)
.await
.unwrap_err();
@@ -1167,12 +1345,14 @@ mod tests {
vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef],
)
.unwrap();
let output = output_from_flight_message_stream(futures_util::stream::iter(vec![
Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())),
Ok(FlightMessage::Metrics("{not-json}".to_string())),
]
as Vec<Result<FlightMessage>>))
let output = output_from_flight_message_stream(
"test-peer".to_string(),
futures_util::stream::iter(vec![
Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())),
Ok(FlightMessage::Metrics("{not-json}".to_string())),
] as Vec<Result<FlightMessage>>),
)
.await
.unwrap();
let terminal_metrics = output.metrics.clone();
@@ -1211,18 +1391,20 @@ mod tests {
vec![Arc::new(Int32Vector::from_slice([2])) as VectorRef],
)
.unwrap();
let output = output_from_flight_message_stream(futures_util::stream::iter(vec![
Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
Ok(FlightMessage::RecordBatch(
first_batch.into_df_record_batch(),
)),
Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(1))),
Ok(FlightMessage::RecordBatch(
second_batch.into_df_record_batch(),
)),
Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(2))),
]
as Vec<Result<FlightMessage>>))
let output = output_from_flight_message_stream(
"test-peer".to_string(),
futures_util::stream::iter(vec![
Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
Ok(FlightMessage::RecordBatch(
first_batch.into_df_record_batch(),
)),
Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(1))),
Ok(FlightMessage::RecordBatch(
second_batch.into_df_record_batch(),
)),
Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(2))),
] as Vec<Result<FlightMessage>>),
)
.await
.unwrap();
let terminal_metrics = output.metrics.clone();
src/client/src/error.rs
+18
View File
@@ -40,6 +40,21 @@ pub enum Error {
source: BoxedError,
},
#[snafu(display(
"Failed to receive Flight data from {}, code: {}: {}",
addr,
tonic_code,
message
))]
FlightStream {
addr: String,
tonic_code: Code,
message: String,
source: BoxedError,
#[snafu(implicit)]
location: Location,
},
#[snafu(display("Failed to convert FlightData"))]
ConvertFlightData {
#[snafu(implicit)]
@@ -154,6 +169,7 @@ impl ErrorExt for Error {
Error::Server { code, .. } | Error::Tonic { code, .. } => *code,
Error::FlightGet { source, .. }
| Error::FlightStream { source, .. }
| Error::RegionServer { source, .. }
| Error::FlowServer { source, .. } => source.status_code(),
Error::CreateChannel { source, .. }
@@ -174,6 +190,7 @@ impl ErrorExt for Error {
match self {
Error::Tonic { retry_hint, .. } => *retry_hint,
Error::FlightGet { source, .. }
| Error::FlightStream { source, .. }
| Error::RegionServer { source, .. }
| Error::FlowServer { source, .. }
| Error::External { source, .. } => source.retry_hint(),
@@ -193,6 +210,7 @@ impl Error {
pub fn tonic_code(&self) -> Option<Code> {
match self {
Self::FlightGet { tonic_code, .. }
| Self::FlightStream { tonic_code, .. }
| Self::RegionServer {
code: tonic_code, ..
}
src/client/src/flight.rs
+72 -2
View File
@@ -12,12 +12,82 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::pin::Pin;
use arrow_flight::FlightData;
use common_grpc::flight::{FlightDecoder, FlightMessage};
use snafu::ResultExt;
use futures_util::stream::Peekable;
use futures_util::{Stream, StreamExt};
use snafu::{OptionExt, ResultExt};
use crate::Result;
use crate::error::{ConvertFlightDataSnafu, Error};
use crate::error::{ConvertFlightDataSnafu, Error, IllegalFlightMessagesSnafu};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum FlightMessageKind {
Schema,
RecordBatch,
AffectedRows,
Metrics,
}
impl From<&FlightMessage> for FlightMessageKind {
fn from(message: &FlightMessage) -> Self {
match message {
FlightMessage::Schema(_) => Self::Schema,
FlightMessage::RecordBatch(_) => Self::RecordBatch,
FlightMessage::AffectedRows { .. } => Self::AffectedRows,
FlightMessage::Metrics(_) => Self::Metrics,
}
}
}
pub(crate) struct FlightMessageReader<S: Stream + Unpin> {
/// Remote Flight peer associated with this response stream.
remote_addr: String,
messages: Peekable<S>,
}
impl<S> FlightMessageReader<S>
where
S: Stream<Item = Result<FlightMessage>> + Unpin,
{
pub(crate) fn new(remote_addr: impl Into<String>, messages: S) -> Self {
Self {
remote_addr: remote_addr.into(),
messages: messages.peekable(),
}
}
pub(crate) fn remote_addr(&self) -> &str {
&self.remote_addr
}
pub(crate) async fn read_first(&mut self) -> Result<FlightMessage> {
self.read_next().await?.context(IllegalFlightMessagesSnafu {
reason: "Expect the response not to be empty",
})
}
pub(crate) async fn read_next(&mut self) -> Result<Option<FlightMessage>> {
self.messages.next().await.transpose()
}
pub(crate) async fn peek_next_message_kind(&mut self) -> Result<Option<FlightMessageKind>> {
match Pin::new(&mut self.messages).peek().await {
Some(Ok(message)) => Ok(Some(message.into())),
None => Ok(None),
Some(Err(_)) => match self.read_next().await {
// `peek` only borrows the error; consume it to preserve the source error.
Err(error) => Err(error),
Ok(_) => IllegalFlightMessagesSnafu {
reason: "Flight stream changed after peek".to_string(),
}
.fail(),
},
}
}
}
pub(crate) fn decode_flight_data(
decoder: &mut FlightDecoder,
src/client/src/region.rs
+135 -54
View File
@@ -47,7 +47,7 @@ use crate::error::{
self, FlightGetSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu,
MissingFieldSnafu, Result, ServerSnafu,
};
use crate::flight::decode_flight_data;
use crate::flight::{FlightMessageKind, FlightMessageReader, decode_flight_data};
use crate::{Client, metrics};
const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10);
@@ -203,26 +203,22 @@ impl RegionRequester {
async fn recordbatches_from_flight_message_stream<S>(
addr: String,
mut flight_message_stream: S,
flight_message_stream: S,
) -> Result<SendableRecordBatchStream>
where
S: Stream<Item = Result<FlightMessage>> + Send + Unpin + 'static,
{
let Some(first_flight_message) = flight_message_stream.next().await else {
return IllegalFlightMessagesSnafu {
reason: "Expect the response not to be empty",
}
.fail()
.map_err(|error| flight_stream_error(&addr, error));
};
let FlightMessage::Schema(schema) =
first_flight_message.map_err(|e| flight_stream_error(&addr, e))?
let mut reader = FlightMessageReader::new(addr.clone(), flight_message_stream);
let FlightMessage::Schema(schema) = reader
.read_first()
.await
.map_err(|error| flight_stream_error(reader.remote_addr(), error))?
else {
return IllegalFlightMessagesSnafu {
reason: "Expect schema to be the first flight message",
}
.fail()
.map_err(|error| flight_stream_error(&addr, error));
.map_err(|error| flight_stream_error(reader.remote_addr(), error));
};
let metrics = Arc::new(ArcSwapOption::from(None));
@@ -233,31 +229,23 @@ where
let schema =
Arc::new(datatypes::schema::Schema::try_from(schema).context(error::ConvertSchemaSnafu)?);
let schema_cloned = schema.clone();
let stream_addr = addr.clone();
let stream_addr = addr;
let stream = Box::pin(stream!({
let _span = tracing_context.attach(common_telemetry::tracing::info_span!(
"poll_flight_data_stream"
));
let mut buffered_message: Option<FlightMessage> = None;
let mut stream_ended = false;
while !stream_ended {
// get the next message from the buffered message or read from the flight message stream
let flight_message_item = if let Some(msg) = buffered_message.take() {
Some(Ok(msg))
} else {
flight_message_stream.next().await
};
let flight_message = match flight_message_item {
Some(Ok(message)) => message,
Some(Err(e)) => {
yield Err(BoxedError::new(flight_stream_error(&stream_addr, e)))
let flight_message = match reader.read_next().await {
Ok(Some(message)) => message,
Ok(None) => break,
Err(error) => {
yield Err(BoxedError::new(flight_stream_error(&stream_addr, error)))
.context(ExternalSnafu);
break;
}
None => break,
};
match flight_message {
@@ -265,42 +253,54 @@ where
let result_to_yield =
RecordBatch::from_df_record_batch(schema_cloned.clone(), record_batch);
// get the next message from the stream. normally it should be a metrics message.
if let Some(next_flight_message_result) = flight_message_stream.next().await {
match next_flight_message_result {
Ok(FlightMessage::Metrics(s)) => {
let m = serde_json::from_str(&s).ok().map(Arc::new);
metrics_ref.swap(m);
}
Ok(FlightMessage::RecordBatch(rb)) => {
// for some reason it's not a metrics message, so we need to buffer this record batch
// and yield it in the next iteration.
buffered_message = Some(FlightMessage::RecordBatch(rb));
}
Ok(_) => {
yield IllegalFlightMessagesSnafu {
reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message"
}
.fail()
.map_err(BoxedError::new)
.context(ExternalSnafu);
break;
}
Err(e) => {
yield Err(BoxedError::new(flight_stream_error(&stream_addr, e)))
// Metrics follow a batch so MergeScan can observe them before yielding it.
match reader.peek_next_message_kind().await {
Ok(Some(FlightMessageKind::Metrics)) => {
let metrics_message = match reader.read_next().await {
Ok(Some(FlightMessage::Metrics(metrics))) => metrics,
Ok(Some(_) | None) => {
yield IllegalFlightMessagesSnafu {
reason: "Flight stream changed after peek",
}
.fail()
.map_err(BoxedError::new)
.context(ExternalSnafu);
break;
break;
}
Err(error) => {
yield Err(BoxedError::new(flight_stream_error(
&stream_addr,
error,
)))
.context(ExternalSnafu);
break;
}
};
let metrics = serde_json::from_str(&metrics_message).ok().map(Arc::new);
metrics_ref.swap(metrics);
}
Ok(Some(FlightMessageKind::RecordBatch)) => {}
Ok(Some(FlightMessageKind::Schema | FlightMessageKind::AffectedRows)) => {
yield IllegalFlightMessagesSnafu {
reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message"
}
.fail()
.map_err(BoxedError::new)
.context(ExternalSnafu);
break;
}
Ok(None) => stream_ended = true,
Err(error) => {
yield Err(BoxedError::new(flight_stream_error(&stream_addr, error)))
.context(ExternalSnafu);
break;
}
} else {
// the stream has ended
stream_ended = true;
}
yield Ok(result_to_yield);
}
FlightMessage::Metrics(s) => {
// just a branch in case of some metrics message comes after other things.
// Metrics may arrive before the next RecordBatch.
let m = serde_json::from_str(&s).ok().map(Arc::new);
metrics_ref.swap(m);
continue;
@@ -412,6 +412,7 @@ mod test {
RemoteDynFilterUnregister, RemoteDynFilterUpdate, region_request, remote_dyn_filter_request,
};
use common_recordbatch::adapter::RecordBatchMetrics;
use datatypes::arrow::array::Int32Array;
use datatypes::prelude::{ConcreteDataType, VectorRef};
use datatypes::schema::{ColumnSchema, Schema};
use datatypes::vectors::Int32Vector;
@@ -598,6 +599,86 @@ mod test {
assert_eq!(metrics.elapsed_compute, 7);
}
#[tokio::test]
async fn test_record_batch_stream_updates_following_metrics_before_yielding_batch() {
let schema = test_schema();
let batch = RecordBatch::new(
schema.clone(),
vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef],
)
.unwrap();
let mut recordbatches = recordbatches_from_flight_message_stream(
"test-peer".to_string(),
stream::iter(vec![
Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())),
Ok(FlightMessage::Metrics(test_metrics_json())),
]),
)
.await
.unwrap();
let batch = recordbatches.next().await.unwrap().unwrap();
assert_eq!(batch.num_rows(), 1);
let metrics = recordbatches.metrics().unwrap();
assert_eq!(metrics.elapsed_compute, 7);
assert!(recordbatches.next().await.is_none());
}
#[tokio::test]
async fn test_record_batch_stream_preserves_peeked_record_batch() {
let schema = test_schema();
let first_batch = RecordBatch::new(
schema.clone(),
vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef],
)
.unwrap();
let second_batch = RecordBatch::new(
schema.clone(),
vec![Arc::new(Int32Vector::from_slice([2])) as VectorRef],
)
.unwrap();
let mut recordbatches = recordbatches_from_flight_message_stream(
"test-peer".to_string(),
stream::iter(vec![
Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
Ok(FlightMessage::RecordBatch(
first_batch.into_df_record_batch(),
)),
Ok(FlightMessage::RecordBatch(
second_batch.into_df_record_batch(),
)),
]),
)
.await
.unwrap();
let first_batch = recordbatches.next().await.unwrap().unwrap();
let second_batch = recordbatches.next().await.unwrap().unwrap();
assert_eq!(
first_batch
.column(0)
.as_any()
.downcast_ref::<Int32Array>()
.unwrap()
.value(0),
1
);
assert_eq!(
second_batch
.column(0)
.as_any()
.downcast_ref::<Int32Array>()
.unwrap()
.value(0),
2
);
assert!(recordbatches.next().await.is_none());
}
#[tokio::test]
async fn test_record_batch_stream_exposes_error_after_pre_batch_metrics() {
let schema = test_schema();
src/cmd/Cargo.toml
+1
View File
@@ -118,6 +118,7 @@ tonic.workspace = true
tracing-appender.workspace = true
[target.'cfg(unix)'.dependencies]
daemonize = "0.5"
pprof = { version = "0.14", features = [
"flamegraph",
] }
src/cmd/src/bin/greptime.rs
+51 -7
View File
@@ -17,6 +17,8 @@
use clap::{Parser, Subcommand};
use cmd::datanode::builder::InstanceBuilder;
#[cfg(unix)]
use cmd::error::DaemonizeSnafu;
use cmd::error::{InitTlsProviderSnafu, Result};
use cmd::options::GlobalOptions;
use cmd::{App, cli, datanode, flownode, frontend, metasrv, standalone, user};
@@ -73,20 +75,24 @@ static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
#[cfg(debug_assertions)]
fn main() -> Result<()> {
use snafu::ResultExt;
let cli = Command::parse();
maybe_daemonize(&cli)?;
// Set the stack size to 8MB for the thread so it wouldn't overflow on large stack usage in debug mode
// see https://github.com/GreptimeTeam/greptimedb/pull/4317
// and https://github.com/rust-lang/rust/issues/34283
std::thread::Builder::new()
.name("main_spawn".to_string())
.stack_size(8 * 1024 * 1024)
.spawn(|| {
.spawn(move || {
{
tokio::runtime::Builder::new_multi_thread()
.thread_stack_size(8 * 1024 * 1024)
.enable_all()
.build()
.expect("Failed building the Runtime")
.block_on(main_body())
.block_on(main_body(cli))
}
})
.context(cmd::error::SpawnThreadSnafu)?
@@ -95,15 +101,53 @@ fn main() -> Result<()> {
}
#[cfg(not(debug_assertions))]
#[tokio::main]
async fn main() -> Result<()> {
main_body().await
fn main() -> Result<()> {
let cli = Command::parse();
maybe_daemonize(&cli)?;
tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.expect("Failed building the Runtime")
.block_on(main_body(cli))
}
async fn main_body() -> Result<()> {
async fn main_body(cli: Command) -> Result<()> {
setup_human_panic();
install_default_crypto_provider().map_err(|msg| InitTlsProviderSnafu { msg }.build())?;
start(Command::parse()).await
start(cli).await
}
/// If `--daemon` was passed for `standalone start`, fork into the background
/// before any Tokio runtime is built. On non-Unix platforms the `--daemon`
/// flag does not exist, so this is a no-op.
fn maybe_daemonize(cli: &Command) -> Result<()> {
let daemon = matches!(&cli.subcmd, SubCommand::Standalone(cmd) if cmd.is_daemon());
if !daemon {
return Ok(());
}
#[cfg(unix)]
{
let cwd = std::env::current_dir().map_err(|e| {
DaemonizeSnafu {
msg: format!("failed to read current working directory: {e}"),
}
.build()
})?;
daemonize::Daemonize::new()
.working_directory(cwd)
.start()
.map_err(|e| {
DaemonizeSnafu {
msg: format!("failed to daemonize: {e}"),
}
.build()
})?;
}
Ok(())
}
async fn start(cli: Command) -> Result<()> {
src/cmd/src/error.rs
+9 -1
View File
@@ -162,6 +162,13 @@ pub enum Error {
location: Location,
},
#[snafu(display("Failed to daemonize: {msg}"))]
Daemonize {
msg: String,
#[snafu(implicit)]
location: Location,
},
#[snafu(display("Unsupported selector type: {}", selector_type))]
UnsupportedSelectorType {
selector_type: String,
@@ -361,7 +368,8 @@ impl ErrorExt for Error {
Error::SerdeJson { .. }
| Error::FileIo { .. }
| Error::SpawnThread { .. }
| Error::InitTlsProvider { .. } => StatusCode::Unexpected,
| Error::InitTlsProvider { .. }
| Error::Daemonize { .. } => StatusCode::Unexpected,
Error::Other { source, .. } => source.status_code(),
src/cmd/src/frontend.rs
+196 -19
View File
@@ -13,7 +13,9 @@
// limitations under the License.
use std::fmt::Debug;
use std::future::Future;
use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
@@ -61,6 +63,7 @@ use crate::options::{GlobalOptions, GreptimeOptions};
use crate::{App, create_resource_limit_metrics, log_versions, maybe_activate_heap_profile};
type FrontendOptions = GreptimeOptions<frontend::frontend::FrontendOptions>;
type HeartbeatExtensionSetupFuture<'a> = Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>;
pub struct Instance {
frontend: Frontend,
@@ -117,7 +120,35 @@ pub struct Command {
impl Command {
pub async fn build(&self, opts: FrontendOptions) -> Result<Instance> {
self.subcmd.build(opts).await
self.build_with_heartbeat_extensions(opts, FrontendHeartbeatExtensions::default())
.await
}
/// Builds a frontend with pre-registered heartbeat extensions.
///
/// Register extensions before calling this method. The supplied registry is installed before
/// normal plugin setup. After plugin heartbeat setup completes, the registry is frozen before
/// heartbeat handlers and the heartbeat task are built, so later registrations are rejected
/// and every heartbeat consumer observes the same extension membership.
pub async fn build_with_heartbeat_extensions(
&self,
opts: FrontendOptions,
heartbeat_extensions: FrontendHeartbeatExtensions,
) -> Result<Instance> {
let plugins = plugins_with_heartbeat_extensions(heartbeat_extensions);
self.forward_build_with_plugins(opts, plugins, |command, opts, plugins| {
command.build_with_plugins(opts, plugins)
})
.await
}
fn forward_build_with_plugins<'a, T>(
&'a self,
opts: FrontendOptions,
plugins: Plugins,
build: impl FnOnce(&'a StartCommand, FrontendOptions, Plugins) -> T,
) -> T {
self.subcmd.forward_build_with_plugins(opts, plugins, build)
}
pub fn load_options(&self, global_options: &GlobalOptions) -> Result<FrontendOptions> {
@@ -131,9 +162,14 @@ pub enum SubCommand {
}
impl SubCommand {
async fn build(&self, opts: FrontendOptions) -> Result<Instance> {
fn forward_build_with_plugins<'a, T>(
&'a self,
opts: FrontendOptions,
plugins: Plugins,
build: impl FnOnce(&'a StartCommand, FrontendOptions, Plugins) -> T,
) -> T {
match self {
SubCommand::Start(cmd) => cmd.build(opts).await,
SubCommand::Start(cmd) => build(cmd, opts, plugins),
}
}
@@ -329,7 +365,11 @@ impl StartCommand {
Ok(())
}
async fn build(&self, opts: FrontendOptions) -> Result<Instance> {
async fn build_with_plugins(
&self,
opts: FrontendOptions,
plugins: Plugins,
) -> Result<Instance> {
let guard = common_telemetry::init_global_logging(
APP_NAME,
&opts.component.logging,
@@ -380,15 +420,8 @@ impl StartCommand {
.await
.context(error::MetaClientInitSnafu)?;
let mut plugins = Plugins::new();
plugins::setup_frontend_plugins_pre_build(
&mut plugins,
&plugin_opts,
&opts,
Some(&meta_config),
)
.await
.context(error::StartFrontendSnafu)?;
let mut plugins =
prepare_frontend_plugins(plugins, &plugin_opts, &opts, Some(&meta_config)).await?;
// now initialize the meta_client with plugins
let meta_client = meta_client::create_meta_client(
@@ -434,6 +467,8 @@ impl StartCommand {
// Some queries are expected to take long time.
let mut channel_config = opts.datanode.client.channel_config();
channel_config.timeout = None;
// Source Flight streams and sink unary responses share pooled connections.
channel_config.http2_adaptive_window = Some(true);
if opts.grpc.flight_compression.transport_compression() {
channel_config.accept_compression = true;
channel_config.send_compression = true;
@@ -497,12 +532,20 @@ impl StartCommand {
let instance = Arc::new(instance);
plugins::setup_frontend_heartbeat_extensions(&mut plugins, &plugin_opts, &instance)
.await
.context(error::StartFrontendSnafu)?;
let heartbeat_extensions = plugins
.get::<FrontendHeartbeatExtensions>()
.unwrap_or_default();
let heartbeat_instance = instance.clone();
let heartbeat_extensions =
setup_and_freeze_frontend_heartbeat_extensions(&mut plugins, move |plugins| {
Box::pin(async move {
plugins::setup_frontend_heartbeat_extensions(
plugins,
&plugin_opts,
&heartbeat_instance,
)
.await
.context(error::StartFrontendSnafu)
})
})
.await?;
let heartbeat_task = Some(create_heartbeat_task_with_extensions(
&opts,
meta_client,
@@ -524,6 +567,36 @@ impl StartCommand {
}
}
async fn prepare_frontend_plugins(
mut plugins: Plugins,
plugin_opts: &[PluginOptions],
opts: &frontend::frontend::FrontendOptions,
meta_config: Option<&[PluginOptions]>,
) -> Result<Plugins> {
plugins::setup_frontend_plugins_pre_build(&mut plugins, plugin_opts, opts, meta_config)
.await
.context(error::StartFrontendSnafu)?;
Ok(plugins)
}
fn plugins_with_heartbeat_extensions(heartbeat_extensions: FrontendHeartbeatExtensions) -> Plugins {
let plugins = Plugins::new();
plugins.insert(heartbeat_extensions);
plugins
}
async fn setup_and_freeze_frontend_heartbeat_extensions(
plugins: &mut Plugins,
setup: impl for<'a> FnOnce(&'a mut Plugins) -> HeartbeatExtensionSetupFuture<'a>,
) -> Result<FrontendHeartbeatExtensions> {
setup(plugins).await?;
let extensions = plugins
.get::<FrontendHeartbeatExtensions>()
.unwrap_or_default();
extensions.freeze();
Ok(extensions)
}
pub fn create_heartbeat_task(
options: &frontend::frontend::FrontendOptions,
meta_client: MetaClientRef,
@@ -563,6 +636,7 @@ fn create_heartbeat_task_with_extensions(
#[cfg(test)]
mod tests {
use std::io::Write;
use std::sync::Arc;
use std::time::Duration;
use auth::{Identity, Password, UserProviderRef};
@@ -576,6 +650,109 @@ mod tests {
use super::*;
use crate::options::GlobalOptions;
#[derive(Debug)]
struct TestExtension(&'static str);
#[async_trait]
impl frontend::heartbeat::FrontendHeartbeatExtension for TestExtension {
fn name(&self) -> &'static str {
self.0
}
}
#[test]
fn test_command_forwards_heartbeat_extensions_to_start_build() {
let command = Command {
subcmd: SubCommand::Start(StartCommand::default()),
};
let extensions = FrontendHeartbeatExtensions::default();
let shared_extensions = extensions.clone();
let forwarded_plugins = command.forward_build_with_plugins(
FrontendOptions::default(),
plugins_with_heartbeat_extensions(extensions),
|_, _, plugins| plugins,
);
let forwarded_extensions = forwarded_plugins
.get::<FrontendHeartbeatExtensions>()
.unwrap();
assert_registry_identity(&shared_extensions, &forwarded_extensions);
}
#[tokio::test]
async fn test_prefilled_heartbeat_extensions_survive_pre_build_setup() {
let extensions = FrontendHeartbeatExtensions::default();
let shared_extensions = extensions.clone();
let options = frontend::frontend::FrontendOptions {
meta_client: Some(MetaClientOptions::default()),
..Default::default()
};
let plugins = prepare_frontend_plugins(
plugins_with_heartbeat_extensions(extensions),
&[],
&options,
None,
)
.await
.unwrap();
let setup_extensions = plugins.get_or_insert(FrontendHeartbeatExtensions::default);
assert_registry_identity(&shared_extensions, &setup_extensions);
}
#[tokio::test]
async fn test_heartbeat_extension_setup_completes_before_freeze() {
let extensions = FrontendHeartbeatExtensions::default();
assert_eq!(
extensions.try_register(Arc::new(TestExtension("caller"))),
Ok(())
);
let mut plugins = plugins_with_heartbeat_extensions(extensions.clone());
let finalized = setup_and_freeze_frontend_heartbeat_extensions(&mut plugins, |plugins| {
Box::pin(async move {
let setup_extensions = plugins.get_or_insert(FrontendHeartbeatExtensions::default);
assert_eq!(
setup_extensions.try_register(Arc::new(TestExtension("plugin"))),
Ok(())
);
Ok(())
})
})
.await
.unwrap();
assert_eq!(
finalized
.extensions()
.iter()
.map(|extension| extension.name())
.collect::<Vec<_>>(),
["caller", "plugin"]
);
assert_eq!(
extensions.try_register(Arc::new(TestExtension("late"))),
Err(frontend::heartbeat::RegistrationError::Frozen)
);
assert_eq!(finalized.len(), 2);
}
fn assert_registry_identity(
expected: &FrontendHeartbeatExtensions,
actual: &FrontendHeartbeatExtensions,
) {
let extension = Arc::new(TestExtension("cmd-test-extension"));
assert!(expected.register(extension.clone()));
let registered = actual.extensions();
assert_eq!(registered.len(), 1);
assert!(Arc::ptr_eq(
&(extension as Arc<dyn frontend::heartbeat::FrontendHeartbeatExtension>),
&registered[0]
));
}
#[test]
fn test_try_from_start_command() {
let command = StartCommand {
src/cmd/src/standalone.rs
+39
View File
@@ -152,6 +152,11 @@ impl Command {
) -> Result<GreptimeOptions<StandaloneOptions>> {
self.subcmd.load_options(global_options)
}
/// Whether the `standalone start` command requested daemonization.
pub fn is_daemon(&self) -> bool {
self.subcmd.is_daemon()
}
}
#[derive(Parser)]
@@ -174,6 +179,12 @@ impl SubCommand {
SubCommand::Start(cmd) => cmd.load_options(global_options),
}
}
fn is_daemon(&self) -> bool {
match self {
SubCommand::Start(cmd) => cmd.is_daemon(),
}
}
}
pub struct Instance {
@@ -291,9 +302,24 @@ pub struct StartCommand {
/// The working home directory of this standalone instance.
#[clap(long)]
data_home: Option<String>,
/// Run in the background as a daemon.
#[cfg(unix)]
#[clap(short, long)]
daemon: bool,
}
impl StartCommand {
/// Whether the `standalone start` command requested daemonization.
#[cfg(unix)]
pub(crate) fn is_daemon(&self) -> bool {
self.daemon
}
#[cfg(not(unix))]
pub(crate) fn is_daemon(&self) -> bool {
false
}
/// Load the GreptimeDB options from various sources (command line, config file or env).
pub fn load_options(
&self,
@@ -1356,6 +1382,19 @@ mod tests {
assert_eq!(command.grpc_bind_addr.as_deref(), Some("127.0.0.1:34001"));
}
#[cfg(unix)]
#[test]
fn test_parse_daemon_flag() {
let command = StartCommand::try_parse_from(["standalone", "--daemon"]).unwrap();
assert!(command.is_daemon());
let command = StartCommand::try_parse_from(["standalone", "-d"]).unwrap();
assert!(command.is_daemon());
let command = StartCommand::try_parse_from(["standalone"]).unwrap();
assert!(!command.is_daemon());
}
#[test]
fn test_help_uses_grpc_option_names() {
let mut cmd = StartCommand::command();
src/cmd/tests/load_config_test.rs
+2 -2
View File
@@ -113,7 +113,7 @@ fn test_load_datanode_example_config() {
},
region_engine: vec![
RegionEngineConfig::Mito(MitoConfig {
auto_flush_interval: Duration::from_secs(3600),
auto_flush_interval: Duration::from_secs(10 * 60),
default_region_write_buffer_size: ReadableSize::mb(0),
write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)),
scan_memory_limit: MemoryLimit::Unlimited,
@@ -354,7 +354,7 @@ fn test_load_standalone_example_config() {
}),
region_engine: vec![
RegionEngineConfig::Mito(MitoConfig {
auto_flush_interval: Duration::from_secs(3600),
auto_flush_interval: Duration::from_secs(10 * 60),
default_region_write_buffer_size: ReadableSize::mb(0),
write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)),
scan_memory_limit: MemoryLimit::Unlimited,
src/cmd/tests/standalone_daemon.rs
+109
View File
@@ -0,0 +1,109 @@
// 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.
#![cfg(unix)]
use std::net::TcpStream;
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};
/// Stops the orphaned background daemon when the test finishes (including on
/// panic). The daemon's command line contains the unique `--data-home` path, so
/// we match on it to locate the process.
struct DaemonGuard(String);
impl Drop for DaemonGuard {
fn drop(&mut self) {
let _ = Command::new("pkill")
.args(["-f", &self.0])
.stdout(Stdio::null())
.stderr(Stdio::null())
.status();
}
}
#[test]
#[cfg(unix)]
fn test_standalone_daemon_mode_detaches_from_shell() {
let bin = env!("CARGO_BIN_EXE_greptime");
let data_home = tempfile::tempdir().unwrap();
let data_home_path = data_home.path().to_str().unwrap().to_string();
let http_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port());
let grpc_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port());
let mysql_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port());
let postgres_addr = format!("127.0.0.1:{}", common_test_util::ports::get_port());
let mut child = Command::new(bin)
.args([
"standalone",
"start",
"--daemon",
"--data-home",
&data_home_path,
"--http-addr",
&http_addr,
"--grpc-bind-addr",
&grpc_addr,
"--mysql-addr",
&mysql_addr,
"--postgres-addr",
&postgres_addr,
])
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()
.unwrap();
let _guard = DaemonGuard(data_home_path);
// When daemonization works, the parent forks and exits right away. When it
// does not, the process stays in the foreground and never returns.
let detach_timeout = Duration::from_secs(30);
let started = Instant::now();
loop {
match child.try_wait().unwrap() {
Some(status) => {
assert!(
status.success(),
"daemon start should exit successfully, got {status:?}"
);
break;
}
None => {
assert!(
started.elapsed() < detach_timeout,
"daemon mode should detach and return quickly, but it blocked for over {detach_timeout:?}"
);
std::thread::sleep(Duration::from_millis(200));
}
}
}
// A successful detach is not enough: a broken implementation could fork,
// have the parent exit, then crash the child. Poll the HTTP address to
// confirm the daemon actually came up and is listening.
let ready_timeout = Duration::from_secs(30);
let started = Instant::now();
loop {
if TcpStream::connect(&http_addr).is_ok() {
break;
}
assert!(
started.elapsed() < ready_timeout,
"daemon should come up and listen on {http_addr}, but it was unreachable for {ready_timeout:?}"
);
std::thread::sleep(Duration::from_millis(200));
}
}
src/common/catalog/src/consts.rs
+5
View File
@@ -220,6 +220,11 @@ pub const SEMANTIC_RELATIONSHIPS_TABLE_NAME: &str = "semantic_relationships";
/// `semantic_relationships`; see [`is_ddl_reserved_table`] for its lifecycle.
pub const SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME: &str = "semantic_relationships_declared";
/// Width of the window the graph derivation bins observations into. Sources
/// synthesizing observations must land a row in every window they describe,
/// so this is shared rather than restated per crate.
pub const SEMANTIC_GRAPH_WINDOW_NANOS: i64 = 60 * 1_000_000_000;
// Column names of the graph tables: `catalog` exposes these schemas and the
// read-time plans in `operator` must project exactly them.
pub const OBSERVED_AT_COLUMN: &str = "observed_at";
src/common/meta/Cargo.toml
+1 -1
View File
@@ -89,7 +89,7 @@ strum.workspace = true
table = { workspace = true, features = ["testing"] }
tokio.workspace = true
tokio-postgres = { workspace = true, optional = true }
tokio-postgres-rustls = { version = "0.12", optional = true }
tokio-postgres-rustls = { version = "0.14", optional = true }
tonic.workspace = true
tracing.workspace = true
typetag.workspace = true
src/common/meta/src/datanode.rs
+2
View File
@@ -31,6 +31,8 @@ use crate::heartbeat::utils::get_datanode_workloads;
const DATANODE_STAT_PREFIX: &str = "__meta_datanode_stat";
pub const REGION_STATISTIC_KEY: &str = "__region_statistic";
/// The private table name for persisted region statistics.
pub const REGION_STATS_HISTORY_TABLE_NAME: &str = "region_statistics_history";
lazy_static! {
pub(crate) static ref DATANODE_LEASE_KEY_PATTERN: Regex =
src/common/meta/src/reconciliation.rs
+1
View File
@@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
pub(crate) mod event;
pub mod manager;
pub(crate) mod reconcile_catalog;
pub(crate) mod reconcile_database;
src/common/meta/src/reconciliation/event.rs
+351
View File
@@ -0,0 +1,351 @@
// 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::any::Any;
use api::v1::value::ValueData;
use api::v1::{ColumnSchema, Row};
use common_event_recorder::Event;
use common_event_recorder::error::{Result, SerializeEventSnafu};
use common_event_recorder::event_table::{
CATALOG_NAME_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, TABLE_ID_COLUMN,
TABLE_NAME_COLUMN, column_schemas, nullable_string, nullable_value,
};
use serde::Serialize;
use snafu::ResultExt;
use store_api::storage::TableId;
use crate::reconciliation::ResolveStrategy;
/// Stable event type stored for physical table reconciliation procedures.
pub(crate) const RECONCILE_TABLE_EVENT_TYPE: &str = "reconcile_table";
const PAYLOAD_VERSION: u8 = 1;
#[derive(Debug, Clone, Copy)]
enum ReconciliationEventType {
Table,
}
impl ReconciliationEventType {
const fn as_str(self) -> &'static str {
match self {
Self::Table => RECONCILE_TABLE_EVENT_TYPE,
}
}
}
/// Nullable object locators shared by all reconciliation event types.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct ReconciliationLocator {
catalog_name: Option<String>,
schema_name: Option<String>,
table_name: Option<String>,
table_id: Option<TableId>,
physical_table_id: Option<TableId>,
}
impl ReconciliationLocator {
/// Creates a locator for a physical table with its fully qualified name and ID.
pub(crate) fn physical_table(
catalog_name: &str,
schema_name: &str,
table_name: &str,
table_id: TableId,
) -> Self {
Self {
catalog_name: Some(catalog_name.to_string()),
schema_name: Some(schema_name.to_string()),
table_name: Some(table_name.to_string()),
table_id: Some(table_id),
..Default::default()
}
}
}
#[derive(Debug, Serialize)]
#[serde(untagged)]
enum ReconciliationPayload {
TableSubmitted(TableSubmittedPayload),
TableResult(TableResultPayload),
}
#[derive(Debug, Serialize)]
struct TableSubmittedPayload {
version: u8,
resolve_strategy: &'static str,
is_subprocedure: bool,
}
#[derive(Debug, Serialize)]
struct TableResultPayload {
version: u8,
complete: bool,
metadata_state: Option<&'static str>,
resolution_strategy_applied: Option<&'static str>,
resolved_column_count: Option<usize>,
scanned_region_count: usize,
updated_region_count: usize,
table_info_updated: bool,
last_completed_phase: Option<&'static str>,
}
/// Event representation shared by reconciliation procedures.
#[derive(Debug)]
pub(crate) struct ReconciliationEvent {
event_type: ReconciliationEventType,
locator: ReconciliationLocator,
payload: Option<ReconciliationPayload>,
}
impl ReconciliationEvent {
/// Builds the bounded intent event emitted when table reconciliation is submitted.
pub(crate) fn table_submitted(
locator: ReconciliationLocator,
resolve_strategy: ResolveStrategy,
is_subprocedure: bool,
) -> Self {
Self {
event_type: ReconciliationEventType::Table,
locator,
payload: Some(ReconciliationPayload::TableSubmitted(
TableSubmittedPayload {
version: PAYLOAD_VERSION,
resolve_strategy: resolve_strategy_name(resolve_strategy),
is_subprocedure,
},
)),
}
}
/// Builds a terminal event from the bounded reconciliation result summary.
#[allow(clippy::too_many_arguments)]
pub(crate) fn table_result(
locator: ReconciliationLocator,
complete: bool,
metadata_state: Option<&'static str>,
resolution_strategy_applied: Option<ResolveStrategy>,
resolved_column_count: Option<usize>,
scanned_region_count: usize,
updated_region_count: usize,
table_info_updated: bool,
last_completed_phase: Option<&'static str>,
) -> Self {
Self {
event_type: ReconciliationEventType::Table,
locator,
payload: Some(ReconciliationPayload::TableResult(TableResultPayload {
version: PAYLOAD_VERSION,
complete,
metadata_state,
resolution_strategy_applied: resolution_strategy_applied.map(resolve_strategy_name),
resolved_column_count,
scanned_region_count,
updated_region_count,
table_info_updated,
last_completed_phase,
})),
}
}
/// Builds a table lifecycle event whose reconciliation payload is null.
pub(crate) fn table_lifecycle(locator: ReconciliationLocator) -> Self {
Self::lifecycle(ReconciliationEventType::Table, locator)
}
fn lifecycle(event_type: ReconciliationEventType, locator: ReconciliationLocator) -> Self {
Self {
event_type,
locator,
payload: None,
}
}
fn schema() -> Vec<ColumnSchema> {
column_schemas([
&CATALOG_NAME_COLUMN,
&SCHEMA_NAME_COLUMN,
&TABLE_NAME_COLUMN,
&TABLE_ID_COLUMN,
&PHYSICAL_TABLE_ID_COLUMN,
])
}
}
impl Event for ReconciliationEvent {
fn event_type(&self) -> &str {
self.event_type.as_str()
}
fn json_payload(&self) -> Result<serde_json::Value> {
match &self.payload {
Some(payload) => serde_json::to_value(payload).context(SerializeEventSnafu),
None => Ok(serde_json::Value::Null),
}
}
fn extra_schema(&self) -> Vec<ColumnSchema> {
Self::schema()
}
fn extra_rows(&self) -> Result<Vec<Row>> {
Ok(vec![Row {
values: vec![
nullable_string(self.locator.catalog_name.as_deref()),
nullable_string(self.locator.schema_name.as_deref()),
nullable_string(self.locator.table_name.as_deref()),
nullable_table_id(self.locator.table_id),
nullable_table_id(self.locator.physical_table_id),
],
}])
}
fn as_any(&self) -> &dyn Any {
self
}
}
fn resolve_strategy_name(strategy: ResolveStrategy) -> &'static str {
match strategy {
ResolveStrategy::UseLatest => "use_latest",
ResolveStrategy::UseMetasrv => "use_metasrv",
ResolveStrategy::AbortOnConflict => "abort_on_conflict",
}
}
fn nullable_table_id(value: Option<TableId>) -> api::v1::Value {
nullable_value(value.map(ValueData::U32Value))
}
#[cfg(test)]
mod tests {
use api::v1::value::ValueData;
use api::v1::{ColumnDataType, Row, SemanticType, Value};
use common_event_recorder::Event;
use serde_json::json;
use super::*;
#[test]
fn reconciliation_events_use_the_shared_locator_contract() {
let table = ReconciliationEvent::table_lifecycle(ReconciliationLocator::physical_table(
"greptime", "public", "metrics", 42,
));
assert_eq!(table.event_type(), RECONCILE_TABLE_EVENT_TYPE);
assert_eq!(
table
.extra_schema()
.into_iter()
.map(|column| {
(
column.column_name,
ColumnDataType::try_from(column.datatype).unwrap(),
SemanticType::try_from(column.semantic_type).unwrap(),
)
})
.collect::<Vec<_>>(),
vec![
(
"catalog_name".to_string(),
ColumnDataType::String,
SemanticType::Field,
),
(
"schema_name".to_string(),
ColumnDataType::String,
SemanticType::Field,
),
(
"table_name".to_string(),
ColumnDataType::String,
SemanticType::Field,
),
(
"table_id".to_string(),
ColumnDataType::Uint32,
SemanticType::Field,
),
(
"physical_table_id".to_string(),
ColumnDataType::Uint32,
SemanticType::Field,
),
]
);
assert_eq!(
table.extra_rows().unwrap(),
vec![Row {
values: vec![
ValueData::StringValue("greptime".to_string()).into(),
ValueData::StringValue("public".to_string()).into(),
ValueData::StringValue("metrics".to_string()).into(),
ValueData::U32Value(42).into(),
Value::default(),
],
}]
);
assert_eq!(table.json_payload().unwrap(), serde_json::Value::Null);
}
#[test]
fn submitted_payloads_are_versioned_and_use_stable_strategy_names() {
for (strategy, expected) in [
(ResolveStrategy::UseLatest, "use_latest"),
(ResolveStrategy::UseMetasrv, "use_metasrv"),
(ResolveStrategy::AbortOnConflict, "abort_on_conflict"),
] {
let table = ReconciliationEvent::table_submitted(
ReconciliationLocator::physical_table("greptime", "public", "metrics", 42),
strategy,
true,
);
assert_eq!(
table.json_payload().unwrap(),
json!({
"version": 1,
"resolve_strategy": expected,
"is_subprocedure": true,
})
);
}
}
#[test]
fn terminal_payloads_distinguish_complete_and_partial_results() {
let table = ReconciliationEvent::table_result(
ReconciliationLocator::physical_table("greptime", "public", "metrics", 42),
false,
Some("inconsistent"),
Some(ResolveStrategy::UseMetasrv),
Some(4),
3,
2,
true,
Some("update_table_info"),
);
assert_eq!(
table.json_payload().unwrap(),
json!({
"version": 1,
"complete": false,
"metadata_state": "inconsistent",
"resolution_strategy_applied": "use_metasrv",
"resolved_column_count": 4,
"scanned_region_count": 3,
"updated_region_count": 2,
"table_info_updated": true,
"last_completed_phase": "update_table_info",
})
);
}
}
src/common/meta/src/reconciliation/reconcile_table.rs
+540 -4
View File
@@ -18,13 +18,12 @@ pub(crate) mod reconciliation_start;
pub(crate) mod resolve_column_metadata;
pub(crate) mod update_table_info;
use std::any::Any;
use std::fmt::Debug;
use common_procedure::error::{FromJsonSnafu, ToJsonSnafu};
use common_procedure::{
Context as ProcedureContext, Error as ProcedureError, LockKey, Procedure,
Result as ProcedureResult, Status,
Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey,
Procedure, Result as ProcedureResult, Status,
};
use serde::{Deserialize, Serialize};
use snafu::ResultExt;
@@ -42,6 +41,9 @@ use crate::key::{DeserializedValueWithBytes, TableMetadataManagerRef};
use crate::lock_key::{CatalogLock, SchemaLock, TableNameLock};
use crate::metrics;
use crate::node_manager::NodeManagerRef;
use crate::reconciliation::event::{
RECONCILE_TABLE_EVENT_TYPE, ReconciliationEvent, ReconciliationLocator,
};
use crate::reconciliation::reconcile_table::reconciliation_start::ReconciliationStart;
use crate::reconciliation::reconcile_table::resolve_column_metadata::ResolveStrategy;
use crate::reconciliation::utils::{
@@ -56,6 +58,105 @@ pub struct ReconcileTableContext {
pub volatile_ctx: VolatileContext,
}
#[derive(Debug, Clone, Copy)]
enum TablePhase {
Start,
ResolveColumnMetadata,
ReconcileRegions,
UpdateTableInfo,
}
impl TablePhase {
const fn as_event_value(self) -> &'static str {
match self {
Self::Start => "start",
Self::ResolveColumnMetadata => "resolve_column_metadata",
Self::ReconcileRegions => "reconcile_regions",
Self::UpdateTableInfo => "update_table_info",
}
}
}
#[derive(Debug, Clone, Copy)]
enum TableMetadataState {
Consistent,
Inconsistent,
}
impl TableMetadataState {
const fn as_event_value(self) -> &'static str {
match self {
Self::Consistent => "consistent",
Self::Inconsistent => "inconsistent",
}
}
}
#[derive(Debug, Default)]
struct ReconcileTableResultSummary {
metadata_state: Option<TableMetadataState>,
resolution_strategy_applied: Option<ResolveStrategy>,
resolved_column_count: Option<usize>,
scanned_region_count: usize,
updated_region_count: usize,
table_info_updated: bool,
last_completed_phase: Option<TablePhase>,
}
impl ReconcileTableResultSummary {
fn record_scanned_regions(&mut self, scanned_region_count: usize) {
self.scanned_region_count = scanned_region_count;
}
fn mark_start_completed(&mut self) {
self.last_completed_phase = Some(TablePhase::Start);
}
fn record_metadata_state(&mut self, metadata_state: TableMetadataState) {
self.metadata_state = Some(metadata_state);
}
fn record_resolution_strategy(&mut self, resolve_strategy: ResolveStrategy) {
self.resolution_strategy_applied = Some(resolve_strategy);
}
fn record_resolved_columns(
&mut self,
metadata_state: TableMetadataState,
resolution_strategy_applied: Option<ResolveStrategy>,
resolved_column_count: Option<usize>,
) {
self.metadata_state = Some(metadata_state);
self.resolution_strategy_applied = resolution_strategy_applied;
self.resolved_column_count = resolved_column_count;
self.last_completed_phase = Some(TablePhase::ResolveColumnMetadata);
}
fn record_updated_regions(&mut self, updated_region_count: usize) {
self.updated_region_count = updated_region_count;
}
fn mark_region_phase_completed(&mut self) {
self.last_completed_phase = Some(TablePhase::ReconcileRegions);
}
fn mark_table_info_updated(&mut self) {
self.table_info_updated = true;
}
fn mark_table_info_phase_completed(&mut self) {
self.last_completed_phase = Some(TablePhase::UpdateTableInfo);
}
fn metadata_state(&self) -> Option<&'static str> {
self.metadata_state.map(TableMetadataState::as_event_value)
}
fn last_completed_phase(&self) -> Option<&'static str> {
self.last_completed_phase.map(TablePhase::as_event_value)
}
}
impl ReconcileTableContext {
/// Creates a new [`ReconcileTableContext`] with the given [`Context`] and [`PersistentContext`].
pub fn new(ctx: Context, persistent_ctx: PersistentContext) -> Self {
@@ -147,6 +248,7 @@ impl PersistentContext {
pub(crate) struct VolatileContext {
pub(crate) table_meta: Option<TableMeta>,
pub(crate) metrics: ReconcileTableMetrics,
result_summary: ReconcileTableResultSummary,
}
pub struct ReconcileTableProcedure {
@@ -256,6 +358,56 @@ impl Procedure for ReconcileTableProcedure {
TableNameLock::new(table_ref.catalog, table_ref.schema, table_ref.table).into(),
])
}
fn event(&self, ctx: &EventContext<'_>) -> Option<Box<dyn common_event_recorder::Event>> {
if !ctx.event_type_filter.allows(RECONCILE_TABLE_EVENT_TYPE) {
return None;
}
let persistent_ctx = &self.context.persistent_ctx;
let table_name = &persistent_ctx.table_name;
let locator = ReconciliationLocator::physical_table(
&table_name.catalog_name,
&table_name.schema_name,
&table_name.table_name,
persistent_ctx.table_id,
);
let event = match ctx.trigger {
EventTrigger::Submitted => ReconciliationEvent::table_submitted(
locator,
persistent_ctx.resolve_strategy,
persistent_ctx.is_subprocedure,
),
EventTrigger::Succeeded => {
Self::result_event(locator, &self.context.volatile_ctx.result_summary, true)
}
EventTrigger::Failed | EventTrigger::Poisoned => {
Self::result_event(locator, &self.context.volatile_ctx.result_summary, false)
}
_ => ReconciliationEvent::table_lifecycle(locator),
};
Some(Box::new(event))
}
}
impl ReconcileTableProcedure {
fn result_event(
locator: ReconciliationLocator,
summary: &ReconcileTableResultSummary,
complete: bool,
) -> ReconciliationEvent {
ReconciliationEvent::table_result(
locator,
complete,
summary.metadata_state(),
summary.resolution_strategy_applied,
summary.resolved_column_count,
summary.scanned_region_count,
summary.updated_region_count,
summary.table_info_updated,
summary.last_completed_phase(),
)
}
}
#[async_trait::async_trait]
@@ -272,6 +424,390 @@ pub(crate) trait State: Sync + Send + Debug {
ctx: &mut ReconcileTableContext,
procedure_ctx: &ProcedureContext,
) -> Result<(Box<dyn State>, Status)>;
}
fn as_any(&self) -> &dyn Any;
#[cfg(test)]
mod tests {
use std::sync::Arc;
use common_event_recorder::{EventTypeFilter, EventTypeFilterRef};
use common_procedure::{
ChildSubmissionOutcome, EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState,
RetryPhase,
};
use common_procedure_test::MockContextProvider;
use serde_json::{Value, json};
use store_api::storage::RegionId;
use super::*;
use crate::ddl::test_util::datanode_handler::PartialSuccessDatanodeHandler;
use crate::key::DeserializedValueWithBytes;
use crate::key::table_info::TableInfoValue;
use crate::key::table_route::PhysicalTableRouteValue;
use crate::key::test_utils::new_test_table_info_with_name;
use crate::peer::Peer;
use crate::reconciliation::reconcile_table::reconcile_regions::ReconcileRegions;
use crate::reconciliation::utils::build_column_metadata_from_table_info;
use crate::rpc::router::{Region, RegionRoute};
use crate::test_util::{MockDatanodeManager, new_ddl_context};
struct TableEventHarness {
procedure_id: ProcedureId,
lifecycle_state: ProcedureState,
event_type_filter: EventTypeFilterRef,
}
impl TableEventHarness {
fn all() -> Self {
Self {
procedure_id: ProcedureId::random(),
lifecycle_state: ProcedureState::Running,
event_type_filter: Arc::new(EventTypeFilter::All),
}
}
fn selected(event_types: impl IntoIterator<Item = &'static str>) -> Self {
Self {
event_type_filter: Arc::new(EventTypeFilter::Only(
event_types.into_iter().map(str::to_string).collect(),
)),
..Self::all()
}
}
fn event(
&self,
procedure: &dyn Procedure,
trigger: EventTrigger,
) -> Option<Box<dyn common_event_recorder::Event>> {
procedure.event(&EventContext {
procedure_id: self.procedure_id,
lifecycle_state: &self.lifecycle_state,
trigger,
event_type_filter: self.event_type_filter.clone(),
event_context: None,
})
}
}
#[test]
fn table_submitted_events_cover_root_and_child_intent() {
let events = TableEventHarness::all();
let root = test_procedure(false);
let child = test_procedure(true);
for (procedure, is_subprocedure) in [(&root, false), (&child, true)] {
let submitted = events.event(procedure, EventTrigger::Submitted).unwrap();
assert_eq!(submitted.event_type(), RECONCILE_TABLE_EVENT_TYPE);
assert_eq!(
submitted.json_payload().unwrap(),
json!({
"version": 1,
"resolve_strategy": "use_latest",
"is_subprocedure": is_subprocedure,
})
);
}
}
#[test]
fn table_non_terminal_lifecycle_events_have_null_payloads() {
let events = TableEventHarness::all();
let mut procedure = test_procedure(true);
procedure.context.volatile_ctx.result_summary = populated_summary();
for trigger in [
EventTrigger::Recovered,
EventTrigger::ChildSubmitted {
procedure_id: ProcedureId::random(),
outcome: ChildSubmissionOutcome::Accepted,
},
EventTrigger::Retrying {
phase: RetryPhase::Execute,
attempt: 2,
},
EventTrigger::RollingBack,
] {
assert_eq!(
events
.event(&procedure, trigger)
.unwrap()
.json_payload()
.unwrap(),
Value::Null
);
}
}
#[test]
fn table_terminal_events_report_bounded_results() {
let events = TableEventHarness::all();
let mut procedure = test_procedure(true);
procedure.context.volatile_ctx.result_summary = populated_summary();
for (trigger, complete) in [
(EventTrigger::Succeeded, true),
(EventTrigger::Failed, false),
(EventTrigger::Poisoned, false),
] {
assert_eq!(
events
.event(&procedure, trigger)
.unwrap()
.json_payload()
.unwrap(),
json!({
"version": 1,
"complete": complete,
"metadata_state": "inconsistent",
"resolution_strategy_applied": "use_latest",
"resolved_column_count": 3,
"scanned_region_count": 2,
"updated_region_count": 1,
"table_info_updated": true,
"last_completed_phase": "update_table_info",
})
);
}
}
#[test]
fn table_event_filtering_uses_the_reconciliation_event_type() {
let procedure = test_procedure(false);
assert!(
TableEventHarness::selected([RECONCILE_TABLE_EVENT_TYPE])
.event(&procedure, EventTrigger::Submitted)
.is_some()
);
assert!(
TableEventHarness::selected(["create_table"])
.event(&procedure, EventTrigger::Submitted)
.is_none()
);
assert!(
TableEventHarness::selected([])
.event(&procedure, EventTrigger::Submitted)
.is_none()
);
}
#[test]
fn table_result_summary_is_not_persisted() {
let events = TableEventHarness::all();
let mut procedure = test_procedure(false);
procedure.context.volatile_ctx.result_summary = populated_summary();
let value: Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap();
assert!(value["persistent_ctx"].get("result_summary").is_none());
let loaded = ReconcileTableProcedure::from_json(
test_context(),
&serde_json::to_string(&value).unwrap(),
)
.unwrap();
assert_eq!(
events
.event(&loaded, EventTrigger::Succeeded)
.unwrap()
.json_payload()
.unwrap(),
json!({
"version": 1,
"complete": true,
"metadata_state": null,
"resolution_strategy_applied": null,
"resolved_column_count": null,
"scanned_region_count": 0,
"updated_region_count": 0,
"table_info_updated": false,
"last_completed_phase": null,
})
);
}
#[test]
fn table_partial_summary_preserves_completed_work_without_advancing_phases() {
let events = TableEventHarness::all();
let mut procedure = test_procedure(false);
let summary = &mut procedure.context.volatile_ctx.result_summary;
summary.record_scanned_regions(2);
assert_eq!(
events
.event(&procedure, EventTrigger::Failed)
.unwrap()
.json_payload()
.unwrap(),
json!({
"version": 1,
"complete": false,
"metadata_state": null,
"resolution_strategy_applied": null,
"resolved_column_count": null,
"scanned_region_count": 2,
"updated_region_count": 0,
"table_info_updated": false,
"last_completed_phase": null,
})
);
let summary = &mut procedure.context.volatile_ctx.result_summary;
summary.mark_start_completed();
summary.record_metadata_state(TableMetadataState::Inconsistent);
assert_eq!(
events
.event(&procedure, EventTrigger::Failed)
.unwrap()
.json_payload()
.unwrap(),
json!({
"version": 1,
"complete": false,
"metadata_state": "inconsistent",
"resolution_strategy_applied": null,
"resolved_column_count": null,
"scanned_region_count": 2,
"updated_region_count": 0,
"table_info_updated": false,
"last_completed_phase": "start",
})
);
let summary = &mut procedure.context.volatile_ctx.result_summary;
summary.record_resolved_columns(
TableMetadataState::Inconsistent,
Some(ResolveStrategy::UseLatest),
Some(3),
);
summary.record_updated_regions(1);
assert_eq!(
events
.event(&procedure, EventTrigger::Failed)
.unwrap()
.json_payload()
.unwrap(),
json!({
"version": 1,
"complete": false,
"metadata_state": "inconsistent",
"resolution_strategy_applied": "use_latest",
"resolved_column_count": 3,
"scanned_region_count": 2,
"updated_region_count": 1,
"table_info_updated": false,
"last_completed_phase": "resolve_column_metadata",
})
);
}
#[tokio::test]
async fn table_mixed_region_outcome_preserves_partial_result() {
let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new(
PartialSuccessDatanodeHandler { retryable: false },
)));
let context = Context {
node_manager: ddl_context.node_manager,
table_metadata_manager: ddl_context.table_metadata_manager,
cache_invalidator: ddl_context.cache_invalidator,
};
let mut procedure = ReconcileTableProcedure::new(
context,
42,
TableName::new("greptime", "public", "metrics"),
ResolveStrategy::UseLatest,
false,
);
let mut table_info = new_test_table_info_with_name(42, "metrics");
table_info.meta.column_ids = vec![0, 1, 2];
let name_to_ids = table_info.name_to_ids().unwrap();
let column_metadatas = build_column_metadata_from_table_info(
table_info.meta.schema.column_schemas(),
&table_info.meta.primary_key_indices,
&name_to_ids,
)
.unwrap();
let region_ids = [RegionId::new(42, 1), RegionId::new(42, 2)];
procedure.context.persistent_ctx.table_info_value = Some(
DeserializedValueWithBytes::from_inner(TableInfoValue::new(table_info)),
);
procedure.context.persistent_ctx.physical_table_route =
Some(PhysicalTableRouteValue::new(vec![
RegionRoute {
region: Region::new_test(region_ids[0]),
leader_peer: Some(Peer::empty(1)),
..Default::default()
},
RegionRoute {
region: Region::new_test(region_ids[1]),
leader_peer: Some(Peer::empty(2)),
..Default::default()
},
]));
let summary = &mut procedure.context.volatile_ctx.result_summary;
summary.record_scanned_regions(region_ids.len());
summary.mark_start_completed();
summary.record_resolved_columns(
TableMetadataState::Inconsistent,
Some(ResolveStrategy::UseLatest),
Some(column_metadatas.len()),
);
procedure.state = Box::new(ReconcileRegions::new(column_metadatas, region_ids.to_vec()));
let procedure_ctx = ProcedureContext {
procedure_id: ProcedureId::random(),
provider: Arc::new(MockContextProvider::default()),
event_context: None,
};
assert!(procedure.execute(&procedure_ctx).await.is_err());
assert_eq!(
TableEventHarness::all()
.event(&procedure, EventTrigger::Failed)
.unwrap()
.json_payload()
.unwrap(),
json!({
"version": 1,
"complete": false,
"metadata_state": "inconsistent",
"resolution_strategy_applied": "use_latest",
"resolved_column_count": 3,
"scanned_region_count": 2,
"updated_region_count": 1,
"table_info_updated": false,
"last_completed_phase": "resolve_column_metadata",
})
);
}
fn populated_summary() -> ReconcileTableResultSummary {
ReconcileTableResultSummary {
metadata_state: Some(TableMetadataState::Inconsistent),
resolution_strategy_applied: Some(ResolveStrategy::UseLatest),
resolved_column_count: Some(3),
scanned_region_count: 2,
updated_region_count: 1,
table_info_updated: true,
last_completed_phase: Some(TablePhase::UpdateTableInfo),
}
}
fn test_procedure(is_subprocedure: bool) -> ReconcileTableProcedure {
ReconcileTableProcedure::new(
test_context(),
42,
TableName::new("greptime", "public", "metrics"),
ResolveStrategy::UseLatest,
is_subprocedure,
)
}
fn test_context() -> Context {
let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new(())));
Context {
node_manager: ddl_context.node_manager,
table_metadata_manager: ddl_context.table_metadata_manager,
cache_invalidator: ddl_context.cache_invalidator,
}
}
}
src/common/meta/src/reconciliation/reconcile_table/reconcile_regions.rs
+13 -12
View File
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::collections::{HashMap, HashSet};
use api::v1::column_def::try_as_column_def;
@@ -65,7 +64,6 @@ impl State for ReconcileRegions {
let table_meta = ctx.build_table_meta(&self.column_metadatas)?;
ctx.volatile_ctx.table_meta = Some(table_meta);
let table_id = ctx.table_id();
let table_name = ctx.table_name();
let primary_keys = self
.column_metadatas
@@ -126,10 +124,12 @@ impl State for ReconcileRegions {
}
}
let mut results = future::join_all(sync_column_tsks)
.await
.into_iter()
.collect::<Result<Vec<_>>>()?;
let results = future::join_all(sync_column_tsks).await;
let updated_region_count = results.iter().filter(|result| result.is_ok()).count();
ctx.volatile_ctx
.result_summary
.record_updated_regions(updated_region_count);
let mut results = results.into_iter().collect::<Result<Vec<_>>>()?;
// Ensures all the column metadatas are the same.
let column_metadatas =
@@ -137,17 +137,22 @@ impl State for ReconcileRegions {
UnexpectedSnafu {
err_msg: format!(
"The table column metadata schemas from datanodes are not the same, table: {}, table_id: {}",
table_name,
ctx.table_name(),
table_id
),
},
)?;
ctx.volatile_ctx
.result_summary
.mark_region_phase_completed();
// Checks all column metadatas are consistent, and updates the table info if needed.
if column_metadatas != self.column_metadatas {
info!(
"Datanode column metadatas are not consistent with metasrv, updating metasrv's column metadatas, table: {}, table_id: {}",
table_name, table_id
ctx.table_name(),
table_id
);
// Safety: fetched in the above.
let table_info_value = ctx.persistent_ctx.table_info_value.clone().unwrap();
@@ -159,10 +164,6 @@ impl State for ReconcileRegions {
Ok((Box::new(ReconciliationEnd), Status::executing(false)))
}
fn as_any(&self) -> &dyn Any {
self
}
}
/// Makes an alter region request to sync columns.
src/common/meta/src/reconciliation/reconcile_table/reconciliation_end.rs
-6
View File
@@ -12,8 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use common_procedure::{Context as ProcedureContext, Status};
use common_telemetry::info;
use serde::{Deserialize, Serialize};
@@ -46,8 +44,4 @@ impl State for ReconciliationEnd {
Ok((Box::new(ReconciliationEnd), Status::done()))
}
fn as_any(&self) -> &dyn Any {
self
}
}
src/common/meta/src/reconciliation/reconcile_table/reconciliation_start.rs
+14 -11
View File
@@ -12,8 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use common_procedure::{Context as ProcedureContext, Status};
use common_telemetry::info;
use serde::{Deserialize, Serialize};
@@ -46,7 +44,6 @@ impl State for ReconciliationStart {
procedure_ctx: &ProcedureContext,
) -> Result<(Box<dyn State>, Status)> {
let table_id = ctx.table_id();
let table_name = ctx.table_name();
let (physical_table_id, physical_table_route) = ctx
.table_metadata_manager
@@ -58,14 +55,17 @@ impl State for ReconciliationStart {
error::UnexpectedSnafu {
err_msg: format!(
"Reconcile table only works for physical table, but got logical table: {}, table_id: {}",
table_name, table_id
ctx.table_name(),
table_id
),
}
);
info!(
"Reconciling table: {}, table_id: {}, procedure_id: {}",
table_name, table_id, procedure_ctx.procedure_id
ctx.table_name(),
table_id,
procedure_ctx.procedure_id
);
// TODO(weny): Repairs the table route if needed.
let region_metadata_lister = RegionMetadataLister::new(ctx.node_manager.clone());
@@ -79,6 +79,9 @@ impl State for ReconciliationStart {
.list(physical_table_id, &physical_table_route.region_routes)
.await?
};
ctx.volatile_ctx
.result_summary
.record_scanned_regions(region_metadatas.len());
ensure!(!region_metadatas.is_empty(), {
metrics::METRIC_META_RECONCILIATION_STATS
@@ -92,7 +95,8 @@ impl State for ReconciliationStart {
error::UnexpectedSnafu {
err_msg: format!(
"No region metadata found for table: {}, table_id: {}",
table_name, table_id
ctx.table_name(),
table_id
),
}
});
@@ -109,11 +113,14 @@ impl State for ReconciliationStart {
error::UnexpectedSnafu {
err_msg: format!(
"Some regions are not opened, table: {}, table_id: {}",
table_name, table_id
ctx.table_name(),
table_id
),
}
});
ctx.volatile_ctx.result_summary.mark_start_completed();
// Persist the physical table route.
// TODO(weny): refetch the physical table route if repair is needed.
ctx.persistent_ctx.physical_table_route = Some(physical_table_route);
@@ -127,8 +134,4 @@ impl State for ReconciliationStart {
Status::executing(false),
))
}
fn as_any(&self) -> &dyn Any {
self
}
}
src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs
+27 -7
View File
@@ -12,8 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use async_trait::async_trait;
use common_procedure::{Context as ProcedureContext, Status};
use common_telemetry::info;
@@ -25,7 +23,7 @@ use strum::AsRefStr;
use crate::error::{self, MissingColumnIdsSnafu, Result};
use crate::reconciliation::reconcile_table::reconcile_regions::ReconcileRegions;
use crate::reconciliation::reconcile_table::update_table_info::UpdateTableInfo;
use crate::reconciliation::reconcile_table::{ReconcileTableContext, State};
use crate::reconciliation::reconcile_table::{ReconcileTableContext, State, TableMetadataState};
use crate::reconciliation::utils::{
ResolveColumnMetadataResult, build_column_metadata_from_table_info,
check_column_metadatas_consistent, resolve_column_metadatas_with_latest,
@@ -101,6 +99,12 @@ impl State for ResolveColumnMetadata {
table_name, table_id
);
ctx.volatile_ctx.result_summary.record_resolved_columns(
TableMetadataState::Consistent,
None,
Some(column_metadatas.len()),
);
// Update metrics.
ctx.mut_metrics().resolve_column_metadata_result =
Some(ResolveColumnMetadataResult::Consistent);
@@ -110,6 +114,10 @@ impl State for ResolveColumnMetadata {
));
};
ctx.volatile_ctx
.result_summary
.record_metadata_state(TableMetadataState::Inconsistent);
match self.strategy {
ResolveStrategy::UseMetasrv => {
let table_info_value = ctx.persistent_ctx.table_info_value.as_ref().unwrap();
@@ -126,6 +134,12 @@ impl State for ResolveColumnMetadata {
let region_ids =
resolve_column_metadatas_with_metasrv(&column_metadata, &self.region_metadata)?;
ctx.volatile_ctx.result_summary.record_resolved_columns(
TableMetadataState::Inconsistent,
Some(self.strategy),
Some(column_metadata.len()),
);
// Update metrics.
let metrics = ctx.mut_metrics();
metrics.resolve_column_metadata_result =
@@ -139,6 +153,12 @@ impl State for ResolveColumnMetadata {
let (column_metadatas, region_ids) =
resolve_column_metadatas_with_latest(&self.region_metadata)?;
ctx.volatile_ctx.result_summary.record_resolved_columns(
TableMetadataState::Inconsistent,
Some(self.strategy),
Some(column_metadatas.len()),
);
// Update metrics.
let metrics = ctx.mut_metrics();
metrics.resolve_column_metadata_result =
@@ -151,6 +171,10 @@ impl State for ResolveColumnMetadata {
ResolveStrategy::AbortOnConflict => {
let table_name = table_name.to_string();
ctx.volatile_ctx
.result_summary
.record_resolution_strategy(self.strategy);
// Update metrics.
let metrics = ctx.mut_metrics();
metrics.resolve_column_metadata_result =
@@ -163,8 +187,4 @@ impl State for ResolveColumnMetadata {
}
}
}
fn as_any(&self) -> &dyn Any {
self
}
}
src/common/meta/src/reconciliation/reconcile_table/update_table_info.rs
+7 -6
View File
@@ -12,8 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use common_procedure::{Context as ProcedureContext, Status};
use common_telemetry::info;
use serde::{Deserialize, Serialize};
@@ -81,6 +79,9 @@ impl State for UpdateTableInfo {
ctx.table_name(),
ctx.table_id()
);
ctx.volatile_ctx
.result_summary
.mark_table_info_phase_completed();
return Ok((Box::new(ReconciliationEnd), Status::executing(true)));
}
@@ -98,6 +99,7 @@ impl State for UpdateTableInfo {
new_table_info,
)
.await?;
ctx.volatile_ctx.result_summary.mark_table_info_updated();
let table_ref = ctx.table_name().table_ref();
let table_id = ctx.table_id();
@@ -119,11 +121,10 @@ impl State for UpdateTableInfo {
// Update metrics.
let metrics = ctx.mut_metrics();
metrics.update_table_info = true;
ctx.volatile_ctx
.result_summary
.mark_table_info_phase_completed();
Ok((Box::new(ReconciliationEnd), Status::executing(true)))
}
fn as_any(&self) -> &dyn Any {
self
}
}
src/common/recordbatch/src/cursor.rs
+8
View File
@@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use datatypes::schema::SchemaRef;
use futures::StreamExt;
use tokio::sync::Mutex;
@@ -28,12 +29,15 @@ struct Inner {
/// A cursor on RecordBatchStream that fetches data batch by batch
pub struct RecordBatchStreamCursor {
schema: SchemaRef,
inner: Mutex<Inner>,
}
impl RecordBatchStreamCursor {
pub fn new(stream: SendableRecordBatchStream) -> RecordBatchStreamCursor {
let schema = stream.schema();
Self {
schema,
inner: Mutex::new(Inner {
stream,
current_row_index: 0,
@@ -43,6 +47,10 @@ impl RecordBatchStreamCursor {
}
}
pub fn schema(&self) -> SchemaRef {
self.schema.clone()
}
/// Take `size` of row from the `RecordBatchStream` and create a new
/// `RecordBatch` for these rows.
pub async fn take(&self, size: usize) -> Result<RecordBatch> {
src/common/time/src/timestamp.rs
+69 -6
View File
@@ -483,6 +483,31 @@ impl Timestamp {
ParseTimestampSnafu { raw: s }.fail()
}
/// Interprets a timezone-less [`NaiveDateTime`] in the given timezone and
/// returns the corresponding timestamp.
///
/// Datetimes that fall into a DST gap (a local time that does not exist)
/// are rejected, while ambiguous datetimes (from a repeated local time)
/// are resolved to the earlier instant. This policy is shared with
/// [`Timestamp::from_str`] so that the text and binary protocols interpret
/// datetimes consistently.
pub fn from_naive_datetime(
datetime: NaiveDateTime,
timezone: &Timezone,
) -> crate::error::Result<Timestamp> {
match datetime_to_utc(&datetime, timezone) {
LocalResult::Single(utc) | LocalResult::Ambiguous(utc, _) => {
Timestamp::from_chrono_datetime(utc).context(ParseTimestampSnafu {
raw: format!("{datetime} (timezone {timezone})"),
})
}
LocalResult::None => ParseTimestampSnafu {
raw: format!("{datetime} (timezone {timezone})"),
}
.fail(),
}
}
pub fn negative(mut self) -> Self {
self.value = -self.value;
self
@@ -531,12 +556,8 @@ fn naive_datetime_to_timestamp(
.context(ParseTimestampSnafu { raw: s });
};
match datetime_to_utc(&datetime, timezone) {
LocalResult::None => ParseTimestampSnafu { raw: s }.fail(),
LocalResult::Single(utc) | LocalResult::Ambiguous(utc, _) => {
Timestamp::from_chrono_datetime(utc).context(ParseTimestampSnafu { raw: s })
}
}
Timestamp::from_naive_datetime(datetime, timezone)
.map_err(|_| ParseTimestampSnafu { raw: s }.build())
}
impl From<i64> for Timestamp {
@@ -923,6 +944,48 @@ mod tests {
);
}
#[test]
fn test_from_naive_datetime() {
let datetime = NaiveDate::from_ymd_opt(2026, 8, 13)
.unwrap()
.and_hms_opt(8, 0, 0)
.unwrap();
// A fixed-offset timezone shifts the datetime by a constant amount.
let shanghai = Timezone::from_tz_string("Asia/Shanghai").unwrap();
assert_eq!(
"2026-08-13 00:00:00",
Timestamp::from_naive_datetime(datetime, &shanghai)
.unwrap()
.to_chrono_datetime()
.unwrap()
.to_string()
);
// 2026-03-08 02:30 does not exist in America/New_York (DST gap).
let new_york = Timezone::from_tz_string("America/New_York").unwrap();
let gap = NaiveDate::from_ymd_opt(2026, 3, 8)
.unwrap()
.and_hms_opt(2, 30, 0)
.unwrap();
assert!(Timestamp::from_naive_datetime(gap, &new_york).is_err());
// 2026-11-01 01:30 is ambiguous in America/New_York; picks the first
// instant (EDT, UTC-4).
let ambiguous = NaiveDate::from_ymd_opt(2026, 11, 1)
.unwrap()
.and_hms_opt(1, 30, 0)
.unwrap();
assert_eq!(
"2026-11-01 05:30:00",
Timestamp::from_naive_datetime(ambiguous, &new_york)
.unwrap()
.to_chrono_datetime()
.unwrap()
.to_string()
);
}
#[test]
fn test_to_iso8601_string() {
set_default_timezone(Some("Asia/Shanghai")).unwrap();
src/common/time/src/timezone.rs
+9
View File
@@ -108,6 +108,15 @@ impl Timezone {
}
}
/// A named zone merely sitting at +00:00 today is not UTC: it may have been
/// elsewhere at the timestamp being converted.
pub fn is_utc(&self) -> bool {
match self {
Self::Offset(offset) => offset.local_minus_utc() == 0,
Self::Named(tz) => matches!(tz, Tz::UTC),
}
}
/// Returns the number of seconds to add to convert from UTC to the local time.
pub fn local_minus_utc(&self) -> i64 {
match self {
src/datanode/src/region_server.rs
+7 -2
View File
@@ -67,7 +67,8 @@ use servers::error::{
};
use servers::grpc::FlightCompression;
use servers::grpc::flight::{
FlightCraft, FlightRecordBatchSource, FlightRecordBatchStream, TonicStream,
FlightCraft, FlightRecordBatchSource, FlightRecordBatchStream, FlightRecordBatchStreamInput,
TonicStream,
};
use servers::grpc::region_server::RegionServerHandler;
use session::context::{
@@ -990,7 +991,11 @@ impl FlightCraft for RegionServer {
};
let stream = Box::pin(FlightRecordBatchStream::new(
FlightRecordBatchSource::initializer(initializer),
FlightRecordBatchStreamInput::initializer(async move {
initializer
.await
.map(FlightRecordBatchSource::RecordBatches)
}),
tracing_context,
self.flight_compression,
query_ctx,
src/flow/src/batching_mode.rs
+3
View File
@@ -44,6 +44,8 @@ pub struct BatchingModeOptions {
/// The gRPC connection timeout
#[serde(with = "humantime_serde")]
pub grpc_conn_timeout: Duration,
#[serde(with = "humantime_serde")]
pub experimental_flight_do_get_timeout: Duration,
/// The gRPC max retry number
pub experimental_grpc_max_retries: u32,
/// Flow wait for available frontend timeout,
@@ -72,6 +74,7 @@ impl Default for BatchingModeOptions {
slow_query_threshold: Duration::from_secs(60),
experimental_min_refresh_duration: Duration::new(5, 0),
grpc_conn_timeout: Duration::from_secs(5),
experimental_flight_do_get_timeout: Duration::from_secs(10),
experimental_grpc_max_retries: 3,
experimental_frontend_scan_timeout: Duration::from_secs(30),
experimental_max_filter_num_per_query: 20,
src/flow/src/batching_mode/frontend_client.rs
+7 -6
View File
@@ -413,12 +413,12 @@ impl FrontendClient {
peer: db.peer.clone(),
});
db.database
.query_with_terminal_metrics_and_flow_extensions(
request,
&hints,
extensions,
snapshot_seqs,
)
.flight_request()
.with_hints(&hints)
.with_flow_extensions(extensions)
.with_snapshot_seqs(snapshot_seqs)
.with_timeout(batch_opts.experimental_flight_do_get_timeout)
.query_with_terminal_metrics(request)
.await
.map_err(BoxedError::new)
.context(ExternalSnafu)
@@ -623,6 +623,7 @@ mod tests {
use std::task::{Context, Poll};
use std::time::Duration;
use api::v1::query_request::Query;
use arrow_flight::flight_service_server::FlightServiceServer;
use arrow_flight::{FlightData, Ticket};
use common_query::{Output, OutputData};
src/frontend/src/heartbeat.rs
+62 -7
View File
@@ -16,8 +16,9 @@
mod tests;
use std::collections::{HashMap, HashSet};
use std::fmt::{Display, Formatter};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex as StdMutex};
use std::sync::{Arc, Mutex as StdMutex, MutexGuard, PoisonError};
use api::v1::meta::heartbeat_request::NodeWorkloads;
use api::v1::meta::{FrontendWorkloads, HeartbeatRequest, HeartbeatResponse, NodeInfo, Peer};
@@ -93,34 +94,88 @@ pub trait FrontendHeartbeatExtension: Send + Sync {
}
/// A shareable, ordered registry of frontend heartbeat extensions.
///
/// Registration is available until [`FrontendHeartbeatExtensions::freeze`] is called. After that,
/// membership is immutable so every heartbeat consumer observes the same extensions.
#[derive(Clone, Default)]
pub struct FrontendHeartbeatExtensions {
inner: Arc<StdMutex<FrontendHeartbeatExtensionsInner>>,
}
/// The reason a frontend heartbeat extension was not registered.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RegistrationError {
/// An extension with the same name is already registered.
Duplicate,
/// The registry is frozen and no longer accepts registrations.
Frozen,
}
impl Display for RegistrationError {
fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::Duplicate => formatter.write_str("heartbeat extension is already registered"),
Self::Frozen => formatter.write_str("heartbeat extension registry is frozen"),
}
}
}
impl std::error::Error for RegistrationError {}
#[derive(Default)]
struct FrontendHeartbeatExtensionsInner {
names: HashSet<String>,
extensions: Vec<Arc<dyn FrontendHeartbeatExtension>>,
frozen: bool,
}
impl FrontendHeartbeatExtensions {
fn lock(&self) -> MutexGuard<'_, FrontendHeartbeatExtensionsInner> {
self.inner.lock().unwrap_or_else(PoisonError::into_inner)
}
/// Registers an extension without replacing an existing extension of the same name.
///
/// Returns `true` for a new registration and `false` for an idempotent duplicate.
/// Returns `true` for a new registration and `false` for an idempotent duplicate or when the
/// registry is frozen. Use [`FrontendHeartbeatExtensions::try_register`] when the rejection
/// reason must not be discarded.
pub fn register(&self, extension: Arc<dyn FrontendHeartbeatExtension>) -> bool {
let mut inner = self.inner.lock().unwrap();
self.try_register(extension).is_ok()
}
/// Registers an extension and returns the reason when registration is rejected.
///
/// If the registry is already frozen, this method returns [`RegistrationError::Frozen`]
/// without calling [`FrontendHeartbeatExtension::name`]. If freezing races with name
/// evaluation, the registry state is checked again before insertion.
pub fn try_register(
&self,
extension: Arc<dyn FrontendHeartbeatExtension>,
) -> std::result::Result<(), RegistrationError> {
if self.lock().frozen {
return Err(RegistrationError::Frozen);
}
let name = extension.name().to_string();
let mut inner = self.lock();
if inner.frozen {
return Err(RegistrationError::Frozen);
}
if !inner.names.insert(name) {
return false;
return Err(RegistrationError::Duplicate);
}
inner.extensions.push(extension);
true
Ok(())
}
/// Prevents further registrations while preserving the current registration order.
pub fn freeze(&self) {
self.lock().frozen = true;
}
/// Returns the registered extensions in registration order.
pub fn extensions(&self) -> Vec<Arc<dyn FrontendHeartbeatExtension>> {
self.inner.lock().unwrap().extensions.clone()
self.lock().extensions.clone()
}
/// Returns response handlers in extension registration order.
@@ -133,7 +188,7 @@ impl FrontendHeartbeatExtensions {
/// Returns the number of registered extensions.
pub fn len(&self) -> usize {
self.inner.lock().unwrap().extensions.len()
self.lock().extensions.len()
}
/// Returns whether no extension is registered.
src/frontend/src/heartbeat/tests.rs
+166 -1
View File
@@ -14,8 +14,9 @@
use std::collections::{HashMap, VecDeque};
use std::future::pending;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::sync::{Arc, Barrier, Mutex};
use std::time::Duration;
use api::v1::meta::heartbeat_request::NodeWorkloads;
@@ -494,6 +495,170 @@ impl FrontendHeartbeatExtension for TestExtension {
}
}
struct PanickingNameExtension;
#[async_trait]
impl FrontendHeartbeatExtension for PanickingNameExtension {
fn name(&self) -> &str {
panic!("extension name panic")
}
}
#[test]
fn test_panicking_extension_name_does_not_corrupt_registry() {
let extensions = FrontendHeartbeatExtensions::default();
let result = catch_unwind(AssertUnwindSafe(|| {
extensions.register(Arc::new(PanickingNameExtension));
}));
assert!(result.is_err());
let valid = TestExtension::new("valid");
assert!(extensions.register(valid.clone()));
let registered = extensions.extensions();
assert_eq!(registered.len(), 1);
assert!(Arc::ptr_eq(
&(valid as Arc<dyn FrontendHeartbeatExtension>),
&registered[0]
));
assert_eq!(extensions.len(), 1);
}
#[test]
fn test_registry_recovers_from_poisoned_lock() {
let extensions = FrontendHeartbeatExtensions::default();
let poisoned = extensions.clone();
let result = catch_unwind(AssertUnwindSafe(move || {
let _guard = poisoned.lock();
panic!("poison registry lock")
}));
assert!(result.is_err());
let valid = TestExtension::new("valid");
assert!(extensions.register(valid.clone()));
let registered = extensions.extensions();
assert_eq!(registered.len(), 1);
assert!(Arc::ptr_eq(
&(valid as Arc<dyn FrontendHeartbeatExtension>),
&registered[0]
));
assert_eq!(extensions.len(), 1);
}
#[test]
fn test_registry_freeze_rejects_late_registration_without_changing_membership() {
let extensions = FrontendHeartbeatExtensions::default();
let calls = Arc::new(AtomicUsize::new(0));
let handler: HeartbeatResponseHandlerRef = Arc::new(ShortCircuitHandler {
result: ShortCircuitResult::Done,
calls,
});
let registered = TestExtension::with_handler("registered", handler.clone());
assert!(extensions.register(registered.clone()));
assert!(!extensions.register(TestExtension::new("registered")));
extensions.freeze();
assert!(!extensions.register(TestExtension::new("late")));
let frozen_extensions = extensions.extensions();
assert_eq!(frozen_extensions.len(), 1);
assert!(Arc::ptr_eq(
&(registered as Arc<dyn FrontendHeartbeatExtension>),
&frozen_extensions[0]
));
let frozen_handlers = extensions.response_handlers();
assert_eq!(frozen_handlers.len(), 1);
assert!(Arc::ptr_eq(&handler, &frozen_handlers[0]));
assert_eq!(extensions.len(), 1);
}
#[test]
fn test_try_register_distinguishes_duplicate_and_frozen() {
let extensions = FrontendHeartbeatExtensions::default();
assert_eq!(
extensions.try_register(TestExtension::new("registered")),
Ok(())
);
assert_eq!(
extensions.try_register(TestExtension::new("registered")),
Err(RegistrationError::Duplicate)
);
extensions.freeze();
assert_eq!(
extensions.try_register(TestExtension::new("late")),
Err(RegistrationError::Frozen)
);
}
#[test]
fn test_frozen_registration_does_not_call_extension_name() {
let extensions = FrontendHeartbeatExtensions::default();
extensions.freeze();
let compatible_result = catch_unwind(AssertUnwindSafe(|| {
extensions.register(Arc::new(PanickingNameExtension))
}));
assert!(!compatible_result.unwrap());
let typed_result = catch_unwind(AssertUnwindSafe(|| {
extensions.try_register(Arc::new(PanickingNameExtension))
}));
assert_eq!(typed_result.unwrap(), Err(RegistrationError::Frozen));
assert!(extensions.is_empty());
}
struct BlockingNameExtension {
entered_name: Arc<Barrier>,
release_name: Arc<Barrier>,
}
#[async_trait]
impl FrontendHeartbeatExtension for BlockingNameExtension {
fn name(&self) -> &str {
self.entered_name.wait();
self.release_name.wait();
"blocking"
}
}
#[test]
fn test_freeze_wins_while_registration_name_is_running() {
let extensions = FrontendHeartbeatExtensions::default();
let entered_name = Arc::new(Barrier::new(2));
let release_name = Arc::new(Barrier::new(2));
let registering = extensions.clone();
let extension = Arc::new(BlockingNameExtension {
entered_name: entered_name.clone(),
release_name: release_name.clone(),
});
let registration = std::thread::spawn(move || registering.try_register(extension));
entered_name.wait();
let freezing = extensions.clone();
let (freeze_done_tx, freeze_done_rx) = std::sync::mpsc::sync_channel(1);
let freeze = std::thread::spawn(move || {
freezing.freeze();
let _ = freeze_done_tx.send(());
});
if let Err(error) = freeze_done_rx.recv_timeout(Duration::from_secs(2)) {
release_name.wait();
let registration = registration.join();
let freeze = freeze.join();
panic!(
"freeze did not complete while extension name was blocked: {error:?}; registration: {registration:?}; freeze: {freeze:?}"
);
}
release_name.wait();
freeze.join().unwrap();
assert_eq!(registration.join().unwrap(), Err(RegistrationError::Frozen));
assert!(extensions.is_empty());
}
fn test_task(
connector: Arc<dyn HeartbeatConnector>,
extensions: FrontendHeartbeatExtensions,
src/frontend/src/instance.rs
+102
View File
@@ -59,6 +59,7 @@ use common_recordbatch::error::StreamTimeoutSnafu;
use common_telemetry::logging::SlowQueryOptions;
use common_telemetry::{debug, error, tracing};
use dashmap::DashMap;
use datafusion::dataframe::DataFrame;
use datafusion::physical_plan::ExecutionPlan;
use datafusion_expr::LogicalPlan;
use futures::{Stream, StreamExt, future};
@@ -134,6 +135,7 @@ pub struct Instance {
slow_query_options: SlowQueryOptions,
influxdb_default_merge_mode: InfluxdbMergeMode,
trace_ingest_chunk_size: usize,
otlp_resource_info: bool,
suspend: Arc<AtomicBool>,
// cache for otlp metrics
@@ -846,6 +848,22 @@ impl Instance {
query_interceptor.pre_execute(stmt.as_ref(), Some(&plan), query_ctx.clone())?;
// TQL EXPLAIN/ANALYZE formats are consumed from the query context at
// execution time (see `optimize_physical_plan`); re-apply the side
// effect of `plan_tql` that was lost when the plan was built during
// Describe. `explain_format` is per-query state, so this never
// overwrites anything.
if let Some(Statement::Tql(tql)) = &stmt {
let format = match tql {
Tql::Explain(explain) => explain.format.as_ref(),
Tql::Analyze(analyze) => analyze.format.as_ref(),
Tql::Eval(_) => None,
};
if let Some(format) = format {
query_ctx.set_explain_format(format.to_string());
}
}
let query = stmt
.as_ref()
.map(|s| s.to_string())
@@ -927,6 +945,59 @@ impl Instance {
vec![result]
}
/// Builds the [`DataFrame`] for an information-schema-backed `SHOW`
/// statement; `None` for other statements. The future is boxed to keep
/// `do_describe_inner`'s state machine small.
fn show_statement_dataframe<'a>(
&'a self,
stmt: &'a Statement,
query_ctx: &'a QueryContextRef,
) -> Pin<Box<dyn Future<Output = Option<query::error::Result<DataFrame>>> + Send + 'a>> {
Box::pin(async move {
let engine = &self.query_engine;
let catalog_manager = self.catalog_manager();
let ctx = query_ctx.clone();
let dataframe = match stmt {
Statement::ShowDatabases(show) => {
query::sql::show_databases_dataframe(show, engine, catalog_manager, ctx).await
}
Statement::ShowTables(show) => {
query::sql::show_tables_dataframe(show, engine, catalog_manager, ctx).await
}
Statement::ShowViews(show) => {
query::sql::show_views_dataframe(show, engine, catalog_manager, ctx).await
}
Statement::ShowFlows(show) => {
query::sql::show_flows_dataframe(show, engine, catalog_manager, ctx).await
}
Statement::ShowColumns(show) => {
query::sql::show_columns_dataframe(show, engine, catalog_manager, ctx).await
}
Statement::ShowTableStatus(show) => {
query::sql::show_table_status_dataframe(show, engine, catalog_manager, ctx)
.await
}
Statement::ShowCharset(kind) => {
query::sql::show_charsets_dataframe(kind, engine, catalog_manager, ctx).await
}
Statement::ShowCollation(kind) => {
query::sql::show_collations_dataframe(kind, engine, catalog_manager, ctx).await
}
Statement::ShowIndex(show) => {
query::sql::show_index_dataframe(show, engine, catalog_manager, ctx).await
}
Statement::ShowRegion(show) => {
query::sql::show_region_dataframe(show, engine, catalog_manager, ctx).await
}
Statement::ShowProcesslist(show) => {
query::sql::show_processlist_dataframe(show, engine, catalog_manager, ctx).await
}
_ => return None,
};
Some(dataframe)
})
}
async fn do_describe_inner(
&self,
stmt: Statement,
@@ -946,6 +1017,37 @@ impl Instance {
let plannable = is_inner_plannable(&stmt)
|| matches!(&stmt, Statement::Explain(explain) if is_inner_plannable(explain.statement.as_ref()));
if let Statement::Tql(tql) = stmt {
// TQL produces a logical plan; describe it from the plan so the
// extended-protocol RowDescription matches the executed DataRows.
self.check_sql_permission(&Statement::Tql(tql.clone()), &query_ctx)
.await?;
let plan = self.statement_executor.plan_tql(tql, &query_ctx).await?;
return self
.query_engine
.describe(plan, query_ctx)
.await
.map(Some)
.context(error::DescribeStatementSnafu);
}
// Describe SHOW statements from the same projection the executor builds.
if let Some(dataframe) = self
.show_statement_dataframe(&stmt, &query_ctx)
.await
.transpose()
.context(PlanStatementSnafu)?
{
self.check_sql_permission(&stmt, &query_ctx).await?;
let plan = dataframe.into_unoptimized_plan();
return self
.query_engine
.describe(plan, query_ctx)
.await
.map(Some)
.context(error::DescribeStatementSnafu);
}
if plannable {
self.check_sql_permission(&stmt, &query_ctx).await?;
src/frontend/src/instance/builder.rs
+1
View File
@@ -367,6 +367,7 @@ impl FrontendBuilder {
slow_query_options: self.options.slow_query.clone(),
influxdb_default_merge_mode: self.options.influxdb.default_merge_mode,
trace_ingest_chunk_size: self.options.otlp.trace_ingest_chunk_size,
otlp_resource_info: self.options.otlp.experimental_enable_resource_info,
suspend: Arc::new(AtomicBool::new(false)),
})
}
src/frontend/src/instance/entity_graph.rs
+211 -19
View File
@@ -23,7 +23,7 @@
//! the query engine. Injected into the catalog manager after the engine is built,
//! breaking the `catalog -> query` cycle.
use std::collections::HashMap;
use std::collections::{BTreeMap, HashMap};
use std::sync::Weak;
use async_trait::async_trait;
@@ -60,8 +60,9 @@ use table::TableRef;
use table::metadata::TableInfo;
use table::predicate::{TimeRangeExtraction, extract_time_range_strict};
use table::requests::{
EntityRole, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, SOURCE_PROMETHEUS,
is_trace_v1_table, parse_entity_columns, parse_entity_option_key,
EntityRole, SEMANTIC_METRIC_TYPE, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC,
SOURCE_OPENTELEMETRY, SOURCE_PROMETHEUS, is_trace_v1_table, parse_entity_columns,
parse_entity_option_key,
};
use crate::error;
@@ -177,6 +178,7 @@ impl EntityGraphProviderImpl {
time_index: time_index.clone(),
entity_type,
id_columns,
id_qualifier: None,
descriptive_columns,
scope_columns,
})
@@ -188,10 +190,11 @@ impl EntityGraphProviderImpl {
/// All entity declarations of one table: the explicit options plus the
/// zero-configuration conventions (`otlp_trace_entities` for trace-v1
/// tables — including the `service` identity of tables created before the
/// ingest-side auto-stamp — and the Prometheus descriptor whitelist). An
/// explicit declaration of a type always suppresses the implicit one, even
/// when the explicit declaration is invalid and skipped: silently falling
/// back would change entity identity behind the user's back.
/// ingest-side auto-stamp — and the Prometheus/OTel descriptor
/// whitelists). An explicit declaration of a type always suppresses the
/// implicit one, even when the explicit declaration is invalid and
/// skipped: silently falling back would change entity identity behind the
/// user's back.
fn declarations_for(
table_info: &TableInfo,
conventions: &Conventions,
@@ -204,7 +207,20 @@ impl EntityGraphProviderImpl {
&mut declarations,
);
}
Self::extend_with_prometheus_conventions(table_info, conventions, &mut declarations);
Self::extend_with_info_metric_conventions(
table_info,
&conventions.prometheus_info_metrics,
SOURCE_PROMETHEUS,
None,
&mut declarations,
);
Self::extend_with_info_metric_conventions(
table_info,
&conventions.otel_info_metrics,
SOURCE_OPENTELEMETRY,
Some(servers::semantic::METRIC_TYPE_INFO),
&mut declarations,
);
declarations
}
@@ -216,29 +232,34 @@ impl EntityGraphProviderImpl {
})
}
/// Implicit declarations of the well-known Prometheus entity-descriptor
/// metrics (the `prometheus_info_metrics` whitelist of `conventions.yaml`),
/// gated on the ingest-stamped `signal_type=metric` + `source=prometheus`
/// options and keyed by table name. The metric engine's physical table
/// Implicit declarations of the well-known entity-descriptor metrics
/// (the `prometheus_info_metrics` / `otel_info_metrics` whitelists of
/// `conventions.yaml`), gated on the ingest-stamped `signal_type=metric`
/// option plus the whitelist's expected `source`; OTel descriptors also
/// require `metric.type=info`. The metric engine's physical table
/// aggregates every logical table's columns and must not contribute a
/// duplicate source.
fn extend_with_prometheus_conventions(
fn extend_with_info_metric_conventions(
table_info: &TableInfo,
conventions: &Conventions,
whitelist: &BTreeMap<String, Vec<ImplicitEntity>>,
expected_source: &str,
expected_metric_type: Option<&str>,
declarations: &mut Vec<EntityDeclaration>,
) {
let Some(implicit_entities) = conventions.prometheus_info_metrics.get(&table_info.name)
else {
let Some(implicit_entities) = whitelist.get(&table_info.name) else {
return;
};
let options = &table_info.meta.options.extra_options;
if options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str) != Some(SIGNAL_TYPE_METRIC)
|| options.get(SEMANTIC_SOURCE).map(String::as_str) != Some(SOURCE_PROMETHEUS)
|| options.get(SEMANTIC_SOURCE).map(String::as_str) != Some(expected_source)
|| expected_metric_type.is_some_and(|expected| {
options.get(SEMANTIC_METRIC_TYPE).map(String::as_str) != Some(expected)
})
|| table_info.is_physical_table()
{
debug!(
"Table `{}` matches the info-metric whitelist but is not a prometheus logical \
metric table; skipping its implicit declarations",
"Table `{}` matches the info-metric whitelist but is not an eligible \
`{expected_source}` info-metric source; skipping its implicit declarations",
table_info.name
);
return;
@@ -296,12 +317,19 @@ impl EntityGraphProviderImpl {
.cloned()
.collect()
};
// A table predating the qualifier column keeps the unqualified
// identity rather than losing the declaration.
let id_qualifier = implicit
.qualified_by
.clone()
.filter(|c| schema.column_schema_by_name(c).is_some());
declarations.push(EntityDeclaration {
schema: table_info.schema_name.clone(),
table: table_info.name.clone(),
time_index: time_index.clone(),
entity_type: implicit.entity.clone(),
id_columns: implicit.id.clone(),
id_qualifier,
descriptive_columns,
scope_columns: vec![],
});
@@ -832,6 +860,24 @@ mod tests {
declarations[2].id_columns,
vec!["service_name", "resource_attributes.service.instance.id"]
);
assert_eq!(declarations[1].id_qualifier, None);
let namespaced = table_info(
&[
"service_name",
"resource_attributes.service.namespace",
"resource_attributes.service.instance.id",
],
&[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)],
);
for declaration in sorted_declarations(&namespaced) {
assert_eq!(
declaration.id_qualifier.as_deref(),
Some("resource_attributes.service.namespace"),
"{} must qualify its identity like the metric side's job",
declaration.entity_type
);
}
// Missing uid column: no pod entity synthesized, no name-based guess.
let no_uid = table_info(
@@ -856,6 +902,47 @@ mod tests {
assert!(sorted_declarations(&invalid_explicit).is_empty());
}
#[test]
fn trace_table_host_and_container_require_stable_ids() {
let with_ids = table_info(
&[
"service_name",
"resource_attributes.host.id",
"resource_attributes.host.name",
"resource_attributes.container.id",
],
&[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)],
);
let declarations = sorted_declarations(&with_ids);
let types: Vec<&str> = declarations
.iter()
.map(|d| d.entity_type.as_str())
.collect();
assert_eq!(types, vec!["container", "host", "service"]);
assert_eq!(
declarations[1].id_columns,
vec!["resource_attributes.host.id"]
);
assert_eq!(
declarations[1].descriptive_columns,
vec!["resource_attributes.host.name"]
);
let names_only = table_info(
&[
"service_name",
"resource_attributes.host.name",
"resource_attributes.container.name",
],
&[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)],
);
let types: Vec<String> = sorted_declarations(&names_only)
.into_iter()
.map(|d| d.entity_type)
.collect();
assert_eq!(types, vec!["service"]);
}
#[test]
fn prometheus_implicit_declarations_are_gated() {
let labels: &[&str] = &["namespace", "pod", "node"];
@@ -900,6 +987,111 @@ mod tests {
assert_eq!(declarations[1].id_columns, vec!["pod"]);
}
const OTEL_STAMPS: &[(&str, &str)] = &[
(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC),
(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY),
(SEMANTIC_METRIC_TYPE, servers::semantic::METRIC_TYPE_INFO),
];
#[test]
fn greptime_otel_resource_info_gets_implicit_declarations() {
let info = prom_table_info(
"greptime_otel_resource_info",
&[
"job",
"instance",
"service.name",
"service.namespace",
"host.id",
"host.name",
"container.id",
"container.name",
"k8s.pod.uid",
"k8s.pod.name",
"k8s.namespace.name",
],
OTEL_STAMPS,
);
let declarations = sorted_declarations(&info);
let types: Vec<&str> = declarations
.iter()
.map(|d| d.entity_type.as_str())
.collect();
assert_eq!(
types,
vec![
"container",
"host",
"k8s.pod",
"service",
"service.instance"
]
);
assert_eq!(declarations[1].id_columns, vec!["host.id"]);
assert_eq!(declarations[1].descriptive_columns, vec!["host.name"]);
assert_eq!(declarations[3].id_columns, vec!["job"]);
assert_eq!(
declarations[3].descriptive_columns,
vec!["service.name", "service.namespace"]
);
assert_eq!(declarations[4].id_columns, vec!["job", "instance"]);
assert!(declarations[4].descriptive_columns.is_empty());
let partial = prom_table_info(
"greptime_otel_resource_info",
&["job", "service.name", "host.id"],
OTEL_STAMPS,
);
let types: Vec<String> = sorted_declarations(&partial)
.into_iter()
.map(|d| d.entity_type)
.collect();
assert_eq!(types, vec!["host", "service"]);
}
#[test]
fn otel_implicit_declarations_are_gated() {
let labels: &[&str] = &["job", "instance", "host.id"];
assert!(
sorted_declarations(&prom_table_info(
"greptime_otel_resource_info",
labels,
PROM_STAMPS
))
.is_empty()
);
let mut stamps = OTEL_STAMPS.to_vec();
stamps.push((PHYSICAL_TABLE_METADATA_KEY, "true"));
assert!(
sorted_declarations(&prom_table_info(
"greptime_otel_resource_info",
labels,
&stamps
))
.is_empty()
);
assert!(
conventions()
.unwrap()
.otel_info_metrics
.contains_key(servers::otlp::metrics::OTEL_RESOURCE_INFO_TABLE_NAME)
);
let wrong_type = [
(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC),
(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY),
(SEMANTIC_METRIC_TYPE, servers::semantic::METRIC_TYPE_GAUGE),
];
assert!(
sorted_declarations(&prom_table_info(
"greptime_otel_resource_info",
labels,
&wrong_type
))
.is_empty()
);
}
#[test]
fn target_info_descriptive_rest_covers_remaining_tags() {
let info = prom_table_info(
src/frontend/src/instance/otlp.rs
+42 -5
View File
@@ -50,7 +50,7 @@ use table::requests::{
use self::trace_ingest::trace_conventions;
use crate::instance::Instance;
use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS};
use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS, OTLP_RESOURCE_INFO_WRITE_ERRORS};
fn trace_permission_targets(
table_name: &str,
@@ -119,9 +119,15 @@ impl OpenTelemetryProtocolHandler for Instance {
.cloned()
.unwrap_or_default();
metric_ctx.is_legacy = is_legacy;
metric_ctx.resource_info = self.otlp_resource_info;
let (requests, rows, semantic_index, mut outcome) =
otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)?;
let otlp::metrics::MetricsConversion {
requests,
rows,
semantic_index,
resource_info,
mut outcome,
} = otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)?;
if outcome.rejected_data_points > 0 {
warn!(
"Rejected {} OTLP exponential histogram data points: {}",
@@ -156,7 +162,7 @@ impl OpenTelemetryProtocolHandler for Instance {
// OTLP tables have one sample field in both the legacy and physical paths.
let output = if metric_ctx.is_legacy || !metric_ctx.with_metric_engine {
self.handle_row_inserts(requests, ctx, false, true)
self.handle_row_inserts(requests, ctx.clone(), false, true)
.await
.map_err(BoxedError::new)
.context(error::ExecuteGrpcQuerySnafu)
@@ -165,12 +171,43 @@ impl OpenTelemetryProtocolHandler for Instance {
.extension(PHYSICAL_TABLE_PARAM)
.unwrap_or(GREPTIME_PHYSICAL_TABLE)
.to_string();
self.handle_metric_row_inserts(requests, ctx, physical_table)
self.handle_metric_row_inserts(requests, ctx.clone(), physical_table)
.await
.map_err(BoxedError::new)
.context(error::ExecuteGrpcQuerySnafu)
}?;
outcome.write_cost = output.meta.cost;
// Derived enrichment, written after the metric data is committed:
// failing here would make the client retry data the server already
// accepted, so every failure degrades to a warning instead.
if let Some(resource_info) = resource_info {
let written = match self.check_row_insert_permission(
&resource_info,
&ctx,
PermissionReq::Action(OTLP_WRITE),
) {
Ok(_) => self
.handle_row_inserts(resource_info, ctx, false, false)
.await
.map_err(BoxedError::new)
.map_err(|e| e.to_string()),
Err(e) => Err(e.to_string()),
};
match written {
Ok(descriptor_output) => outcome.write_cost += descriptor_output.meta.cost,
Err(e) => {
OTLP_RESOURCE_INFO_WRITE_ERRORS.inc();
warn!("Failed to write the OTLP resource descriptor table: {e}");
outcome.error_message.get_or_insert(format!(
"metric data was accepted, but writing the resource \
descriptor table `{}` failed: {e}",
otlp::metrics::OTEL_RESOURCE_INFO_TABLE_NAME
));
}
}
}
Ok(outcome)
}
src/frontend/src/lib.rs
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#![recursion_limit = "256"]
pub mod error;
pub mod events;
pub mod frontend;
src/frontend/src/metrics.rs
+8
View File
@@ -45,6 +45,14 @@ lazy_static! {
)
.unwrap();
/// Failed writes of the synthesized OTLP resource descriptor table; these
/// surface as OTLP partial-success warnings, not request failures.
pub static ref OTLP_RESOURCE_INFO_WRITE_ERRORS: IntCounter = register_int_counter!(
"greptime_frontend_otlp_resource_info_write_errors",
"frontend otlp resource descriptor write errors"
)
.unwrap();
/// The number of OpenTelemetry traces send by frontend node.
pub static ref OTLP_TRACES_ROWS: IntCounter = register_int_counter!(
"greptime_frontend_otlp_traces_rows",
src/frontend/src/service_config/otlp.rs
+7
View File
@@ -23,6 +23,11 @@ pub struct OtlpOptions {
pub experimental_enable_exponential_histogram: bool,
/// Maximum spans per trace ingest chunk. Set to 0 to disable splitting.
pub trace_ingest_chunk_size: usize,
/// Whether to synthesize the `greptime_otel_resource_info` descriptor
/// table from the resource attributes of OTLP metrics, so metrics-only
/// services reach the semantic graph. Off by default: it creates and
/// writes a table the user did not send.
pub experimental_enable_resource_info: bool,
}
impl Default for OtlpOptions {
@@ -31,6 +36,7 @@ impl Default for OtlpOptions {
enable: true,
experimental_enable_exponential_histogram: false,
trace_ingest_chunk_size: DEFAULT_TRACE_INGEST_CHUNK_SIZE,
experimental_enable_resource_info: false,
}
}
}
@@ -45,6 +51,7 @@ mod tests {
assert!(default.enable);
assert!(!default.experimental_enable_exponential_histogram);
assert_eq!(default.trace_ingest_chunk_size, 512);
assert!(!default.experimental_enable_resource_info);
let options: OtlpOptions = toml::from_str("enable = false").unwrap();
assert!(!options.enable);
src/meta-srv/src/handler/persist_stats_handler.rs
+3 -5
View File
@@ -22,7 +22,7 @@ use client::inserter::{Context as InserterContext, Inserter};
use common_catalog::consts::DEFAULT_PRIVATE_SCHEMA_NAME;
use common_macro::{Schema, ToRow};
use common_meta::DatanodeId;
use common_meta::datanode::RegionStat;
use common_meta::datanode::{REGION_STATS_HISTORY_TABLE_NAME, RegionStat};
use common_telemetry::warn;
use dashmap::DashMap;
use store_api::region_engine::RegionRole;
@@ -40,8 +40,6 @@ pub struct PersistStatsHandler {
persist_interval: Duration,
}
/// The name of the table to persist region stats.
const META_REGION_STATS_HISTORY_TABLE_NAME: &str = "region_statistics_history";
/// The default context to persist region stats.
const DEFAULT_CONTEXT: InserterContext = InserterContext {
catalog: DEFAULT_CATALOG_NAME,
@@ -207,7 +205,7 @@ impl PersistStatsHandler {
&DEFAULT_CONTEXT,
RowInsertRequests {
inserts: vec![RowInsertRequest {
table_name: META_REGION_STATS_HISTORY_TABLE_NAME.to_string(),
table_name: REGION_STATS_HISTORY_TABLE_NAME.to_string(),
rows: Some(Rows {
schema: PersistRegionStat::schema(),
rows,
@@ -537,7 +535,7 @@ mod tests {
};
assert_eq!(
request.table_name,
META_REGION_STATS_HISTORY_TABLE_NAME.to_string()
REGION_STATS_HISTORY_TABLE_NAME.to_string()
);
assert_eq!(request.rows.unwrap().rows, vec![expected_row]);
src/metric-engine/src/engine/flush.rs
+7 -6
View File
@@ -121,6 +121,7 @@ mod tests {
.map(|path| path.replace(&e.file_id, "<file_id>"));
e.file_id = "<file_id>".to_string();
e.index_version = 0;
e.max_row_group_uncompressed_size = 0;
// Round down sizes to nearest 1000 to avoid exact size
// comparisons that break when the SST format changes.
e.file_size = e.file_size / 1000 * 1000;
@@ -133,12 +134,12 @@ mod tests {
assert_eq!(
debug_format,
r#"
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/data/<file_id>.parquet", file_size: 3000, index_file_path: Some("test_metric_region/11_0000000001/data/index/<file_id>.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(20), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/data/<file_id>.parquet", file_size: 3000, index_file_path: Some("test_metric_region/11_0000000002/data/index/<file_id>.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417473(11, 16777217), table_id: 11, region_number: 16777217, region_group: 1, region_sequence: 1, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/metadata/<file_id>.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 8, num_row_groups: 1, num_series: Some(8), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(8), origin_region_id: 47261417473(11, 16777217), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665867_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665868\0\0\x06") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417474(11, 16777218), table_id: 11, region_number: 16777218, region_group: 1, region_sequence: 2, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/metadata/<file_id>.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), origin_region_id: 47261417474(11, 16777218), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665858_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665858\0\0\x06") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/data/<file_id>.parquet", file_size: 3000, index_file_path: Some("test_metric_region/22_0000000042/data/index/<file_id>.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94506057770(22, 16777258), table_id: 22, region_number: 16777258, region_group: 1, region_sequence: 42, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/metadata/<file_id>.parquet", file_size: 4000, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), origin_region_id: 94506057770(22, 16777258), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_9543417\t331754_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_9543417\t331754\0\0\x06") }"#,
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/data/<file_id>.parquet", file_size: 3000, max_row_group_uncompressed_size: 0, index_file_path: Some("test_metric_region/11_0000000001/data/index/<file_id>.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(20), partition_expr: None, origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/data/<file_id>.parquet", file_size: 3000, max_row_group_uncompressed_size: 0, index_file_path: Some("test_metric_region/11_0000000002/data/index/<file_id>.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x04W\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417473(11, 16777217), table_id: 11, region_number: 16777217, region_group: 1, region_sequence: 1, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000001/metadata/<file_id>.parquet", file_size: 4000, max_row_group_uncompressed_size: 0, index_file_path: None, index_file_size: None, num_rows: 8, num_row_groups: 1, num_series: Some(8), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(8), partition_expr: None, origin_region_id: 47261417473(11, 16777217), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665867_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665868\0\0\x06") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 47261417474(11, 16777218), table_id: 11, region_number: 16777218, region_group: 1, region_sequence: 2, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/11_0000000002/metadata/<file_id>.parquet", file_size: 4000, max_row_group_uncompressed_size: 0, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), partition_expr: None, origin_region_id: 47261417474(11, 16777218), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_4771708\t665858_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_4771708\t665858\0\0\x06") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/data/<file_id>.parquet", file_size: 3000, max_row_group_uncompressed_size: 0, index_file_path: Some("test_metric_region/22_0000000042/data/index/<file_id>.puffin"), index_file_size: Some(0), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05"), primary_key_max: Some(b"\x80\0\0\x04\x01\0\0\x08\xae\x80\0\0\x03\x01)\x02K\x9f\xa5A\xb7\xb3\0\0\0\x02\x01\x01tag_0\0\0\0\x05") }
ManifestSstEntry { table_dir: "test_metric_region/", region_id: 94506057770(22, 16777258), table_id: 22, region_number: 16777258, region_group: 1, region_sequence: 42, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test_metric_region/22_0000000042/metadata/<file_id>.parquet", file_size: 4000, max_row_group_uncompressed_size: 0, index_file_path: None, index_file_size: None, num_rows: 4, num_row_groups: 1, num_series: Some(4), min_ts: 0::Millisecond, max_ts: 0::Millisecond, sequence: Some(4), partition_expr: None, origin_region_id: 94506057770(22, 16777258), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01__column\t_9543417\t331754_Z\t3JlcHRpb\tWVfdGltZ\tXN0YW1w\0\x07"), primary_key_max: Some(b"\x01\x01__region\t_9543417\t331754\0\0\x06") }"#,
);
// list from storage
let storage_entries = mito
src/mito2/benches/bench_compaction_picker.rs
-66
View File
@@ -17,7 +17,6 @@ use std::hint::black_box;
use criterion::{Criterion, criterion_group, criterion_main};
use mito2::compaction::run::{
Item, Ranged, SortedRun, find_overlapping_items, find_sorted_runs, find_sorted_runs_original,
merge_seq_files, reduce_runs,
};
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
@@ -49,20 +48,6 @@ impl Item for MockFile {
}
}
fn generate_test_files(n: usize) -> Vec<MockFile> {
let mut files = Vec::with_capacity(n);
for _ in 0..n {
// Create slightly overlapping ranges to force multiple sorted runs
files.push(MockFile {
start: 0,
end: 10,
pk: 0,
size: 10,
});
}
files
}
fn generate_same_timestamp_files(total_files: usize, files_per_timestamp: usize) -> Vec<MockFile> {
const TIMESTAMP_INTERVAL_SECS: i64 = 10 * 60;
@@ -107,21 +92,6 @@ fn bench_find_sorted_runs(c: &mut Criterion) {
group.finish();
}
fn bench_reduce_runs(c: &mut Criterion) {
let mut group = c.benchmark_group("reduce_runs");
for size in [10, 100, 1000].iter() {
group.bench_function(format!("size_{}", size), |b| {
let mut files = generate_test_files(*size);
let runs = find_sorted_runs(&mut files);
b.iter(|| {
reduce_runs(black_box(runs.clone()));
});
});
}
group.finish();
}
fn bench_find_overlapping_items(c: &mut Criterion) {
let mut group = c.benchmark_group("find_overlapping_items");
@@ -158,45 +128,9 @@ fn bench_find_overlapping_items(c: &mut Criterion) {
group.finish();
}
fn bench_merge_seq_files(c: &mut Criterion) {
let mut group = c.benchmark_group("merge_seq_files");
for size in [10, 100, 1000].iter() {
group.bench_function(format!("size_{}", size), |b| {
// Create a set of files with varying sizes
let mut files = Vec::with_capacity(*size);
for i in 0..*size {
// Create files with different sizes to test the scoring algorithm
let file_size = if i % 3 == 0 {
5
} else if i % 3 == 1 {
10
} else {
15
};
files.push(MockFile {
start: i as i64,
end: (i + 1) as i64,
pk: 0,
size: file_size,
});
}
b.iter(|| {
merge_seq_files(black_box(&files), black_box(Some(50)));
});
});
}
group.finish();
}
criterion_group!(
benches,
bench_find_sorted_runs,
bench_reduce_runs,
bench_find_overlapping_items,
bench_merge_seq_files
);
criterion_main!(benches);
src/mito2/src/compaction/run.rs
+59 -422
View File
@@ -20,14 +20,9 @@ use std::collections::BinaryHeap;
use bytes::{Buf, Bytes};
use common_base::BitVec;
use common_base::readable_size::ReadableSize;
use common_time::Timestamp;
use smallvec::{SmallVec, smallvec};
use crate::sst::file::{FileHandle, RegionFileId};
/// Default max compaction output file size when not specified.
const DEFAULT_MAX_OUTPUT_SIZE: u64 = ReadableSize::gb(2).as_bytes();
use crate::sst::file::FileHandle;
/// Trait for any items with specific range (both boundaries are inclusive).
pub trait Ranged {
@@ -152,71 +147,11 @@ pub trait Item: Ranged + Clone {
fn size(&self) -> usize;
}
/// A group of files that are created by the same compaction task.
#[derive(Debug, Clone)]
pub struct FileGroup {
files: SmallVec<[FileHandle; 2]>,
size: usize,
num_rows: usize,
min_timestamp: Timestamp,
max_timestamp: Timestamp,
primary_key_range: Option<(Bytes, Bytes)>,
}
impl FileGroup {
pub(crate) fn new_with_file(file: FileHandle) -> Self {
let size = file.size() as usize;
let (min_timestamp, max_timestamp) = file.time_range();
let num_rows = file.num_rows();
let primary_key_range = file.primary_key_range();
Self {
files: smallvec![file],
size,
num_rows,
min_timestamp,
max_timestamp,
primary_key_range,
}
}
pub(crate) fn num_rows(&self) -> usize {
self.num_rows
}
pub(crate) fn add_file(&mut self, file: FileHandle) {
self.size += file.size() as usize;
self.num_rows += file.num_rows();
let (min_timestamp, max_timestamp) = file.time_range();
self.min_timestamp = self.min_timestamp.min(min_timestamp);
self.max_timestamp = self.max_timestamp.max(max_timestamp);
self.primary_key_range =
merge_primary_key_ranges(self.primary_key_range.take(), file.primary_key_range());
self.files.push(file);
}
pub(crate) fn num_files(&self) -> usize {
self.files.len()
}
#[cfg(test)]
pub(crate) fn files(&self) -> &[FileHandle] {
&self.files[..]
}
pub(crate) fn file_ids(&self) -> SmallVec<[RegionFileId; 2]> {
SmallVec::from_iter(self.files.iter().map(|f| f.file_id()))
}
pub(crate) fn into_files(self) -> impl Iterator<Item = FileHandle> {
self.files.into_iter()
}
}
impl Ranged for FileGroup {
impl Ranged for FileHandle {
type BoundType = Timestamp;
fn range(&self) -> (Self::BoundType, Self::BoundType) {
(self.min_timestamp, self.max_timestamp)
self.time_range()
}
fn overlap(&self, other: &Self) -> bool {
@@ -226,7 +161,7 @@ impl Ranged for FileGroup {
return false;
}
match (&self.primary_key_range, &other.primary_key_range) {
match (&self.primary_key_range(), &other.primary_key_range()) {
(Some(lhs), Some(rhs)) => primary_key_ranges_overlap(lhs, rhs),
_ => true,
}
@@ -239,16 +174,16 @@ impl Ranged for FileGroup {
return false;
}
match (&self.primary_key_range, &other.primary_key_range) {
match (&self.primary_key_range(), &other.primary_key_range()) {
(Some(lhs), Some(rhs)) => primary_key_ranges_overlap(lhs, rhs),
_ => true,
}
}
}
impl Item for FileGroup {
impl Item for FileHandle {
fn size(&self) -> usize {
self.size
self.size() as usize
}
}
@@ -553,124 +488,8 @@ where
runs.into_iter().map(|x| x.run).collect()
}
/// Finds a set of files with minimum penalty to merge that can reduce the total num of runs.
/// The penalty of merging is defined as the size of all overlapping files between two runs.
pub fn reduce_runs<T: Item>(mut runs: Vec<SortedRun<T>>) -> Vec<T> {
assert!(runs.len() > 1);
// sort runs by size
runs.sort_unstable_by_key(|a| a.size);
// limit max probe runs to 100
let probe_end = runs.len().min(100);
let mut min_penalty = usize::MAX;
let mut files = vec![];
let mut temp_files = vec![];
for i in 0..probe_end {
for j in i + 1..probe_end {
let (a, b) = runs.split_at_mut(j);
find_overlapping_items(&mut a[i], &mut b[0], &mut temp_files);
let penalty = temp_files.iter().map(|e| e.size()).sum();
if penalty < min_penalty {
min_penalty = penalty;
files.clear();
files.extend_from_slice(&temp_files);
}
}
}
files
}
/// Finds the optimal set of adjacent files to merge based on a scoring system.
///
/// This function evaluates all possible contiguous subsets of files to find the best
/// candidates for merging, considering:
///
/// 1. File reduction - prioritizes merging more files to reduce the total count
/// 2. Write amplification - minimizes the ratio of largest file to total size
/// 3. Size efficiency - prefers merges that utilize available space effectively
///
/// When multiple merge candidates have the same score, older files (those with lower indices)
/// are preferred.
///
/// # Arguments
/// * `input_files` - Slice of files to consider for merging
/// * `max_file_size` - Optional maximum size constraint for the merged file.
/// If None, uses 1.5 times the average file size.
///
/// # Returns
/// A vector containing the best set of adjacent files to merge.
/// Returns an empty vector if input is empty or contains only one file.
pub fn merge_seq_files<T: Item>(input_files: &[T], max_file_size: Option<u64>) -> Vec<T> {
if input_files.is_empty() || input_files.len() == 1 {
return vec![];
}
// Limit the number of files to process to 100 to control time complexity
let files_to_process = if input_files.len() > 100 {
&input_files[0..100]
} else {
input_files
};
// Calculate target size based on max_file_size or average file size
let target_size = match max_file_size {
Some(size) => size as usize,
None => {
// Calculate 1.5*average_file_size if max_file_size is not provided and clamp to 2GB
let total_size: usize = files_to_process.iter().map(|f| f.size()).sum();
((((total_size as f64) / (files_to_process.len() as f64)) * 1.5) as usize)
.min(DEFAULT_MAX_OUTPUT_SIZE as usize)
}
};
// Find the best group of adjacent files to merge
let mut best_group = Vec::new();
let mut best_score = f64::NEG_INFINITY;
// Try different starting positions - iterate from end to start to prefer older files
for start_idx in (0..files_to_process.len()).rev() {
// Try different ending positions - also iterate from end to start
for end_idx in (start_idx + 1..files_to_process.len()).rev() {
let group = &files_to_process[start_idx..=end_idx];
let total_size: usize = group.iter().map(|f| f.size()).sum();
// Skip if total size exceeds target size
if total_size > target_size {
continue; // Use continue instead of break to check smaller ranges
}
// Calculate amplification factor (largest file size / total size)
let largest_file_size = group.iter().map(|f| f.size()).max().unwrap_or(0);
let amplification_factor = largest_file_size as f64 / total_size as f64;
// Calculate file reduction (number of files that will be reduced)
let file_reduction = group.len() - 1;
// Calculate score based on multiple factors:
// 1. File reduction (higher is better)
// 2. Amplification factor (lower is better)
// 3. Size efficiency (how close to target size)
let file_reduction_score = file_reduction as f64 / files_to_process.len() as f64;
let amp_factor_score = (1.0 - amplification_factor) * 1.5; // Lower amplification is better
let size_efficiency = (total_size as f64 / target_size as f64).min(1.0); // Reward using available space
let score = file_reduction_score + amp_factor_score + size_efficiency;
// Check if this group is better than our current best
// Use >= instead of > to prefer older files (which we encounter first due to reverse iteration)
if score >= best_score {
best_score = score;
best_group = group.to_vec();
}
}
}
best_group
}
#[cfg(test)]
mod tests {
use std::collections::HashSet;
use bytes::Bytes;
use store_api::storage::FileId;
@@ -709,17 +528,6 @@ mod tests {
.collect()
}
fn build_items_with_size(items: &[(i64, i64, usize)]) -> Vec<MockFile> {
items
.iter()
.map(|(start, end, size)| MockFile {
start: *start,
end: *end,
size: *size,
})
.collect()
}
fn pk_range(min: &'static [u8], max: &'static [u8]) -> Option<(Bytes, Bytes)> {
Some((Bytes::from_static(min), Bytes::from_static(max)))
}
@@ -844,99 +652,6 @@ mod tests {
}
}
fn check_reduce_runs(
files: &[(i64, i64)],
expected_runs: &[Vec<(i64, i64)>],
expected: &[(i64, i64)],
) {
let runs = check_sorted_runs(files, expected_runs);
if runs.len() <= 1 {
assert!(expected.is_empty());
return;
}
let files_to_merge = reduce_runs(runs);
let file_to_merge_timestamps = files_to_merge
.into_iter()
.map(|f| (f.start, f.end))
.collect::<HashSet<_>>();
let expected = expected.iter().cloned().collect::<HashSet<_>>();
assert_eq!(&expected, &file_to_merge_timestamps);
}
#[test]
fn test_reduce_runs() {
// [1..3] [5..6]
// [2..4]
check_reduce_runs(
&[(1, 3), (2, 4), (5, 6)],
&[vec![(1, 3), (5, 6)], vec![(2, 4)]],
&[(1, 3), (2, 4)],
);
// [1..2][3..5]
// [4..6]
check_reduce_runs(
&[(1, 2), (3, 5), (4, 6)],
&[vec![(1, 2), (3, 5)], vec![(4, 6)]],
&[(3, 5), (4, 6)],
);
// [1..2][3..4] [7..8]
// [4..6]
check_reduce_runs(
&[(1, 2), (3, 4), (4, 6), (7, 8)],
&[vec![(1, 2), (3, 4), (4, 6), (7, 8)]],
&[],
);
// [1..2][3........6][7..8][8..9]
// [3..4][5..6]
check_reduce_runs(
&[(1, 2), (3, 4), (5, 6), (3, 6), (7, 8), (8, 9)],
&[vec![(1, 2), (3, 6), (7, 8), (8, 9)], vec![(3, 4), (5, 6)]],
&[(5, 6), (3, 4), (3, 6)], // already satisfied
);
// [1..2][3........6][7..8][8..9]
// [3..4][5..6]
check_reduce_runs(
&[(1, 2), (3, 4), (5, 6), (3, 6), (7, 8), (8, 9)],
&[vec![(1, 2), (3, 6), (7, 8), (8, 9)], vec![(3, 4), (5, 6)]],
&[(3, 4), (3, 6), (5, 6)],
);
// [10..20] [30..40] [50........80][80...100][100..110]
// [50..60] [80..90]
//
check_reduce_runs(
&[
(10, 20),
(30, 40),
(50, 60),
(50, 80),
(80, 90),
(80, 100),
(100, 110),
],
&[
vec![(10, 20), (30, 40), (50, 80), (80, 100), (100, 110)],
vec![(50, 60), (80, 90)],
],
&[(50, 80), (80, 100), (50, 60), (80, 90)],
);
// [0..10]
// [0...11]
// [0....12]
// [0.....13]
check_reduce_runs(
&[(0, 10), (0, 11), (0, 12), (0, 13)],
&[vec![(0, 13)], vec![(0, 12)], vec![(0, 11)], vec![(0, 10)]],
&[(0, 10), (0, 11)],
);
}
#[test]
fn test_find_overlapping_items() {
let mut result = Vec::new();
@@ -1038,27 +753,25 @@ mod tests {
}
#[test]
fn test_file_group_overlap_time_overlap_pk_disjoint() {
let lhs =
FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
0,
100,
0,
1,
10,
pk_range(b"a", b"f"),
));
let rhs =
FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
50,
150,
0,
2,
10,
pk_range(b"x", b"z"),
));
fn test_file_overlap_time_overlap_pk_disjoint() {
let lhs = new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
0,
100,
0,
1,
10,
pk_range(b"a", b"f"),
);
let rhs = new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
50,
150,
0,
2,
10,
pk_range(b"x", b"z"),
);
assert!(!lhs.overlap(&rhs));
}
@@ -1066,7 +779,7 @@ mod tests {
#[test]
fn test_find_sorted_runs_collapses_pk_disjoint_files_into_one_run() {
let mut files = vec![
FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range(
new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
0,
100,
@@ -1074,8 +787,8 @@ mod tests {
1,
10,
pk_range(b"a", b"f"),
)),
FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range(
),
new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
50,
150,
@@ -1083,7 +796,7 @@ mod tests {
2,
10,
pk_range(b"x", b"z"),
)),
),
];
let runs = find_sorted_runs(&mut files);
@@ -1095,7 +808,7 @@ mod tests {
#[test]
fn test_find_sorted_runs_handles_2d_transitivity_break() {
let mut files = vec![
FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range(
new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
0,
100,
@@ -1103,8 +816,8 @@ mod tests {
1,
10,
pk_range(b"a", b"f"),
)),
FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range(
),
new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
50,
150,
@@ -1112,8 +825,8 @@ mod tests {
2,
10,
pk_range(b"x", b"z"),
)),
FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range(
),
new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
50,
150,
@@ -1121,7 +834,7 @@ mod tests {
3,
10,
pk_range(b"a", b"f"),
)),
),
];
let runs = find_sorted_runs(&mut files);
@@ -1133,7 +846,7 @@ mod tests {
#[test]
fn test_find_overlapping_items_skips_pk_disjoint_pairs() {
let mut left = SortedRun::from(vec![FileGroup::new_with_file(
let mut left = SortedRun::from(vec![
new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
0,
@@ -1143,8 +856,8 @@ mod tests {
10,
pk_range(b"a", b"f"),
),
)]);
let mut right = SortedRun::from(vec![FileGroup::new_with_file(
]);
let mut right = SortedRun::from(vec![
new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
50,
@@ -1154,7 +867,7 @@ mod tests {
10,
pk_range(b"x", b"z"),
),
)]);
]);
let mut result = Vec::new();
find_overlapping_items(&mut left, &mut right, &mut result);
@@ -1163,102 +876,26 @@ mod tests {
}
#[test]
fn test_file_group_touching_time_boundary_with_same_pk_is_not_overlap() {
let lhs =
FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
0,
100,
0,
1,
10,
pk_range(b"a", b"f"),
));
let rhs =
FileGroup::new_with_file(new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
100,
150,
0,
2,
10,
pk_range(b"a", b"f"),
));
fn test_file_touching_time_boundary_with_same_pk_is_not_overlap() {
let lhs = new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
0,
100,
0,
1,
10,
pk_range(b"a", b"f"),
);
let rhs = new_file_handle_with_size_sequence_and_primary_key_range(
FileId::random(),
100,
150,
0,
2,
10,
pk_range(b"a", b"f"),
);
assert!(!lhs.overlap(&rhs));
}
#[test]
fn test_merge_seq_files() {
// Test empty input
let files = Vec::<MockFile>::new();
assert_eq!(merge_seq_files(&files, None), Vec::<MockFile>::new());
// Test single file input (should return empty vec as no merge needed)
let files = build_items(&[(1, 5)]);
assert_eq!(merge_seq_files(&files, None), Vec::<MockFile>::new());
// Test the example case: [10, 1, 1, 1] - should merge the last three files
let files = build_items_with_size(&[(1, 2, 10), (3, 4, 1), (5, 6, 1), (7, 8, 1)]);
let result = merge_seq_files(&files, None);
assert_eq!(result.len(), 3);
assert_eq!(result[0].size, 1);
assert_eq!(result[1].size, 1);
assert_eq!(result[2].size, 1);
// Test with files of equal size - should merge as many as possible
let files = build_items_with_size(&[(1, 2, 5), (3, 4, 5), (5, 6, 5), (7, 8, 5)]);
let result = merge_seq_files(&files, Some(20));
assert_eq!(result.len(), 4); // Should merge all 4 files as total size is 20
// Test with max_file_size constraint
let files = build_items_with_size(&[(1, 2, 5), (3, 4, 5), (5, 6, 5), (7, 8, 5)]);
let result = merge_seq_files(&files, Some(10));
assert_eq!(result.len(), 2); // Should merge only 2 files as max size is 10
// Test with uneven file sizes - should prioritize reducing file count
let files = build_items_with_size(&[(1, 2, 2), (3, 4, 3), (5, 6, 4), (7, 8, 10)]);
let result = merge_seq_files(&files, Some(10));
assert_eq!(result.len(), 3); // Should merge the first 3 files (total size 9)
// Test amplification factor prioritization
// Two possible merges: [5, 5] (amp factor 0.5) vs [10, 1, 1] (amp factor 0.83)
let files =
build_items_with_size(&[(1, 2, 5), (3, 4, 5), (5, 6, 10), (7, 8, 1), (9, 10, 1)]);
let result = merge_seq_files(&files, Some(12));
assert_eq!(result.len(), 2);
assert_eq!(result[0].size, 5);
assert_eq!(result[1].size, 5);
// Test with large file preventing merges
let files = build_items_with_size(&[(1, 2, 100), (3, 4, 1), (5, 6, 1), (7, 8, 1)]);
let result = merge_seq_files(&files, Some(10));
assert_eq!(result.len(), 3); // Should merge the last 3 small files
assert_eq!(result[0].size, 1);
assert_eq!(result[1].size, 1);
assert_eq!(result[2].size, 1);
let files = build_items_with_size(&[(1, 2, 100), (3, 4, 20), (5, 6, 20), (7, 8, 20)]);
let result = merge_seq_files(&files, Some(200));
assert_eq!(result.len(), 4);
let files = build_items_with_size(&[(1, 2, 160), (3, 4, 20), (5, 6, 20), (7, 8, 20)]);
let result = merge_seq_files(&files, None);
assert_eq!(result.len(), 3);
assert_eq!(result[0].size, 20);
assert_eq!(result[1].size, 20);
assert_eq!(result[2].size, 20);
let files = build_items_with_size(&[(1, 2, 100), (3, 4, 1)]);
let result = merge_seq_files(&files, Some(200));
assert_eq!(result.len(), 2);
assert_eq!(result[0].size, 100);
assert_eq!(result[1].size, 1);
let files = build_items_with_size(&[(1, 2, 20), (3, 4, 20), (5, 6, 20), (7, 8, 20)]);
let result = merge_seq_files(&files, Some(40));
assert_eq!(result.len(), 2);
assert_eq!(result[0].start, 1);
assert_eq!(result[1].start, 3);
}
}
src/mito2/src/compaction/scheduler.rs
+21 -6
View File
@@ -264,9 +264,13 @@ impl CompactionScheduler {
///
/// # Effects
///
/// Notifies waiters, schedules a pending manual or explicit automatic
/// follow-up, or removes the region status. The returned transition reports
/// a dispatched automatic follow-up or DDLs that are now safe to execute.
/// Notifies waiters, schedules a pending manual or automatic follow-up, or
/// removes the region status. A follow-up is scheduled whenever the cycle
/// latched an automatic trigger, or when `made_progress` is true (the
/// execution removed more files than it added): successful compaction keeps
/// draining the region until the picker returns no plan. The returned
/// transition reports a dispatched automatic follow-up or DDLs that are now
/// safe to execute.
///
/// # Constraints
///
@@ -279,12 +283,18 @@ impl CompactionScheduler {
execution: &CompactionExecution,
manifest_ctx: &ManifestContextRef,
schema_metadata_manager: SchemaMetadataManagerRef,
made_progress: bool,
) -> CompactionTransition {
if !self.is_current_execution(region_id, execution) {
return CompactionTransition::NoAction;
}
self.on_compaction_finished(region_id, manifest_ctx, schema_metadata_manager)
.await
self.on_compaction_finished(
region_id,
manifest_ctx,
schema_metadata_manager,
made_progress,
)
.await
}
/// Completes a cooperatively canceled execution.
@@ -571,6 +581,7 @@ impl CompactionScheduler {
region_id: RegionId,
manifest_ctx: &ManifestContextRef,
schema_metadata_manager: SchemaMetadataManagerRef,
made_progress: bool,
) -> CompactionTransition {
if !self.region_status.contains_key(&region_id) {
return CompactionTransition::NoAction;
@@ -605,7 +616,11 @@ impl CompactionScheduler {
return CompactionTransition::DdlReady(pending_ddl_requests);
}
if status.active.reset_automatic_followup()
// Keep draining when the cycle latched an automatic trigger, or when the
// execution reduced the file count. A no-progress rewrite stops here so a
// split-heavy output cannot loop forever; the next flush trigger resumes.
let should_continue = status.active.reset_automatic_followup() || made_progress;
if should_continue
&& self.schedule_automatic_followup(region_id, manifest_ctx, schema_metadata_manager)
{
return CompactionTransition::AutomaticFollowupScheduled;
src/mito2/src/compaction/scheduler_test.rs
+152 -24
View File
@@ -826,6 +826,7 @@ async fn test_stale_plan_execution_does_not_affect_replacement_status() {
&stale_execution,
&manifest_ctx,
schema_metadata_manager,
true,
)
.await;
assert!(pending_ddls.is_empty());
@@ -907,7 +908,7 @@ async fn test_schedule_compaction_skips_task_exceeding_memory_limit() {
}
#[tokio::test]
async fn test_execution_finished_without_followup_removes_status() {
async fn test_execution_finished_drains_until_no_plan() {
common_telemetry::init_default_ut_logging();
let job_scheduler = Arc::new(VecScheduler::default());
let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone());
@@ -975,23 +976,59 @@ async fn test_execution_finished_without_followup_removes_status() {
&file_metas,
purger.clone(),
);
// A completed execution without an explicit follow-up removes its status.
scheduler
.on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone())
// A completed execution that reduced the file count chains another pick even
// without an explicit trigger, because the layout is still compactable.
let transition = scheduler
.on_compaction_finished(
region_id,
&manifest_ctx,
schema_metadata_manager.clone(),
true,
)
.await;
assert_matches!(transition, CompactionTransition::AutomaticFollowupScheduled);
assert!(scheduler.region_status.contains_key(&region_id));
let finished = recv_compaction_pick_finished(&mut rx).await;
scheduler
.handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone())
.await;
assert_eq!(2, job_scheduler.num_jobs());
// Replace the layout with a single file: the drain must stop once the picker
// finds nothing to do.
let data = version_control.current();
let file_metas: Vec<_> = data.version.ssts.levels()[0]
.files
.values()
.map(|file| file.meta_ref().clone())
.collect();
apply_edit(&version_control, &[(0, end)], &file_metas, purger.clone());
let transition = scheduler
.on_compaction_finished(
region_id,
&manifest_ctx,
schema_metadata_manager.clone(),
true,
)
.await;
assert_matches!(transition, CompactionTransition::AutomaticFollowupScheduled);
let finished = recv_compaction_pick_finished(&mut rx).await;
let transition = scheduler
.handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone())
.await;
assert_matches!(transition, CompactionTransition::NoAction);
assert!(scheduler.region_status.is_empty());
assert_eq!(1, job_scheduler.num_jobs());
assert_eq!(2, job_scheduler.num_jobs());
}
#[tokio::test]
async fn test_time_range_compaction_when_compaction_in_progress() {
async fn test_execution_finished_without_progress_removes_status() {
common_telemetry::init_default_ut_logging();
let job_scheduler = Arc::new(VecScheduler::default());
let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone());
let (tx, mut rx) = mpsc::channel(4);
let mut scheduler = env.mock_compaction_scheduler(tx);
let mut builder = VersionControlBuilder::new();
let purger = builder.file_purger();
let region_id = builder.region_id();
let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager();
@@ -1017,6 +1054,70 @@ async fn test_time_range_compaction_when_compaction_in_progress() {
.push_l0_file(90, end)
.build(),
);
let manifest_ctx = env
.mock_manifest_context(version_control.current().version.metadata.clone())
.await;
let scheduled = scheduler
.schedule_automatic_compaction(
compact_request::Options::Regular(Default::default()),
&version_control,
&env.access_layer,
&manifest_ctx,
schema_metadata_manager.clone(),
)
.unwrap();
assert!(scheduled);
let finished = recv_compaction_pick_finished(&mut rx).await;
scheduler
.handle_compaction_pick_finished(finished, &manifest_ctx, schema_metadata_manager.clone())
.await;
assert_eq!(1, job_scheduler.num_jobs());
// The execution rewrote files without reducing the file count (e.g. its output
// was split into more files than its input). Chaining would loop without making
// progress, so the lifecycle ends here; the next flush trigger resumes compaction.
let transition = scheduler
.on_compaction_finished(
region_id,
&manifest_ctx,
schema_metadata_manager.clone(),
false,
)
.await;
assert_matches!(transition, CompactionTransition::NoAction);
assert!(scheduler.region_status.is_empty());
assert_eq!(1, job_scheduler.num_jobs());
}
#[tokio::test]
async fn test_time_range_compaction_when_compaction_in_progress() {
common_telemetry::init_default_ut_logging();
let job_scheduler = Arc::new(VecScheduler::default());
let env = SchedulerEnv::new().await.scheduler(job_scheduler.clone());
let (tx, mut rx) = mpsc::channel(4);
let mut scheduler = env.mock_compaction_scheduler(tx);
let mut builder = VersionControlBuilder::new();
let purger = builder.file_purger();
let region_id = builder.region_id();
let (schema_metadata_manager, kv_backend) = mock_schema_metadata_manager();
schema_metadata_manager
.register_region_table_info(
builder.region_id().table_id(),
"test_table",
"test_catalog",
"test_schema",
None,
kv_backend,
)
.await;
// 40 files to compact. The first task picks 32, leaving 8 for the pending request.
let end = 1000 * 1000;
for offset in 0..40 {
builder.push_l0_file(offset * 10, end);
}
let version_control = Arc::new(builder.build());
let manifest_ctx = env
.mock_manifest_context(version_control.current().version.metadata.clone())
.await;
@@ -1027,13 +1128,9 @@ async fn test_time_range_compaction_when_compaction_in_progress() {
.map(|file| file.meta_ref().clone())
.collect();
// 5 files for next compaction and removes old files.
apply_edit(
&version_control,
&[(0, end), (20, end), (40, end), (60, end), (80, end)],
&file_metas,
purger.clone(),
);
// Replaces the files before scheduling the first compaction.
let next_files = (0..40).map(|offset| (offset * 20, end)).collect::<Vec<_>>();
apply_edit(&version_control, &next_files, &file_metas, purger.clone());
scheduler
.schedule_automatic_compaction(
@@ -1094,7 +1191,12 @@ async fn test_time_range_compaction_when_compaction_in_progress() {
// On compaction finished and schedule next compaction.
scheduler
.on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone())
.on_compaction_finished(
region_id,
&manifest_ctx,
schema_metadata_manager.clone(),
true,
)
.await;
assert_eq!(1, scheduler.region_status.len());
assert!(
@@ -1285,10 +1387,21 @@ async fn test_ranged_compaction_stops_without_followup() {
assert_eq!(1, job_scheduler.num_jobs());
scheduler
.on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone())
.on_compaction_finished(
region_id,
&manifest_ctx,
schema_metadata_manager.clone(),
true,
)
.await;
assert!(!scheduler.region_status.contains_key(&region_id));
assert_eq!(1, job_scheduler.num_jobs());
// The execution reduced the file count and the remaining window is still
// compactable, so the scheduler drains it with an unrestricted follow-up.
assert!(scheduler.region_status.contains_key(&region_id));
let followup = recv_compaction_pick_finished(&mut rx).await;
scheduler
.handle_compaction_pick_finished(followup, &manifest_ctx, schema_metadata_manager.clone())
.await;
assert_eq!(2, job_scheduler.num_jobs());
}
#[tokio::test]
@@ -1369,7 +1482,12 @@ async fn test_automatic_trigger_during_execution_clears_continuation_scope() {
// Finishing the manual execution must chain an unrestricted follow-up instead of
// continuing with the manual request's time range.
scheduler
.on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone())
.on_compaction_finished(
region_id,
&manifest_ctx,
schema_metadata_manager.clone(),
true,
)
.await;
let continuation = recv_compaction_pick_finished(&mut rx).await;
@@ -1936,7 +2054,7 @@ async fn test_on_compaction_finished_returns_pending_ddl_requests() {
});
let pending_ddls = scheduler
.on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager)
.on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager, true)
.await;
assert_eq!(pending_ddls.len(), 1);
@@ -1985,7 +2103,12 @@ async fn test_planning_terminal_prioritizes_pending_ddl_over_automatic_followup(
assert!(result.is_ok());
let pending_ddls = scheduler
.on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone())
.on_compaction_finished(
region_id,
&manifest_ctx,
schema_metadata_manager.clone(),
true,
)
.await;
assert!(pending_ddls.is_empty());
@@ -2034,7 +2157,7 @@ async fn test_on_compaction_finished_dispatches_pending_ddl_before_chained_regul
});
let pending_ddls = scheduler
.on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager)
.on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager, true)
.await;
assert_eq!(pending_ddls.len(), 1);
@@ -2056,7 +2179,7 @@ async fn test_on_compaction_finished_returns_empty_when_region_absent() {
let (schema_metadata_manager, _kv_backend) = mock_schema_metadata_manager();
let pending_ddls = scheduler
.on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager)
.on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager, true)
.await;
assert!(pending_ddls.is_empty());
@@ -2111,7 +2234,12 @@ async fn test_on_compaction_finished_manual_schedule_error_cleans_status() {
});
let pending_ddls = scheduler
.on_compaction_finished(region_id, &manifest_ctx, schema_metadata_manager.clone())
.on_compaction_finished(
region_id,
&manifest_ctx,
schema_metadata_manager.clone(),
true,
)
.await;
assert!(pending_ddls.is_empty());
src/mito2/src/compaction/twcs.rs
+718 -145
View File
File diff suppressed because it is too large Load Diff
src/mito2/src/config.rs
+10 -2
View File
@@ -102,7 +102,7 @@ pub struct MitoConfig {
pub experimental_compaction_on_exhausted: OnExhaustedPolicy,
// Flush configs:
/// Interval to auto flush a region if it has not flushed yet (default 30 min).
/// Interval to auto flush a region if it has not flushed yet (default 10 min).
#[serde(with = "humantime_serde")]
pub auto_flush_interval: Duration,
/// Global write buffer size threshold to trigger flush.
@@ -210,7 +210,7 @@ impl Default for MitoConfig {
max_background_purges: get_total_cpu_cores(),
experimental_compaction_memory_limit: MemoryLimit::Unlimited,
experimental_compaction_on_exhausted: OnExhaustedPolicy::default(),
auto_flush_interval: Duration::from_secs(30 * 60),
auto_flush_interval: Duration::from_secs(10 * 60),
global_write_buffer_size: ReadableSize::gb(1),
global_write_buffer_reject_size: ReadableSize::gb(2),
default_region_write_buffer_size: ReadableSize::mb(0),
@@ -381,6 +381,14 @@ impl MitoConfig {
mod tests {
use super::*;
#[test]
fn test_default_auto_flush_interval() {
assert_eq!(
Duration::from_secs(10 * 60),
MitoConfig::default().auto_flush_interval
);
}
#[test]
fn test_adjust_sst_metadata_and_prefilter_cache_caps_independently() {
let mut config = MitoConfig::default();
src/mito2/src/engine/apply_staging_manifest_test.rs
+212 -3
View File
@@ -12,20 +12,24 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashSet;
use std::sync::Arc;
use std::{assert_matches, fs};
use api::v1::Rows;
use api::v1::region::{StrictWindow, compact_request};
use common_function::utils::partition_expr_version;
use common_recordbatch::RecordBatches;
use datatypes::arrow::array::AsArray;
use datatypes::arrow::datatypes::Float64Type;
use datatypes::value::Value;
use partition::expr::{PartitionExpr, col};
use store_api::region_engine::{
RegionEngine, RegionRole, RemapManifestsRequest, SettableRegionRoleState,
};
use store_api::region_request::{
ApplyStagingManifestRequest, EnterStagingRequest, RegionFlushRequest, RegionPutRequest,
RegionRequest, StagingPartitionDirective,
ApplyStagingManifestRequest, EnterStagingRequest, RegionCompactRequest, RegionFlushRequest,
RegionPutRequest, RegionRequest, StagingPartitionDirective,
};
use store_api::storage::{FileId, RegionId};
@@ -37,7 +41,9 @@ use crate::manifest::action::{
};
use crate::sst::FormatType;
use crate::sst::file::FileMeta;
use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema};
use crate::test_util::{
CreateRequestBuilder, TestEnv, build_rows, build_rows_for_key, put_rows, rows_schema,
};
fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr {
col(col_name)
@@ -45,6 +51,209 @@ fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr {
.and(col(col_name).lt(Value::Int64(end)))
}
#[tokio::test]
async fn test_apply_staging_manifest_sequence_domain() {
common_telemetry::init_default_ut_logging();
test_apply_staging_manifest_sequence_domain_with_format(false).await;
test_apply_staging_manifest_sequence_domain_with_format(true).await;
}
async fn test_apply_staging_manifest_sequence_domain_with_format(flat_format: bool) {
let mut env = TestEnv::with_prefix("apply-staging-sequence-domain").await;
let engine = env
.create_engine(MitoConfig {
default_flat_format: flat_format,
..Default::default()
})
.await;
let source = RegionId::new(1, 1);
let target = RegionId::new(1, 2);
let request = CreateRequestBuilder::new().build();
let schema = rows_schema(&request);
engine
.handle_request(source, RegionRequest::Create(request.clone()))
.await
.unwrap();
for value in 0..3 {
put_rows(
&engine,
source,
Rows {
schema: schema.clone(),
rows: build_rows_for_key("0", 0, 1, value),
},
)
.await;
}
engine
.handle_request(source, RegionRequest::Flush(RegionFlushRequest::default()))
.await
.unwrap();
let source_manifest = engine
.get_region(source)
.unwrap()
.manifest_ctx
.manifest()
.await;
assert_eq!(source_manifest.files.len(), 1);
assert_eq!(
source_manifest.files.values().next().unwrap().sequence,
Some(std::num::NonZeroU64::new(3).unwrap())
);
engine
.set_region_role_state_gracefully(source, SettableRegionRoleState::StagingLeader)
.await
.unwrap();
let partition_expr = float_range_expr("field_0", 0.1, 100.1)
.as_json_str()
.unwrap();
let result = engine
.remap_manifests(RemapManifestsRequest {
region_id: source,
input_regions: vec![source],
region_mapping: [(source, vec![target])].into_iter().collect(),
new_partition_exprs: [(target, partition_expr.clone())].into_iter().collect(),
})
.await
.unwrap();
engine
.handle_request(target, RegionRequest::Create(request.clone()))
.await
.unwrap();
engine
.handle_request(
target,
RegionRequest::EnterStaging(EnterStagingRequest {
partition_directive: StagingPartitionDirective::UpdatePartitionExpr(
partition_expr.clone(),
),
}),
)
.await
.unwrap();
engine
.handle_request(
target,
RegionRequest::ApplyStagingManifest(ApplyStagingManifestRequest {
partition_expr,
central_region_id: source,
manifest_path: result.manifest_paths[&target].clone(),
}),
)
.await
.unwrap();
let manifest = engine
.get_region(target)
.unwrap()
.manifest_ctx
.manifest()
.await;
assert_eq!(manifest.files.len(), 1);
assert_eq!(manifest.committed_sequence, Some(1));
let imported_file = manifest.files.values().next().unwrap();
assert_eq!(imported_file.region_id, source);
assert_eq!(
imported_file.sequence,
Some(std::num::NonZeroU64::new(1).unwrap())
);
put_rows(
&engine,
target,
Rows {
schema: schema.clone(),
rows: build_rows_for_key("0", 0, 1, 99),
},
)
.await;
assert_target_value(&engine, target, 99.0).await;
engine
.handle_request(target, RegionRequest::Flush(RegionFlushRequest::default()))
.await
.unwrap();
let input_ids = current_file_ids(&engine, target);
assert_eq!(
input_ids.len(),
2,
"imported and target-write SSTs must both exist"
);
engine
.handle_request(
target,
RegionRequest::Compact(RegionCompactRequest {
options: compact_request::Options::StrictWindow(StrictWindow {
window_seconds: 60,
}),
parallelism: None,
time_range: None,
}),
)
.await
.unwrap();
let output_version = engine.get_region(target).unwrap().version();
let output_files = output_version
.ssts
.levels()
.iter()
.flat_map(|level| level.files.values())
.collect::<Vec<_>>();
let output_ids = output_files
.iter()
.map(|file| file.meta_ref().file_id)
.collect::<HashSet<_>>();
assert_eq!(output_ids.len(), 1);
assert!(
output_files
.iter()
.all(|file| file.meta_ref().region_id == target)
);
assert!(
input_ids.iter().all(|id| !output_ids.contains(id)),
"real compaction must replace both input SSTs"
);
assert_target_value(&engine, target, 99.0).await;
}
fn current_file_ids(engine: &crate::engine::MitoEngine, region_id: RegionId) -> HashSet<FileId> {
engine
.get_region(region_id)
.unwrap()
.version()
.ssts
.levels()
.iter()
.flat_map(|level| level.files.values())
.map(|file| file.meta_ref().file_id)
.collect()
}
async fn assert_target_value(
engine: &crate::engine::MitoEngine,
region_id: RegionId,
expected: f64,
) {
let scan = engine
.scan_to_stream(region_id, ScanRequest::default())
.await
.unwrap();
let batches = RecordBatches::try_collect(scan).await.unwrap();
assert_eq!(
batches.iter().map(|batch| batch.num_rows()).sum::<usize>(),
1
);
let batch = batches.iter().next().unwrap();
let values = batch
.column_by_name("field_0")
.unwrap()
.as_primitive::<Float64Type>();
assert_eq!(values.value(0), expected);
}
fn float_range_expr(col_name: &str, start: f64, end: f64) -> PartitionExpr {
col(col_name)
.gt_eq(Value::Float64(start.into()))
src/mito2/src/engine/basic_test.rs
+7 -6
View File
@@ -882,9 +882,9 @@ async fn test_cache_null_primary_key_with_format(flat_format: bool) {
#[tokio::test]
async fn test_list_ssts() {
test_list_ssts_with_format(false, r#"
ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000001/<file_id>.parquet", file_size: 2951, index_file_path: Some("test/11_0000000001/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }
ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000002/<file_id>.parquet", file_size: 2951, index_file_path: Some("test/11_0000000002/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }
ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test/22_0000000042/<file_id>.parquet", file_size: 2951, index_file_path: Some("test/22_0000000042/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"# ,
ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000001/<file_id>.parquet", file_size: 2951, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000001/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }
ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000002/<file_id>.parquet", file_size: 2951, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000002/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }
ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test/22_0000000042/<file_id>.parquet", file_size: 2951, max_row_group_uncompressed_size: 0, index_file_path: Some("test/22_0000000042/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"# ,
r#"
StorageSstEntry { file_path: "test/11_0000000001/<file_id>.parquet", file_size: None, last_modified_ms: None, node_id: None }
StorageSstEntry { file_path: "test/11_0000000001/index/<file_id>.puffin", file_size: None, last_modified_ms: None, node_id: None }
@@ -893,9 +893,9 @@ StorageSstEntry { file_path: "test/11_0000000002/index/<file_id>.puffin", file_s
StorageSstEntry { file_path: "test/22_0000000042/<file_id>.parquet", file_size: None, last_modified_ms: None, node_id: None }
StorageSstEntry { file_path: "test/22_0000000042/index/<file_id>.puffin", file_size: None, last_modified_ms: None, node_id: None }"#).await;
test_list_ssts_with_format(true, r#"
ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000001/<file_id>.parquet", file_size: 3365, index_file_path: Some("test/11_0000000001/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }
ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000002/<file_id>.parquet", file_size: 3365, index_file_path: Some("test/11_0000000002/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }
ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test/22_0000000042/<file_id>.parquet", file_size: 3365, index_file_path: Some("test/22_0000000042/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"#,
ManifestSstEntry { table_dir: "test/", region_id: 47244640257(11, 1), table_id: 11, region_number: 1, region_group: 0, region_sequence: 1, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000001/<file_id>.parquet", file_size: 3365, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000001/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640257(11, 1), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }
ManifestSstEntry { table_dir: "test/", region_id: 47244640258(11, 2), table_id: 11, region_number: 2, region_group: 0, region_sequence: 2, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test/11_0000000002/<file_id>.parquet", file_size: 3365, max_row_group_uncompressed_size: 0, index_file_path: Some("test/11_0000000002/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 47244640258(11, 2), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }
ManifestSstEntry { table_dir: "test/", region_id: 94489280554(22, 42), table_id: 22, region_number: 42, region_group: 0, region_sequence: 42, file_id: "<file_id>", index_version: 0, level: 0, file_path: "test/22_0000000042/<file_id>.parquet", file_size: 3365, max_row_group_uncompressed_size: 0, index_file_path: Some("test/22_0000000042/index/<file_id>.puffin"), index_file_size: Some(250), num_rows: 10, num_row_groups: 1, num_series: Some(1), min_ts: 0::Millisecond, max_ts: 9000::Millisecond, sequence: Some(10), partition_expr: None, origin_region_id: 94489280554(22, 42), node_id: None, visible: true, primary_key_min: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01"), primary_key_max: Some(b"\x01\x01x\0\0\0\0\0\0\0\x01") }"#,
r#"
StorageSstEntry { file_path: "test/11_0000000001/<file_id>.parquet", file_size: None, last_modified_ms: None, node_id: None }
StorageSstEntry { file_path: "test/11_0000000001/index/<file_id>.puffin", file_size: None, last_modified_ms: None, node_id: None }
@@ -967,6 +967,7 @@ async fn test_list_ssts_with_format(
.map(|p| p.replace(&e.file_id, "<file_id>"));
e.file_id = "<file_id>".to_string();
e.index_version = 0;
e.max_row_group_uncompressed_size = 0;
format!("\n{:?}", e)
})
.sorted()
src/mito2/src/engine/compaction_test.rs
+16 -9
View File
@@ -768,6 +768,7 @@ async fn test_truncate_waits_for_non_cancellable_compaction_commit() {
.await;
let create = CreateRequestBuilder::new()
.insert_option("compaction.type", "twcs")
.insert_option("compaction.twcs.trigger_file_num", "4")
.build();
let column_schemas = create
.column_metadatas
@@ -779,7 +780,9 @@ async fn test_truncate_waits_for_non_cancellable_compaction_commit() {
.await
.unwrap();
put_and_flush(&engine, region_id, &column_schemas, 0..10).await;
put_and_flush(&engine, region_id, &column_schemas, 5..20).await;
put_and_flush(&engine, region_id, &column_schemas, 10..20).await;
put_and_flush(&engine, region_id, &column_schemas, 20..30).await;
put_and_flush(&engine, region_id, &column_schemas, 30..40).await;
let commit_guard = gate.arm_commit();
let compact_engine = engine.clone();
@@ -803,7 +806,7 @@ async fn test_truncate_waits_for_non_cancellable_compaction_commit() {
RegionRequest::Truncate(RegionTruncateRequest::ByTimeRanges {
time_ranges: vec![(
Timestamp::new_millisecond(0),
Timestamp::new_millisecond(19_000),
Timestamp::new_millisecond(39_000),
)],
}),
)
@@ -981,12 +984,10 @@ async fn test_compaction_region_with_format(flat_format: bool) {
// [20....29]
// -[15.........29]- (delete)
// [15.....24]
// Output:
// [0..9]
// [10............29] (contains delete)
// [15....24]
// Count-first compaction consumes the first 4 SSTs as soon as the trigger is reached.
// The compacted output and final flush leave 2 SSTs.
assert_eq!(
3,
2,
scanner.num_files(),
"unexpected files: {:?}",
scanner.file_ids()
@@ -1530,6 +1531,7 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) {
let request = CreateRequestBuilder::new()
.insert_option("compaction.type", "twcs")
.insert_option("compaction.twcs.trigger_file_num", "4")
.build();
let column_schemas = request
@@ -1542,7 +1544,9 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) {
.await
.unwrap();
let listener_guard = CompactionListenerGuard::new(listener.clone());
// Flush 2 SSTs for compaction.
// Flush 4 balanced SSTs for compaction.
put_and_flush(&engine, region_id, &column_schemas, 0..10).await;
put_and_flush(&engine, region_id, &column_schemas, 5..20).await;
put_and_flush(&engine, region_id, &column_schemas, 0..10).await;
put_and_flush(&engine, region_id, &column_schemas, 5..20).await;
@@ -1573,7 +1577,7 @@ async fn test_readonly_during_compaction_with_format(flat_format: bool) {
.await
.unwrap();
assert_eq!(
2,
4,
scanner.num_files(),
"unexpected files: {:?}",
scanner.file_ids()
@@ -1614,6 +1618,7 @@ async fn test_local_compaction_cancellation_notifies_before_pending_ddl_dispatch
let request = CreateRequestBuilder::new()
.insert_option("compaction.type", "twcs")
.insert_option("compaction.twcs.trigger_file_num", "4")
.build();
let column_schemas = request
.column_metadatas
@@ -1628,6 +1633,8 @@ async fn test_local_compaction_cancellation_notifies_before_pending_ddl_dispatch
let merge_guard = gate.arm_merge();
put_and_flush(&engine, region_id, &column_schemas, 0..10).await;
put_and_flush(&engine, region_id, &column_schemas, 5..20).await;
put_and_flush(&engine, region_id, &column_schemas, 0..10).await;
put_and_flush(&engine, region_id, &column_schemas, 5..20).await;
tokio::time::timeout(Duration::from_secs(5), gate.wait_until_merge_entered())
.await
src/mito2/src/engine/merge_mode_test.rs
+2 -1
View File
@@ -127,6 +127,7 @@ async fn test_merge_mode_compaction_with_format(flat_format: bool) {
let request = CreateRequestBuilder::new()
.field_num(2)
.insert_option("compaction.type", "twcs")
.insert_option("compaction.twcs.trigger_file_num", "3")
.insert_option("merge_mode", "last_non_null")
.build();
let table_dir = request.table_dir.clone();
@@ -208,7 +209,7 @@ async fn test_merge_mode_compaction_with_format(flat_format: bool) {
.scanner(region_id, ScanRequest::default())
.await
.unwrap();
assert_eq!(2, scanner.num_files());
assert_eq!(1, scanner.num_files());
assert_eq!(1, scanner.num_memtables());
scanner.set_target_partitions(2);
let stream = scanner.scan().await.unwrap();
src/mito2/src/gc/worker_test.rs
+7 -3
View File
@@ -399,7 +399,8 @@ async fn test_gc_worker_basic_compact() {
let region = engine.get_region(region_id).unwrap();
let manifest = region.manifest_ctx.manifest().await;
assert_eq!(manifest.removed_files.removed_files[0].files.len(), 3);
let removed_file_num = manifest.removed_files.removed_files[0].files.len();
assert!(removed_file_num > 0);
let version = manifest.manifest_version;
@@ -413,7 +414,10 @@ async fn test_gc_worker_basic_compact() {
let gc_worker = create_gc_worker(&engine, regions, &file_ref_manifest, true).await;
let report = gc_worker.run().await.unwrap();
assert_eq!(report.deleted_files.get(&region_id).unwrap().len(), 3,);
assert_eq!(
report.deleted_files.get(&region_id).unwrap().len(),
removed_file_num
);
assert!(report.need_retry_regions.is_empty());
}
@@ -476,7 +480,7 @@ async fn test_gc_worker_compact_with_ref() {
let region = engine.get_region(region_id).unwrap();
let manifest = region.manifest_ctx.manifest().await;
assert_eq!(manifest.removed_files.removed_files[0].files.len(), 3);
assert!(!manifest.removed_files.removed_files[0].files.is_empty());
let version = manifest.manifest_version;
src/mito2/src/region.rs
+2
View File
@@ -839,6 +839,7 @@ impl MitoRegion {
level: meta.level,
file_path: sst_file_path(table_dir, meta.file_id(), path_type),
file_size: meta.file_size,
max_row_group_uncompressed_size: meta.max_row_group_uncompressed_size,
index_file_path,
index_file_size,
num_rows: meta.num_rows,
@@ -847,6 +848,7 @@ impl MitoRegion {
min_ts: meta.time_range.0,
max_ts: meta.time_range.1,
sequence: meta.sequence.map(|s| s.get()),
partition_expr: meta.partition_expr.as_ref().map(ToString::to_string),
origin_region_id,
node_id: None,
visible,
src/mito2/src/sst/index.rs
+77 -9
View File
@@ -706,6 +706,14 @@ pub struct IndexBuildTask {
pub region_id: RegionId,
/// The SST file handle to build index for.
pub file: FileHandle,
/// The target region metadata used to decode rows from the SST.
///
/// An SST may originate in another region while being visible in the target
/// manifest. This metadata defines the target schema and sequence domain;
/// applying the staging manifest only makes imported files visible. Index
/// rebuild happens later when a flush, compaction, schema change, or manual
/// index build request schedules it.
pub(crate) target_region_metadata: RegionMetadataRef,
/// The manifest state this build is based on.
pub(crate) source: IndexBuildSource,
pub reason: IndexBuildType,
@@ -846,6 +854,7 @@ impl IndexBuildTask {
let mut parquet_reader = self
.access_layer
.read_sst(self.file.clone()) // use the latest file handle instead of creating a new one
.expected_metadata(Some(self.target_region_metadata.clone()))
.build()
.await?;
@@ -1494,8 +1503,12 @@ mod tests {
use datatypes::schema::{
ColumnSchema, FulltextOptions, SkippingIndexOptions, SkippingIndexType,
};
use datatypes::value::Value;
use index::inverted_index::format::reader::InvertedIndexReader;
use object_store::ObjectStore;
use object_store::services::Memory;
use partition::expr::col;
use puffin::puffin_manager::{PuffinManager, PuffinReader};
use puffin_manager::PuffinManagerFactory;
use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
use tokio::sync::mpsc;
@@ -2088,6 +2101,7 @@ mod tests {
let task = IndexBuildTask {
region_id,
file,
target_region_metadata: version_control.current().version.metadata.clone(),
source: IndexBuildSource::new(
file_meta,
version_control.current().version.metadata.schema_version,
@@ -2120,16 +2134,34 @@ mod tests {
}
#[tokio::test]
async fn test_index_build_task_increments_legacy_index_version() {
async fn test_index_build_task_foreign_file_uses_target_metadata() {
let env = SchedulerEnv::new().await;
let mut scheduler = env.mock_index_build_scheduler(4);
let metadata = Arc::new(sst_region_metadata());
let manifest_ctx = env.mock_manifest_context(metadata.clone()).await;
let region_id = metadata.region_id;
let source_metadata = Arc::new(sst_region_metadata());
let mut target_metadata = (*source_metadata).clone();
target_metadata.region_id = RegionId::new(1, 3);
let mut target_builder = RegionMetadataBuilder::new(target_metadata.region_id);
for mut column_metadata in target_metadata.column_metadatas.clone() {
if column_metadata.column_id == 2 {
column_metadata.column_schema =
column_metadata.column_schema.with_inverted_index(true);
}
target_builder.push_column_metadata(column_metadata);
}
let partition_expr = col("field_0")
.gt_eq(Value::UInt64(100))
.and(col("field_0").lt(Value::UInt64(200)));
target_builder
.primary_key(target_metadata.primary_key.clone())
.partition_expr_json(Some(partition_expr.as_json_str().unwrap()))
.bump_version();
let target_metadata = Arc::new(target_builder.build().unwrap());
let manifest_ctx = env.mock_manifest_context(target_metadata.clone()).await;
let region_id = target_metadata.region_id;
let file_purger = Arc::new(NoopFilePurger {});
let sst_info = mock_sst_file(metadata.clone(), &env, IndexBuildMode::Async).await;
let sst_info = mock_sst_file(source_metadata.clone(), &env, IndexBuildMode::Async).await;
let file_meta = FileMeta {
region_id,
region_id: source_metadata.region_id,
file_id: sst_info.file_id,
file_size: sst_info.file_size,
max_row_group_uncompressed_size: sst_info.max_row_group_uncompressed_size,
@@ -2144,8 +2176,8 @@ mod tests {
seed_manifest_file(&manifest_ctx, &file_meta).await;
let files = HashMap::from([(file_meta.file_id, file_meta.clone())]);
let version_control =
mock_version_control(metadata.clone(), file_purger.clone(), files).await;
let indexer_builder = mock_indexer_builder(metadata.clone(), &env).await;
mock_version_control(target_metadata.clone(), file_purger.clone(), files).await;
let indexer_builder = mock_indexer_builder(target_metadata.clone(), &env).await;
let file = FileHandle::new(file_meta.clone(), file_purger.clone());
@@ -2155,6 +2187,7 @@ mod tests {
let task = IndexBuildTask {
region_id,
file,
target_region_metadata: version_control.current().version.metadata.clone(),
source: IndexBuildSource::new(
file_meta.clone(),
version_control.current().version.metadata.schema_version,
@@ -2199,9 +2232,40 @@ mod tests {
assert!(updated_meta.index_file_size > 0);
assert_eq!(updated_meta.file_id, file_meta.file_id);
assert_eq!(updated_meta.index_version, 1);
let field_0_index = updated_meta
.indexes
.iter()
.find(|index| index.column_id == 2)
.expect("field_0 should have an inverted index");
assert_eq!(
field_0_index.created_indexes.as_slice(),
[IndexType::InvertedIndex]
);
}
_ => panic!("Unexpected worker request: {:?}", worker_req),
}
let puffin_reader = env
.access_layer
.build_puffin_manager()
.reader(&RegionIndexId::new(
RegionFileId::new(source_metadata.region_id, file_meta.file_id),
1,
))
.await
.unwrap();
let blob = puffin_reader
.blob(inverted_index::INDEX_BLOB_TYPE)
.await
.unwrap();
let blob_reader = blob.reader().await.unwrap();
let index_metadata =
index::inverted_index::format::reader::InvertedIndexBlobReader::new(blob_reader)
.metadata(None)
.await
.unwrap();
assert!(index_metadata.metas.contains_key("2"));
assert_eq!(index_metadata.total_row_count, 100);
}
async fn schedule_index_build_task_with_mode(build_mode: IndexBuildMode) {
@@ -2236,6 +2300,7 @@ mod tests {
let task = IndexBuildTask {
region_id,
file,
target_region_metadata: version_control.current().version.metadata.clone(),
source: IndexBuildSource::new(
file_meta.clone(),
version_control.current().version.metadata.schema_version,
@@ -2346,6 +2411,7 @@ mod tests {
let task = IndexBuildTask {
region_id,
file,
target_region_metadata: version_control.current().version.metadata.clone(),
source: IndexBuildSource::new(
file_meta.clone(),
version_control.current().version.metadata.schema_version,
@@ -2444,6 +2510,7 @@ mod tests {
let task = IndexBuildTask {
region_id,
file,
target_region_metadata: version_control.current().version.metadata.clone(),
source: IndexBuildSource::new(
file_meta.clone(),
version_control.current().version.metadata.schema_version,
@@ -2505,7 +2572,7 @@ mod tests {
let schema_version = metadata.schema_version;
let manifest_ctx = env.mock_manifest_context(metadata.clone()).await;
let file_purger = Arc::new(NoopFilePurger {});
let indexer_builder = mock_indexer_builder(metadata, env).await;
let indexer_builder = mock_indexer_builder(metadata.clone(), env).await;
let (tx, _rx) = mpsc::channel(4);
let (result_tx, result_rx) = mpsc::channel::<Result<IndexBuildOutcome>>(4);
@@ -2521,6 +2588,7 @@ mod tests {
let task = IndexBuildTask {
region_id,
file,
target_region_metadata: metadata,
source: IndexBuildSource::new(file_meta, schema_version),
reason,
access_layer: env.access_layer.clone(),
src/mito2/src/sst/parquet.rs
+157 -70
View File
@@ -118,8 +118,8 @@ mod tests {
use datafusion_expr::{BinaryExpr, Expr, Literal, Operator, col, lit};
use datatypes::arrow;
use datatypes::arrow::array::{
ArrayRef, BinaryDictionaryBuilder, RecordBatch, StringArray, StringDictionaryBuilder,
TimestampMillisecondArray, UInt8Array, UInt64Array,
ArrayRef, AsArray, BinaryDictionaryBuilder, RecordBatch, StringArray,
StringDictionaryBuilder, TimestampMillisecondArray, UInt8Array, UInt64Array,
};
use datatypes::arrow::datatypes::{DataType, Field, Schema, UInt32Type};
use datatypes::arrow::util::pretty::pretty_format_batches;
@@ -1625,91 +1625,178 @@ mod tests {
#[tokio::test]
async fn test_read_with_override_sequence() {
test_read_with_override_sequence_with_format(false).await;
test_read_with_override_sequence_with_format(true).await;
}
async fn test_read_with_override_sequence_with_format(flat_format: bool) {
let mut env = TestEnv::new().await;
let object_store = env.init_object_store_manager();
let handle = sst_file_handle(0, 1000);
let file_path = FixedPathProvider {
region_file_id: handle.file_id(),
};
let metadata = Arc::new(sst_region_metadata());
// Create batches with sequence 0 to trigger override functionality.
let source = new_flat_source_from_record_batches(vec![
new_record_batch_with_custom_sequence(&["a", "d"], 0, 60, 0),
new_record_batch_with_custom_sequence(&["b", "f"], 0, 40, 0),
]);
async fn read_sequences(builder: ParquetReaderBuilder) -> Vec<u64> {
let mut reader = builder.build().await.unwrap().unwrap();
let mut sequences = Vec::new();
while let Some(batch) = reader.next_record_batch().await.unwrap() {
let sequence = batch
.column(batch.num_columns() - 2)
.as_primitive::<datatypes::arrow::datatypes::UInt64Type>();
sequences.extend((0..sequence.len()).map(|idx| sequence.value(idx)));
}
sequences
}
let write_opts = WriteOptions {
row_group_size: 50,
..Default::default()
};
async fn write_sst(
object_store: ObjectStore,
metadata: Arc<RegionMetadata>,
handle: FileHandle,
flat_format: bool,
sequence: u64,
) {
let file_path = FixedPathProvider {
region_file_id: handle.file_id(),
};
let source = new_flat_source_from_record_batches(vec![
new_record_batch_with_custom_sequence(&["a", "d"], 0, 60, sequence),
new_record_batch_with_custom_sequence(&["b", "f"], 0, 40, sequence),
]);
let write_opts = WriteOptions {
row_group_size: 50,
..Default::default()
};
let mut metrics = Metrics::new(WriteType::Flush);
let mut writer = ParquetWriter::new_with_object_store(
object_store,
metadata,
IndexConfig::default(),
NoopIndexBuilder,
file_path,
&mut metrics,
)
.await;
if flat_format {
writer
.write_all_flat(source, None, &write_opts)
.await
.unwrap();
} else {
writer
.write_all_flat_as_primary_key(source, None, &write_opts)
.await
.unwrap();
}
}
let mut metrics = Metrics::new(WriteType::Flush);
let mut writer = ParquetWriter::new_with_object_store(
let custom_sequence = 12345u64;
let local_zero_handle = sst_file_handle(0, 1000);
write_sst(
object_store.clone(),
metadata.clone(),
IndexConfig::default(),
NoopIndexBuilder,
file_path,
&mut metrics,
local_zero_handle.clone(),
flat_format,
0,
)
.await;
writer
.write_all_flat_as_primary_key(source, None, &write_opts)
.await
.unwrap()
.remove(0);
// Local all-zero SSTs retain the compatibility override.
let local_zero_none = read_sequences(
ParquetReaderBuilder::new(
FILE_DIR.to_string(),
PathType::Bare,
local_zero_handle.clone(),
object_store.clone(),
)
.expected_metadata(Some(metadata.clone())),
)
.await;
assert!(local_zero_none.iter().all(|sequence| *sequence == 0));
// Read without override sequence (should read sequence 0)
let builder = ParquetReaderBuilder::new(
FILE_DIR.to_string(),
PathType::Bare,
handle.clone(),
object_store.clone(),
);
let mut reader = builder.build().await.unwrap().unwrap();
let mut normal_batches = Vec::new();
while let Some(batch) = reader.next_record_batch().await.unwrap() {
normal_batches.push(batch);
}
// Read with override sequence using FileMeta.sequence
let custom_sequence = 12345u64;
let file_meta = handle.meta_ref();
let mut override_file_meta = file_meta.clone();
override_file_meta.sequence = Some(std::num::NonZero::new(custom_sequence).unwrap());
let override_handle = FileHandle::new(
override_file_meta,
let mut local_zero_meta = local_zero_handle.meta_ref().clone();
local_zero_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap());
let local_zero_override_handle = FileHandle::new(
local_zero_meta,
Arc::new(crate::sst::file_purger::NoopFilePurger),
);
let builder = ParquetReaderBuilder::new(
FILE_DIR.to_string(),
PathType::Bare,
override_handle,
object_store.clone(),
let local_zero_override = read_sequences(
ParquetReaderBuilder::new(
FILE_DIR.to_string(),
PathType::Bare,
local_zero_override_handle,
object_store.clone(),
)
.expected_metadata(Some(metadata.clone())),
)
.await;
assert!(
local_zero_override
.iter()
.all(|sequence| *sequence == custom_sequence)
);
let mut reader = builder.build().await.unwrap().unwrap();
let mut override_batches = Vec::new();
while let Some(batch) = reader.next_record_batch().await.unwrap() {
override_batches.push(batch);
}
// Compare the results
assert_eq!(normal_batches.len(), override_batches.len());
for (normal, override_batch) in normal_batches.into_iter().zip(override_batches.iter()) {
let expected_batch = {
let mut columns = normal.columns().to_vec();
let num_cols = columns.len();
columns[num_cols - 2] =
Arc::new(UInt64Array::from_value(custom_sequence, normal.num_rows()));
RecordBatch::try_new(normal.schema(), columns).unwrap()
};
let local_nonzero_handle = sst_file_handle(0, 1000);
write_sst(
object_store.clone(),
metadata.clone(),
local_nonzero_handle.clone(),
flat_format,
7,
)
.await;
// Override batch should match expected batch
assert_eq!(*override_batch, expected_batch);
}
// Local nonzero SSTs retain physical per-row sequences, even with FileMeta.sequence.
let mut local_nonzero_meta = local_nonzero_handle.meta_ref().clone();
local_nonzero_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap());
let local_nonzero_override_handle = FileHandle::new(
local_nonzero_meta,
Arc::new(crate::sst::file_purger::NoopFilePurger),
);
let local_nonzero_override = read_sequences(
ParquetReaderBuilder::new(
FILE_DIR.to_string(),
PathType::Bare,
local_nonzero_override_handle,
object_store.clone(),
)
.expected_metadata(Some(metadata.clone())),
)
.await;
assert!(local_nonzero_override.iter().all(|sequence| *sequence == 7));
// None never overrides a local nonzero physical sequence.
let local_nonzero_none = read_sequences(
ParquetReaderBuilder::new(
FILE_DIR.to_string(),
PathType::Bare,
local_nonzero_handle.clone(),
object_store.clone(),
)
.expected_metadata(Some(metadata.clone())),
)
.await;
assert!(local_nonzero_none.iter().all(|sequence| *sequence == 7));
// A source-owned handle is foreign when read against target metadata, so the
// target-local manifest barrier is applied even for nonzero physical sequences.
let mut target_metadata = (*metadata).clone();
target_metadata.region_id = RegionId::new(0, 1);
let target_metadata = Arc::new(target_metadata);
let mut foreign_meta = local_nonzero_handle.meta_ref().clone();
foreign_meta.sequence = Some(std::num::NonZeroU64::new(custom_sequence).unwrap());
let foreign_handle = FileHandle::new(
foreign_meta,
Arc::new(crate::sst::file_purger::NoopFilePurger),
);
let foreign = read_sequences(
ParquetReaderBuilder::new(
FILE_DIR.to_string(),
PathType::Bare,
foreign_handle,
object_store,
)
.expected_metadata(Some(target_metadata)),
)
.await;
assert!(foreign.iter().all(|sequence| *sequence == custom_sequence));
}
#[tokio::test]
src/mito2/src/sst/parquet/reader.rs
+19 -2
View File
@@ -22,7 +22,7 @@ use std::time::{Duration, Instant};
use api::v1::SemanticType;
use common_recordbatch::filter::{SimpleFilterEvaluator, TimestampUnitCast};
use common_telemetry::{error, tracing, warn};
use common_telemetry::{debug, error, tracing, warn};
use datafusion::physical_plan::PhysicalExpr;
use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
use datafusion_expr::utils::expr_to_columns;
@@ -471,7 +471,24 @@ impl ParquetReaderBuilder {
&file_path,
skip_auto_convert,
)?;
if need_override_sequence(&parquet_meta) {
// `region_meta` comes from the Parquet/source file and must not be used as the
// target identity. When the caller has no current metadata, the handle is the
// only local identity available and therefore denotes a local read.
let expected_region_id = self
.expected_metadata
.as_ref()
.map(|metadata| metadata.region_id)
.unwrap_or(self.file_handle.region_id());
let is_foreign = self.file_handle.region_id() != expected_region_id;
if is_foreign {
debug!(
"Reading foreign SST, file_id: {}, source_region_id: {}, expected_region_id: {}",
self.file_handle.file_id().file_id(),
self.file_handle.region_id(),
expected_region_id,
);
}
if is_foreign || need_override_sequence(&parquet_meta) {
read_format
.set_override_sequence(self.file_handle.meta_ref().sequence.map(|x| x.get()));
}
src/mito2/src/worker/handle_compaction.rs
+6
View File
@@ -131,6 +131,11 @@ impl<S> RegionWorkerLoop<S> {
return;
}
let execution = request.execution.clone();
// Whether this execution reduced the SST file count. The scheduler keeps
// draining while compaction makes progress; a rewrite that did not reduce
// files (e.g. its output was split into more files than its input) must end
// the chain to avoid a no-progress loop.
let made_progress = request.edit.files_to_remove.len() > request.edit.files_to_add.len();
region.version_control.apply_edit(
Some(request.edit.clone()),
@@ -167,6 +172,7 @@ impl<S> RegionWorkerLoop<S> {
&execution,
&region.manifest_ctx,
self.schema_metadata_manager.clone(),
made_progress,
)
.await;
match transition {
src/mito2/src/worker/handle_rebuild_index.rs
+1
View File
@@ -79,6 +79,7 @@ impl<S> RegionWorkerLoop<S> {
IndexBuildTask {
region_id: region.region_id,
file: file.clone(),
target_region_metadata: version.metadata.clone(),
source: IndexBuildSource::new(file_meta, version.metadata.schema_version),
reason: build_type,
access_layer: access_layer.clone(),
src/operator/src/insert.rs
+53 -15
View File
@@ -29,12 +29,15 @@ use api::v1::{
use catalog::CatalogManagerRef;
use client::{OutputData, OutputMeta};
use common_catalog::consts::{
PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, TRACE_ID_COLUMN, TRACE_TABLE_NAME,
TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table,
DEFAULT_PRIVATE_SCHEMA_NAME, PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, TRACE_ID_COLUMN,
TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table,
trace_operations_table_name, trace_services_table_name,
};
use common_event_recorder::DEFAULT_EVENTS_TABLE_NAME;
use common_frontend::slow_query_event::SLOW_QUERY_TABLE_NAME;
use common_grpc_expr::util::ColumnExpr;
use common_meta::cache::TableFlownodeSetCacheRef;
use common_meta::datanode::REGION_STATS_HISTORY_TABLE_NAME;
use common_meta::node_manager::{AffectedRows, NodeManagerRef};
use common_meta::peer::Peer;
use common_meta::rpc::ddl::TriggerReason;
@@ -490,9 +493,9 @@ impl Inserter {
Ok(inserts)
}
/// Returns `None` if auto table creation is allowed, or `Some(reason)` if
/// disabled by either the global config or the request hint. The reason tells
/// which one, for a clearer error.
/// Returns `Some(reason)` if the config or request hint disables automatic
/// table creation. Exempt private system tables are handled by
/// [`Self::is_auto_create_exempt_private_table`].
fn auto_create_disabled_reason(&self, ctx: &QueryContextRef) -> Result<Option<&'static str>> {
let auto_create_table_hint = ctx
.extension(AUTO_CREATE_TABLE_KEY)
@@ -514,6 +517,16 @@ impl Inserter {
})
}
/// Returns whether a private system table may infer and reconcile its schema
/// even when automatic table creation is disabled.
fn is_auto_create_exempt_private_table(schema: &str, table: &str) -> bool {
schema == DEFAULT_PRIVATE_SCHEMA_NAME
&& matches!(
table,
DEFAULT_EVENTS_TABLE_NAME | SLOW_QUERY_TABLE_NAME | REGION_STATS_HISTORY_TABLE_NAME
)
}
/// Ensures a trace table has the request-global schema without requiring a
/// padded data row to drive on-demand creation or alteration. When
/// `alter_existing` is false, a table created after planning is left for the
@@ -571,12 +584,21 @@ impl Inserter {
let _timer = crate::metrics::CREATE_ALTER_ON_DEMAND
.with_label_values(&[auto_create_table_type.as_str()])
.start_timer();
let catalog = ctx.current_catalog();
let schema = ctx.current_schema();
let auto_create_disabled_reason = self.auto_create_disabled_reason(ctx)?;
// Enabled batches permit every table, so only disabled batches need a whitelist scan.
let has_auto_create_exempt_table = auto_create_disabled_reason.is_some()
&& requests
.inserts
.iter()
.any(|req| Self::is_auto_create_exempt_private_table(&schema, &req.table_name));
let mut table_infos = HashMap::new();
if let Some(disabled_reason) = self.auto_create_disabled_reason(ctx)? {
// Without exempt tables, verify existing tables and reject missing ones without inferring schemas.
if let Some(disabled_reason) = auto_create_disabled_reason
&& !has_auto_create_exempt_table
{
let mut instant_table_ids = HashSet::new();
for req in &requests.inserts {
let table = match self.get_table(catalog, &schema, &req.table_name).await? {
@@ -618,20 +640,25 @@ impl Inserter {
let mut per_table_semantics: Option<Option<PerTableSemanticIndex>> = None;
for req in &mut requests.inserts {
// Mixed batches need a per-table decision so an exempt table cannot authorize others.
let auto_create_allowed = auto_create_disabled_reason.is_none()
|| Self::is_auto_create_exempt_private_table(&schema, &req.table_name);
match self.get_table(catalog, &schema, &req.table_name).await? {
Some(table) => {
let table_info = table.table_info();
if table_info.is_ttl_instant_table() {
instant_table_ids.insert(table_info.table_id());
}
if let Some(alter_expr) = self.get_alter_table_expr_on_demand(
req,
&table,
ctx,
accommodate_existing_schema,
is_single_value,
auto_create_table_type.alter_existing(),
)? {
if auto_create_allowed
&& let Some(alter_expr) = self.get_alter_table_expr_on_demand(
req,
&table,
ctx,
accommodate_existing_schema,
is_single_value,
auto_create_table_type.alter_existing(),
)?
{
alter_tables.push(alter_expr);
need_refresh_table_infos.insert((
catalog.to_string(),
@@ -655,6 +682,17 @@ impl Inserter {
}
table_infos.insert(table_info.table_id(), table_info);
}
None if !auto_create_allowed
&& let Some(disabled_reason) = auto_create_disabled_reason =>
{
return InvalidInsertRequestSnafu {
reason: format!(
"Table `{}` does not exist, and {}",
req.table_name, disabled_reason,
),
}
.fail();
}
None => {
let semantic_index = per_table_semantics
.get_or_insert_with(|| parse_per_table_semantic_index(ctx))
src/operator/src/statement.rs
+1 -1
View File
@@ -84,7 +84,7 @@ use table::table_reference::TableReference;
pub use self::admin::{
AdminEventRecorderHandle, AdminFunctionLayer, AdminFunctionLayerRef,
AdminFunctionRecordingLayer, AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService,
AdminFunctionServiceRef,
AdminFunctionServiceRef, admin_output_schema,
};
use self::set::{
set_bytea_output, set_datestyle, set_intervalstyle, set_timezone, validate_client_encoding,
src/operator/src/statement/admin.rs
+114 -57
View File
@@ -19,6 +19,7 @@ use std::sync::Arc;
use common_function::function::FunctionContext;
use common_function::function_registry::{FUNCTION_REGISTRY, get_admin_function};
use common_function::state::FunctionState;
use common_query::Output;
use common_recordbatch::{RecordBatch, RecordBatches};
use common_sql::convert::sql_value_to_value;
@@ -77,6 +78,103 @@ struct CoreAdminFunctionService {
query_engine: query::QueryEngineRef,
}
/// Parts of an `ADMIN` call needed both for execution and schema derivation.
struct ResolvedAdminFunction {
admin_udf: datafusion_expr::ScalarUDF,
fn_name: String,
args: Vec<VectorRef>,
arg_types: Vec<ArrowDataType>,
ret_type: ArrowDataType,
}
/// Resolves the function, parses its literal arguments and derives its
/// return type, without executing it.
fn resolve_admin_function(
stmt: &Admin,
query_ctx: &QueryContextRef,
state: Arc<FunctionState>,
) -> Result<ResolvedAdminFunction> {
let Admin::Func(func) = stmt;
// the function name should be in lower case.
let func_name = func.name.to_string().to_lowercase();
let factory = get_admin_function(&func_name)
.or_else(|| FUNCTION_REGISTRY.get_function(&func_name))
.context(error::AdminFunctionNotFoundSnafu {
name: func_name.clone(),
})?;
let func_ctx = FunctionContext {
query_ctx: query_ctx.clone(),
state,
};
let admin_udf = factory.provide(func_ctx);
admin_udf
.as_async()
.context(error::AdminFunctionNotFoundSnafu { name: func_name })?;
let fn_name = admin_udf.name().to_string();
let signature = admin_udf.signature();
// Parse function arguments
let FunctionArguments::List(args) = &func.args else {
return error::BuildAdminFunctionArgsSnafu {
msg: format!("unsupported function args {} for {}", func.args, fn_name),
}
.fail();
};
let arg_values = args
.args
.iter()
.map(|arg| {
let FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) = arg else {
return error::BuildAdminFunctionArgsSnafu {
msg: format!("unsupported function arg {arg} for {}", fn_name),
}
.fail();
};
Ok(&value.value)
})
.collect::<Result<Vec<_>>>()?;
let args = args_to_vector(&signature.type_signature, &arg_values, query_ctx)?;
let arg_types = args
.iter()
.map(|arg| arg.data_type().as_arrow_type())
.collect::<Vec<_>>();
let ret_type =
admin_udf
.return_type(&arg_types)
.map_err(|e| error::Error::BuildAdminFunctionArgs {
msg: format!(
"Failed to get return type of admin function {}: {}",
fn_name, e
),
})?;
Ok(ResolvedAdminFunction {
admin_udf,
fn_name,
args,
arg_types,
ret_type,
})
}
/// Output schema of an `ADMIN` statement, mirroring what
/// [`CoreAdminFunctionService::execute`] produces. `None` if the function
/// or arguments are unresolvable; execution will then surface the error.
pub fn admin_output_schema(stmt: &Admin, query_ctx: &QueryContextRef) -> Option<Schema> {
let resolved =
resolve_admin_function(stmt, query_ctx, Arc::new(FunctionState::default())).ok()?;
Some(Schema::new(vec![ColumnSchema::new(
// Use statement as the result column name
stmt.to_string(),
ConcreteDataType::from_arrow_type(&resolved.ret_type),
false,
)]))
}
impl CoreAdminFunctionService {
fn new(query_engine: query::QueryEngineRef) -> Self {
Self { query_engine }
@@ -88,62 +186,23 @@ impl CoreAdminFunctionService {
query_ctx,
} = request;
let Admin::Func(func) = &stmt;
// the function name should be in lower case.
let func_name = func.name.to_string().to_lowercase();
let factory = get_admin_function(&func_name)
.or_else(|| FUNCTION_REGISTRY.get_function(&func_name))
.context(error::AdminFunctionNotFoundSnafu {
name: func_name.clone(),
})?;
let func_ctx = FunctionContext {
query_ctx: query_ctx.clone(),
state: self.query_engine.engine_state().function_state(),
};
let admin_udf = factory.provide(func_ctx);
let resolved = resolve_admin_function(
&stmt,
&query_ctx,
self.query_engine.engine_state().function_state(),
)?;
let ResolvedAdminFunction {
admin_udf,
fn_name,
args,
arg_types,
ret_type,
} = resolved;
let admin_async_fn = admin_udf
.as_async()
.context(error::AdminFunctionNotFoundSnafu { name: func_name })?;
let fn_name = admin_udf.name();
let signature = admin_udf.signature();
// Parse function arguments
let FunctionArguments::List(args) = &func.args else {
return error::BuildAdminFunctionArgsSnafu {
msg: format!("unsupported function args {} for {}", func.args, fn_name),
}
.fail();
};
let arg_values = args
.args
.iter()
.map(|arg| {
let FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) = arg else {
return error::BuildAdminFunctionArgsSnafu {
msg: format!("unsupported function arg {arg} for {}", fn_name),
}
.fail();
};
Ok(&value.value)
})
.collect::<Result<Vec<_>>>()?;
let args = args_to_vector(&signature.type_signature, &arg_values, &query_ctx)?;
let arg_types = args
.iter()
.map(|arg| arg.data_type().as_arrow_type())
.collect::<Vec<_>>();
let ret_type = admin_udf.return_type(&arg_types).map_err(|e| {
error::Error::BuildAdminFunctionArgs {
msg: format!(
"Failed to get return type of admin function {}: {}",
fn_name, e
),
}
})?;
.context(error::AdminFunctionNotFoundSnafu {
name: fn_name.clone(),
})?;
// Convert arguments to DataFusion ColumnarValue format
let columnar_args: Vec<datafusion_expr::ColumnarValue> = args
@@ -174,9 +233,7 @@ impl CoreAdminFunctionService {
let result_columnar = admin_async_fn
.invoke_async_with_args(func_args)
.await
.with_context(|_| ExecuteAdminFunctionSnafu {
msg: fn_name.to_string(),
})?;
.with_context(|_| ExecuteAdminFunctionSnafu { msg: fn_name })?;
// Convert result back to VectorRef
let result_columnar: common_query::prelude::ColumnarValue =
src/operator/src/statement/semantic_graph.rs
+171 -31
View File
@@ -41,8 +41,9 @@ use common_catalog::consts::{
ENTITY_DESCRIPTIVE_COLUMN, ENTITY_ID_ATTRS_COLUMN, ENTITY_ID_COLUMN, ENTITY_SCOPE_COLUMN,
ENTITY_TYPE_COLUMN, ERROR_COUNT_COLUMN, FRESH_UNTIL_COLUMN, GENERATION_ID_COLUMN,
OBSERVED_AT_COLUMN, PROVENANCE_COLUMN, REL_TYPE_COLUMN, REQUEST_COUNT_COLUMN,
SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SOURCE_TABLES_COLUMN, SRC_ID_COLUMN,
SRC_TYPE_COLUMN, VALID_FROM_COLUMN, VALID_UNTIL_COLUMN, WINDOW_END_COLUMN, WINDOW_START_COLUMN,
SEMANTIC_GRAPH_WINDOW_NANOS, SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SOURCE_TABLES_COLUMN,
SRC_ID_COLUMN, SRC_TYPE_COLUMN, VALID_FROM_COLUMN, VALID_UNTIL_COLUMN, WINDOW_END_COLUMN,
WINDOW_START_COLUMN,
};
use common_function::function::FunctionContext;
use common_function::function_registry::FUNCTION_REGISTRY;
@@ -60,7 +61,7 @@ use datafusion::dataframe::DataFrame;
use datafusion::functions::{core as core_fns, datetime as datetime_fns, string as string_fns};
use datafusion::functions_nested::expr_fn::make_array;
use datafusion_common::{Column, Result as DfResult, ScalarValue};
use datafusion_expr::{Expr, LogicalPlan, ScalarUDF, cast, ident, lit};
use datafusion_expr::{Case, Expr, LogicalPlan, ScalarUDF, cast, ident, lit};
pub use relationships::{
CallsSource, CoDeclaredSource, DeclaredSource, RelationshipSources, build_relationships_plan,
};
@@ -122,9 +123,10 @@ pub fn declared_relationships_schema_matches(table_info: &table::metadata::Table
})
}
/// Bin width for the temporal window of derived rows: 60s buckets, matching the
/// service-graph convention.
const BIN_NANOS: i64 = 60 * 1_000_000_000;
/// Bin width for the temporal window of derived rows, matching the
/// service-graph convention. Shared so ingestion-synthesized observations
/// land in the same buckets.
const BIN_NANOS: i64 = SEMANTIC_GRAPH_WINDOW_NANOS;
/// Default retention for the declared-edge table; expiry slides the topology window.
const DEFAULT_DECLARED_RELATIONSHIPS_TTL: &str = "90d";
@@ -284,8 +286,12 @@ pub struct EntityDeclaration {
/// The table's time index column, used for the temporal window filter.
pub time_index: String,
pub entity_type: String,
/// Identifying columns (>= 1). One column → id verbatim; several → composite.
/// Identifying columns (>= 1), ordered broad to narrow.
pub id_columns: Vec<String>,
/// Optional column qualifying the first identity component, so a source
/// carrying the parts separately (a trace table's `service.namespace`)
/// yields the same id as one carrying them pre-composed (`job`).
pub id_qualifier: Option<String>,
/// Descriptive columns snapshotted into the `descriptive` JSON (may be empty).
pub descriptive_columns: Vec<String>,
/// Scope columns (namespace/environment). One column → scope verbatim;
@@ -417,28 +423,91 @@ fn cast_string_or_empty(column: &str) -> Expr {
core_fns::coalesce().call(vec![cast(ident(column), DataType::Utf8), lit("")])
}
/// The canonical entity-id expression for `id_columns`: the value verbatim for
/// a single column, the sorted `k=v,k=v` rendering for a composite. `col`
/// constructs the column reference (unqualified for registry branches,
/// join-side-qualified for the calls derivation).
/// A row identifies an entity only when every identity component is present
/// and non-empty: kube-state-metrics descriptors emit empty-string labels (an
/// unscheduled pod's `node`, an owner-less pod's `owner_*`), and an empty
/// string is never a meaningful entity id.
/// string is never a meaningful entity id. The qualifier is optional by
/// construction and so is not guarded.
pub(crate) fn identifies(column: &str) -> Expr {
ident(column)
.is_not_null()
.and(cast(ident(column), DataType::Utf8).not_eq(lit("")))
}
fn entity_id_expr(id_columns: &[String], col: &dyn Fn(&str) -> Expr) -> Expr {
if let [id] = id_columns {
cast(col(id), DataType::Utf8)
} else {
let mut cols = id_columns.to_vec();
cols.sort();
sorted_kv_expr_with(&cols, false, col)
const ID_SEPARATOR: &str = ",";
const ID_ESCAPE: &str = "\\";
/// Left unescaped: `<namespace>/<name>` is how Prometheus renders `job`.
const ID_QUALIFIER_SEPARATOR: &str = "/";
/// Escapes so a composite id decodes back to its components. The escape
/// character goes first, or it would double the escapes the separator pass
/// introduces.
fn escaped_id_value(value: Expr) -> Expr {
let escape_escape = string_fns::replace().call(vec![
value,
lit(ID_ESCAPE),
lit(format!("{ID_ESCAPE}{ID_ESCAPE}")),
]);
string_fns::replace().call(vec![
escape_escape,
lit(ID_SEPARATOR),
lit(format!("{ID_ESCAPE}{ID_SEPARATOR}")),
])
}
/// The identity values in declared order (broad to narrow), escaped and
/// joined.
///
/// The identifying *column names* are deliberately absent: the same identity
/// reaches us under different names per source — a trace table's
/// `service_name` against a metric table's `job` — and encoding them would
/// split one entity into one per signal. `entity_id_attrs` keeps the
/// structured form.
///
/// `col` constructs the column reference (unqualified for registry branches,
/// join-side-qualified for the calls derivation).
fn entity_id_expr(
id_columns: &[String],
qualifier: Option<&str>,
col: &dyn Fn(&str) -> Expr,
) -> Expr {
let mut parts = Vec::with_capacity(id_columns.len() * 2);
for (index, column) in id_columns.iter().enumerate() {
if index > 0 {
parts.push(lit(ID_SEPARATOR));
}
let value = escaped_id_value(cast(col(column), DataType::Utf8));
parts.push(match qualifier {
Some(qualifier) if index == 0 => qualified_id_expr(qualifier, value, col),
_ => value,
});
}
if let [single] = parts.as_slice() {
single.clone()
} else {
concat_expr(parts)
}
}
/// `<qualifier>/<value>`, or bare `value` when the qualifier is empty on this
/// row — the spec's rule composing `job` from `service.namespace` and
/// `service.name`.
fn qualified_id_expr(qualifier: &str, value: Expr, col: &dyn Fn(&str) -> Expr) -> Expr {
let qualifier = escaped_id_value(
core_fns::coalesce().call(vec![cast(col(qualifier), DataType::Utf8), lit("")]),
);
Expr::Case(Case::new(
None,
vec![(
Box::new(qualifier.clone().eq(lit(""))),
Box::new(value.clone()),
)],
Some(Box::new(concat_expr(vec![
qualifier,
lit(ID_QUALIFIER_SEPARATOR),
value,
]))),
))
}
/// The `parse_json` UDF, shared by all derivation plans. Resolved from the
@@ -583,14 +652,22 @@ fn registry_source(
let mut rows = Vec::with_capacity(1 + rest.len());
for decl in std::iter::once(first).chain(rest) {
// CAST even a single-column id: id columns need not be strings, and
// the computed table declares entity_id STRING. Composite ids
// additionally carry a JSON object of the id columns in
// entity_id_attrs.
let entity_id = entity_id_expr(&decl.id_columns, &|c| ident(c));
let entity_id_attrs = if decl.id_columns.len() == 1 {
// the computed table declares entity_id STRING. An id assembled from
// more than one column additionally carries its parts as a JSON
// object in entity_id_attrs, the structured form of the id string.
let entity_id = entity_id_expr(&decl.id_columns, decl.id_qualifier.as_deref(), &|c| {
ident(c)
});
let id_parts = decl
.id_qualifier
.iter()
.chain(&decl.id_columns)
.cloned()
.collect::<Vec<_>>();
let entity_id_attrs = if id_parts.len() == 1 {
null_json()
} else {
json_object_expr(&decl.id_columns)
json_object_expr(&id_parts)
};
let scope = match decl.scope_columns.as_slice() {
@@ -811,11 +888,74 @@ mod tests {
time_index: "ts".to_string(),
entity_type: entity_type.to_string(),
id_columns: id_columns.iter().map(|s| s.to_string()).collect(),
id_qualifier: None,
descriptive_columns: vec![],
scope_columns: vec![],
}
}
/// The id string must decode back to its components, and a qualifier must
/// reproduce the identity a source that pre-composed it already emits.
#[tokio::test]
async fn registry_id_escaping_and_qualifier() {
let schema = Arc::new(Schema::new(vec![
Field::new(
"ts",
DataType::Timestamp(TimeUnit::Millisecond, None),
false,
),
Field::new("service_name", DataType::Utf8, false),
Field::new("instance", DataType::Utf8, false),
Field::new("namespace", DataType::Utf8, true),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(TimestampMillisecondArray::from(vec![1_000, 2_000, 3_000])) as ArrayRef,
Arc::new(StringArray::from(vec!["cart", "a,b", "we\\ird"])),
Arc::new(StringArray::from(vec!["i-1", "c", "i-3"])),
Arc::new(StringArray::from(vec![Some("shop"), None, Some("")])),
],
)
.unwrap();
let ctx = SessionContext::new();
ctx.register_table(
"svc",
Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()),
)
.unwrap();
let mut declaration = decl("service.instance", &["service_name", "instance"]);
declaration.id_qualifier = Some("namespace".to_string());
declaration.time_index = "ts".to_string();
let plan = build_registry_plan(
vec![RegistrySource {
declarations: vec![declaration],
scan: ctx.table("svc").await.unwrap(),
}],
&test_window(),
)
.unwrap()
.unwrap();
let mut ids: Vec<String> = collect(&ctx, plan)
.await
.iter()
.flat_map(|batch| strings(batch, 5))
.collect();
ids.sort();
assert_eq!(
ids,
vec![
// an absent qualifier leaves the identity bare, and a value
// holding the separator stays distinguishable from two values
"a\\,b,c".to_string(),
"shop/cart,i-1".to_string(),
"we\\\\ird,i-3".to_string(),
]
);
}
#[tokio::test]
async fn registry_single_column_identity() {
let ctx = metric_table_ctx();
@@ -900,24 +1040,24 @@ mod tests {
.collect();
rows.sort();
// Composite id -> sorted `k=v,k=v` plus a JSON object of the id columns;
// descriptive JSON keeps `\`, `"` and control characters intact in
// runtime values, NULL -> "".
// Composite id -> values in declared order plus a JSON object of the
// id columns; descriptive JSON keeps `\`, `"` and control characters
// intact in runtime values, NULL -> "".
assert_eq!(
rows,
vec![
(
"pid=42,service_name=cart".to_string(),
"cart,42".to_string(),
Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()),
Some(r#"{"host":""}"#.to_string()),
),
(
"pid=42,service_name=cart".to_string(),
"cart,42".to_string(),
Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()),
Some(r#"{"host":"h2"}"#.to_string()),
),
(
"pid=42,service_name=cart".to_string(),
"cart,42".to_string(),
Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()),
Some(r#"{"host":"we\"ird\\\nhost"}"#.to_string()),
),
src/operator/src/statement/semantic_graph/conventions.rs
+62 -12
View File
@@ -19,7 +19,7 @@
//! with the binary, not an operator-editable configuration surface; explicit
//! `greptime.semantic.entity.*` declarations always override it.
use std::collections::{BTreeMap, HashSet};
use std::collections::{BTreeMap, HashMap, HashSet};
use std::sync::LazyLock;
use serde::Deserialize;
@@ -43,15 +43,20 @@ pub struct VirtualDstCandidate {
pub connection_type: String,
}
/// One implicit entity declaration: of a whitelisted Prometheus info metric,
/// or of a trace-v1 table's flattened resource attributes.
/// One implicit entity declaration: of a whitelisted Prometheus or OTel info
/// metric, or of a trace-v1 table's flattened resource attributes.
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ImplicitEntity {
pub entity: String,
/// Identifying label columns; every one must exist on the table for the
/// declaration to apply.
/// Identifying label columns, ordered broad to narrow; every one must
/// exist on the table for the declaration to apply.
pub id: Vec<String>,
/// Column qualifying `id[0]` as `<qualifier>/<id[0]>`, so a source
/// carrying the parts separately matches one carrying them pre-composed.
/// Skipped when the column is absent or empty.
#[serde(default)]
pub qualified_by: Option<String>,
/// Descriptive label columns, filtered to those present (kube-state-metrics
/// label sets vary across versions).
#[serde(default)]
@@ -70,7 +75,12 @@ pub struct Conventions {
pub trace_co_declared_edges: Vec<EdgeRule>,
pub virtual_dst_candidates: Vec<VirtualDstCandidate>,
pub otlp_trace_entities: Vec<ImplicitEntity>,
/// Table name -> the entities that table declares, for Prometheus-sourced
/// descriptor metrics (`source = prometheus`).
pub prometheus_info_metrics: BTreeMap<String, Vec<ImplicitEntity>>,
/// Table name -> the entities that table declares, for OTLP-sourced
/// descriptor tables (`source = opentelemetry`).
pub otel_info_metrics: BTreeMap<String, Vec<ImplicitEntity>>,
}
/// The built-in entity-type vocabulary. User-declared types are open-ended;
@@ -78,6 +88,7 @@ pub struct Conventions {
pub const ENTITY_TYPE_SERVICE: &str = "service";
pub const ENTITY_TYPE_SERVICE_INSTANCE: &str = "service.instance";
pub const ENTITY_TYPE_HOST: &str = "host";
pub const ENTITY_TYPE_CONTAINER: &str = "container";
pub const ENTITY_TYPE_PROCESS: &str = "process";
pub const ENTITY_TYPE_K8S_POD: &str = "k8s.pod";
pub const ENTITY_TYPE_K8S_NODE: &str = "k8s.node";
@@ -88,10 +99,11 @@ pub const ENTITY_TYPE_GEN_AI_AGENT: &str = "gen_ai.agent";
pub const ENTITY_TYPE_GEN_AI_MODEL: &str = "gen_ai.model";
pub const ENTITY_TYPE_GEN_AI_TOOL: &str = "gen_ai.tool";
const ENTITY_TYPES: [&str; 12] = [
const ENTITY_TYPES: [&str; 13] = [
ENTITY_TYPE_SERVICE,
ENTITY_TYPE_SERVICE_INSTANCE,
ENTITY_TYPE_HOST,
ENTITY_TYPE_CONTAINER,
ENTITY_TYPE_PROCESS,
ENTITY_TYPE_K8S_POD,
ENTITY_TYPE_K8S_NODE,
@@ -200,11 +212,16 @@ fn validate(conventions: &Conventions) -> Result<(), String> {
let per_table = conventions
.prometheus_info_metrics
.iter()
.chain(&conventions.otel_info_metrics)
.map(|(table, entities)| (table.as_str(), entities))
.chain(std::iter::once((
"otlp traces",
&conventions.otlp_trace_entities,
)));
// An entity type declared with a different number of id columns by two
// sources yields ids that can never match, silently splitting one entity
// in two.
let mut arity = HashMap::new();
for (table, entities) in per_table {
let mut seen_types = HashSet::new();
for implicit in entities {
@@ -226,6 +243,12 @@ fn validate(conventions: &Conventions) -> Result<(), String> {
implicit.entity
));
}
if implicit.qualified_by.as_ref().is_some_and(String::is_empty) {
return Err(format!(
"entity `{}` of info metric `{table}` has an empty qualified_by",
implicit.entity
));
}
if implicit.descriptive_rest && !implicit.descriptive.is_empty() {
return Err(format!(
"entity `{}` of info metric `{table}` sets both descriptive and \
@@ -233,6 +256,17 @@ fn validate(conventions: &Conventions) -> Result<(), String> {
implicit.entity
));
}
match arity.insert(implicit.entity.as_str(), implicit.id.len()) {
Some(previous) if previous != implicit.id.len() => {
return Err(format!(
"entity `{}` is declared with {previous} id columns elsewhere but \
{} for `{table}`",
implicit.entity,
implicit.id.len()
));
}
_ => {}
}
}
}
Ok(())
@@ -249,25 +283,28 @@ mod tests {
/// Each case must fail on exactly the rule it names, so the shared
/// boilerplate is valid and the mutated part is inside the vocabulary.
fn broken(edges: &str, info_metrics: &str) -> String {
fn broken(edges: &str, info_metrics: &str, otel_metrics: &str) -> String {
format!(
"co_declared_edges: [{edges}]\ntrace_co_declared_edges: []\n\
virtual_dst_candidates: []\notlp_trace_entities: []\n\
prometheus_info_metrics: {{{info_metrics}}}"
prometheus_info_metrics: {{{info_metrics}}}\n\
otel_info_metrics: {{{otel_metrics}}}"
)
}
#[test]
fn validation_rejects_broken_conventions() {
let err = |edges, info| parse(&broken(edges, info)).unwrap_err();
let err = |edges, info, otel| parse(&broken(edges, info, otel)).unwrap_err();
assert!(err("{src: host, dst: service, rel: pets}", "").contains("unknown rel_type"));
assert!(err("{src: host, dst: service, rel: pets}", "", "").contains("unknown rel_type"));
assert!(
err("{src: k8s.pods, dst: k8s.node, rel: runs_on}", "").contains("unknown entity type")
err("{src: k8s.pods, dst: k8s.node, rel: runs_on}", "", "")
.contains("unknown entity type")
);
assert!(
err(
"{src: host, dst: service, rel: uses}, {src: host, dst: service, rel: uses}",
"",
""
)
.contains("duplicate edge rule")
@@ -275,14 +312,27 @@ mod tests {
assert!(
err(
"",
"t: [{entity: host, id: [x], descriptive: [y], descriptive_rest: true}]"
"t: [{entity: host, id: [x], descriptive: [y], descriptive_rest: true}]",
""
)
.contains("descriptive_rest")
);
// the otel map runs through the same per-table validation
assert!(err("", "", "t: [{entity: hosts, id: [x]}]").contains("unknown entity type"));
// ids of a different arity for one type can never match each other
assert!(
err(
"",
"a: [{entity: host, id: [x]}]",
"b: [{entity: host, id: [x, y]}]"
)
.contains("id columns")
);
// Unknown YAML keys are rejected, catching typos in the embedded file.
assert!(
parse(&broken(
"{src: host, dst: service, rel: uses, direction: down}",
"",
""
))
.is_err()
src/operator/src/statement/semantic_graph/conventions.yaml
+49 -1
View File
@@ -2,6 +2,10 @@
# versioned with the binary (include_str!), not an operator-editable surface;
# non-standard deployments override per table with explicit
# greptime.semantic.entity.* declarations, which always win.
#
# Every `id` list is ordered broad to narrow, and that order is part of the
# identity: the id renders as the values joined in this order, so all
# declarations of one entity type must agree on it or they name two entities.
co_declared_edges:
- { src: service.instance, dst: host, rel: runs_on }
@@ -11,6 +15,8 @@ co_declared_edges:
- { src: service.instance, dst: service, rel: part_of }
- { src: k8s.pod, dst: k8s.workload, rel: part_of }
- { src: service.instance, dst: k8s.pod, rel: runs_on }
- { src: service.instance, dst: container, rel: runs_on }
- { src: container, dst: host, rel: runs_on }
# Applied only to trace sources; the endpoints are Greptime entity types
# derived from GenAI semantic-convention attributes.
@@ -30,12 +36,25 @@ virtual_dst_candidates:
- { column: span_attributes.server.address, connection_type: virtual_node }
# Implicit declarations for greptime_trace_v1 tables; the pod UID bridges
# trace-side pods onto the kube-state-metrics entities below.
# trace-side pods onto the kube-state-metrics entities below. host.name and
# container.name vary by SDK and resource detector, so identity is the id.
# qualified_by folds the namespace into the service identity the same way the
# metric sources' `job` already carries it, so both signals name one entity.
otlp_trace_entities:
- entity: service
id: [service_name]
qualified_by: resource_attributes.service.namespace
- entity: service.instance
id: [service_name, resource_attributes.service.instance.id]
qualified_by: resource_attributes.service.namespace
- entity: host
id: [resource_attributes.host.id]
descriptive:
- resource_attributes.host.name
- entity: container
id: [resource_attributes.container.id]
descriptive:
- resource_attributes.container.name
- entity: k8s.pod
id: [resource_attributes.k8s.pod.uid]
descriptive:
@@ -113,3 +132,32 @@ prometheus_info_metrics:
- entity: service.instance
id: [job, instance]
descriptive_rest: true
# Implicit declarations for the OTLP resource descriptor synthesized at
# metrics ingestion, gated on source=opentelemetry and metric.type=info.
# job/instance are the derived Prometheus-compatible identity, so services
# line up with target_info. Unlike target_info the attribute set is known, so
# descriptive columns are listed instead of swept up by descriptive_rest,
# which would misfile other entities' ids under service.instance.
otel_info_metrics:
greptime_otel_resource_info:
- entity: service
id: [job]
descriptive:
- service.name
- service.namespace
- entity: service.instance
id: [job, instance]
- entity: host
id: [host.id]
descriptive:
- host.name
- entity: container
id: [container.id]
descriptive:
- container.name
- entity: k8s.pod
id: [k8s.pod.uid]
descriptive:
- k8s.pod.name
- k8s.namespace.name
src/operator/src/statement/semantic_graph/relationships.rs
+55 -24
View File
@@ -200,9 +200,9 @@ fn co_declared_branch(
bin.clone() + bin_interval(),
bin.clone() + bin_interval(),
lit(src.entity_type.as_str()),
entity_id_expr(&src.id_columns, &|c| ident(c)),
entity_id_expr(&src.id_columns, src.id_qualifier.as_deref(), &|c| ident(c)),
lit(dst.entity_type.as_str()),
entity_id_expr(&dst.id_columns, &|c| ident(c)),
entity_id_expr(&dst.id_columns, dst.id_qualifier.as_deref(), &|c| ident(c)),
lit(*rel_type),
lit(*provenance),
lit(1.0_f64),
@@ -501,6 +501,13 @@ fn span_predicate(service: &EntityDeclaration, window: &GraphQueryWindow, strict
}
/// One trace table's client spans, normalized to
/// A trace table's `duration_nano` is UInt64 on tables created before the
/// signed-integer ingest change and Int64 after it. The per-table selects are
/// unioned, and those two have no common integer type.
fn duration_nano_expr() -> Expr {
cast(ident(DURATION_NANO_COLUMN), DataType::Int64).alias(DURATION_NANO_COLUMN)
}
/// `(timestamp, trace_id, span_id, src_id, status_code, duration_nano,
/// virtual_dst, virtual_conn)`. `src_id` is built from the table's `service`
/// declaration, so edges land on exactly the entity ids the registry emits (a
@@ -553,9 +560,12 @@ fn client_spans(
ident(SPAN_ID_COLUMN),
// The cast inside entity_id_expr also normalizes tag columns, which
// come out of the storage engine dictionary-encoded.
entity_id_expr(&service.id_columns, &|c| ident(c)).alias(SRC_ID_COLUMN),
entity_id_expr(&service.id_columns, service.id_qualifier.as_deref(), &|c| {
ident(c)
})
.alias(SRC_ID_COLUMN),
ident(SPAN_STATUS_CODE_COLUMN),
ident(DURATION_NANO_COLUMN),
duration_nano_expr(),
virtual_dst.alias("virtual_dst"),
virtual_conn.alias("virtual_conn"),
])
@@ -578,9 +588,12 @@ fn server_spans(
ident(TRACE_TIMESTAMP_COLUMN),
ident(TRACE_ID_COLUMN),
ident(PARENT_SPAN_ID_COLUMN),
entity_id_expr(&service.id_columns, &|c| ident(c)).alias(DST_ID_COLUMN),
entity_id_expr(&service.id_columns, service.id_qualifier.as_deref(), &|c| {
ident(c)
})
.alias(DST_ID_COLUMN),
ident(SPAN_STATUS_CODE_COLUMN),
ident(DURATION_NANO_COLUMN),
duration_nano_expr(),
])
}
@@ -609,7 +622,10 @@ fn agent_calls_branch(
ident(TRACE_TIMESTAMP_COLUMN),
ident(TRACE_ID_COLUMN),
ident(SPAN_ID_COLUMN),
entity_id_expr(&agent.id_columns, &|c| ident(c)).alias(SRC_ID_COLUMN),
entity_id_expr(&agent.id_columns, agent.id_qualifier.as_deref(), &|c| {
ident(c)
})
.alias(SRC_ID_COLUMN),
])?;
let child = trace
.scan
@@ -619,9 +635,12 @@ fn agent_calls_branch(
ident(TRACE_TIMESTAMP_COLUMN),
ident(TRACE_ID_COLUMN),
ident(PARENT_SPAN_ID_COLUMN),
entity_id_expr(&agent.id_columns, &|c| ident(c)).alias(DST_ID_COLUMN),
entity_id_expr(&agent.id_columns, agent.id_qualifier.as_deref(), &|c| {
ident(c)
})
.alias(DST_ID_COLUMN),
ident(SPAN_STATUS_CODE_COLUMN),
ident(DURATION_NANO_COLUMN),
duration_nano_expr(),
])?;
parents = union_all(parents, parent)?;
children = union_all(children, child)?;
@@ -732,6 +751,16 @@ mod tests {
name: &str,
extra: &[(&str, &[Option<&str>])],
spans: &[Span<'_>],
) {
register_typed_trace_table(ctx, name, extra, spans, DataType::UInt64)
}
fn register_typed_trace_table(
ctx: &SessionContext,
name: &str,
extra: &[(&str, &[Option<&str>])],
spans: &[Span<'_>],
duration_type: DataType,
) {
let mut fields = vec![
Field::new(
@@ -745,7 +774,7 @@ mod tests {
Field::new(SPAN_KIND_COLUMN, DataType::Utf8, false),
Field::new(SPAN_STATUS_CODE_COLUMN, DataType::Utf8, false),
Field::new(SERVICE_NAME_COLUMN, DataType::Utf8, false),
Field::new(DURATION_NANO_COLUMN, DataType::UInt64, false),
Field::new(DURATION_NANO_COLUMN, duration_type.clone(), false),
];
for (column, _) in extra {
fields.push(Field::new(*column, DataType::Utf8, true));
@@ -774,9 +803,14 @@ mod tests {
Arc::new(StringArray::from(
spans.iter().map(|s| s.6).collect::<Vec<_>>(),
)),
Arc::new(UInt64Array::from(
spans.iter().map(|s| s.7).collect::<Vec<_>>(),
)),
match duration_type {
DataType::Int64 => Arc::new(Int64Array::from(
spans.iter().map(|s| s.7 as i64).collect::<Vec<_>>(),
)) as ArrayRef,
_ => Arc::new(UInt64Array::from(
spans.iter().map(|s| s.7).collect::<Vec<_>>(),
)),
},
];
for (_, values) in extra {
columns.push(Arc::new(StringArray::from(values.to_vec())));
@@ -843,6 +877,7 @@ mod tests {
time_index: TRACE_TIMESTAMP_COLUMN.to_string(),
entity_type: "service".to_string(),
id_columns: id_columns.iter().map(|s| s.to_string()).collect(),
id_qualifier: None,
descriptive_columns: vec![],
scope_columns: vec![],
}
@@ -972,7 +1007,7 @@ mod tests {
),
],
);
register_trace_table(
register_typed_trace_table(
&ctx,
"trace_b",
&[],
@@ -989,6 +1024,7 @@ mod tests {
1_500_000_000,
),
],
DataType::Int64,
);
let a = ctx.table("trace_a").await.unwrap();
let b = ctx.table("trace_b").await.unwrap();
@@ -1294,6 +1330,7 @@ mod tests {
time_index: "ts".to_string(),
entity_type: entity_type.to_string(),
id_columns: id_columns.iter().map(|s| s.to_string()).collect(),
id_qualifier: None,
descriptive_columns: vec![],
scope_columns: vec![],
}
@@ -1544,16 +1581,10 @@ mod tests {
let total: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(total, 1);
let batch = &batches[0];
// Composite service identity renders the same sorted `k=v` form the
// registry emits, so edges land on registry entity ids.
assert_eq!(
strings(batch, 5),
vec!["service_name=frontend,service_namespace=ns1"]
);
assert_eq!(
strings(batch, 7),
vec!["service_name=cart,service_namespace=ns1"]
);
// Composite service identity renders exactly as the registry emits it,
// so edges land on registry entity ids.
assert_eq!(strings(batch, 5), vec!["frontend,ns1"]);
assert_eq!(strings(batch, 7), vec!["cart,ns1"]);
}
fn build_relationships_plan_for_test(
src/query/src/analyze.rs
+12 -5
View File
@@ -46,6 +46,17 @@ const STAGE: &str = "stage";
const NODE: &str = "node";
const PLAN: &str = "plan";
/// Fixed output schema of [`DistAnalyzeExec`], for `Describe` handlers:
/// execution rewrites the plan in `optimize_physical_plan`, so this schema
/// differs from the logical `Analyze` plan's.
pub fn dist_analyze_output_schema() -> SchemaRef {
SchemaRef::new(Schema::new(vec![
Field::new(STAGE, DataType::UInt32, true),
Field::new(NODE, DataType::UInt32, true),
Field::new(PLAN, DataType::Utf8, true),
]))
}
#[derive(Debug)]
pub struct DistAnalyzeExec {
input: Arc<dyn ExecutionPlan>,
@@ -58,11 +69,7 @@ pub struct DistAnalyzeExec {
impl DistAnalyzeExec {
/// Create a new DistAnalyzeExec
pub fn new(input: Arc<dyn ExecutionPlan>, verbose: bool, format: AnalyzeFormat) -> Self {
let schema = SchemaRef::new(Schema::new(vec![
Field::new(STAGE, DataType::UInt32, true),
Field::new(NODE, DataType::UInt32, true),
Field::new(PLAN, DataType::Utf8, true),
]));
let schema = dist_analyze_output_schema();
let properties = Arc::new(Self::compute_properties(&input, schema.clone()));
Self {
input,
src/query/src/datafusion.rs
+101 -25
View File
@@ -30,7 +30,7 @@ use common_function::function::FunctionContext;
use common_function::function_factory::ScalarFunctionFactory;
use common_query::{Output, OutputData, OutputMeta};
use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters};
use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream};
use common_recordbatch::{EmptyRecordBatchStream, RecordBatch, SendableRecordBatchStream};
use common_telemetry::tracing;
use datafusion::catalog::TableFunction;
use datafusion::dataframe::DataFrame;
@@ -44,6 +44,7 @@ use datafusion_expr::{
use datatypes::prelude::VectorRef;
use datatypes::schema::Schema;
use futures_util::StreamExt;
use futures_util::future::try_join;
use session::context::QueryContextRef;
use snafu::{OptionExt, ResultExt, ensure};
use sqlparser::ast::AnalyzeFormat;
@@ -80,6 +81,29 @@ pub const QUERY_PARALLELISM_HINT: &str = "query_parallelism";
/// Whether to fallback to the original plan when failed to push down.
pub const QUERY_FALLBACK_HINT: &str = "query_fallback";
// An unbounded queue keeps draining source RPCs while mutation RPCs on a shared
// HTTP/2 connection are pending, trading bounded memory for request liveness.
async fn forward_record_batches(
mut stream: SendableRecordBatchStream,
batch_tx: tokio::sync::mpsc::UnboundedSender<Result<RecordBatch>>,
) -> Result<()> {
while let Some(batch) = stream.next().await {
match batch.context(CreateRecordBatchSnafu) {
Ok(batch) => {
if batch_tx.send(Ok(batch)).is_err() {
break;
}
tokio::task::yield_now().await;
}
Err(error) => {
let _ = batch_tx.send(Err(error));
break;
}
}
}
Ok(())
}
fn query_load_region_id(plan: &Arc<dyn ExecutionPlan>) -> Option<u64> {
let mut region_id = None;
let mut stack = vec![plan.clone()];
@@ -205,7 +229,7 @@ impl DatafusionQueryEngine {
let Output { data, meta } = self
.exec_query_plan((*dml.input).clone(), query_ctx.clone())
.await?;
let mut stream = match data {
let stream = match data {
OutputData::RecordBatches(batches) => batches.as_stream(),
OutputData::Stream(stream) => stream,
_ => unreachable!(),
@@ -214,30 +238,35 @@ impl DatafusionQueryEngine {
let mut affected_rows = 0;
let mut insert_cost = 0;
while let Some(batch) = stream.next().await {
let batch = batch.context(CreateRecordBatchSnafu)?;
let column_vectors = batch
.column_vectors(&table_name.to_string(), table.schema())
.map_err(BoxedError::new)
.context(QueryExecutionSnafu)?;
match dml.op {
WriteOp::Insert(_) => {
// We ignore the insert op.
let output = self
.insert(&table_name, column_vectors, query_ctx.clone())
.await?;
let (rows, cost) = output.extract_rows_and_cost();
affected_rows += rows;
insert_cost += cost;
}
WriteOp::Delete => {
affected_rows += self
.delete(&table_name, &table, column_vectors, query_ctx.clone())
.await?;
}
_ => unreachable!("guarded by the 'ensure!' at the beginning"),
match dml.op {
WriteOp::Insert(_) => {
let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel();
let producer = forward_record_batches(stream, batch_tx);
let consumer = self.consume_insert_record_batches(
batch_rx,
&table_name,
table.schema(),
query_ctx.clone(),
);
let ((), (rows, cost)) = try_join(producer, consumer).await?;
affected_rows += rows;
insert_cost += cost;
}
WriteOp::Delete => {
// Keep DELETE on the same producer/consumer schedule as INSERT so the source
// stream can continue draining while mutation RPCs are pending.
let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel();
let producer = forward_record_batches(stream, batch_tx);
let consumer = self.consume_delete_record_batches(
batch_rx,
&table_name,
&table,
query_ctx.clone(),
);
let (rows, ()) = try_join(consumer, producer).await?;
affected_rows += rows;
}
_ => unreachable!("guarded by the 'ensure!' at the beginning"),
}
Ok(Output::new(
OutputData::AffectedRows(affected_rows),
@@ -245,6 +274,53 @@ impl DatafusionQueryEngine {
))
}
async fn consume_insert_record_batches(
&self,
mut batch_rx: tokio::sync::mpsc::UnboundedReceiver<Result<RecordBatch>>,
table_name: &ResolvedTableReference,
table_schema: Arc<Schema>,
query_ctx: QueryContextRef,
) -> Result<(usize, usize)> {
let mut affected_rows = 0;
let mut insert_cost = 0;
while let Some(batch) = batch_rx.recv().await {
let batch = batch?;
let column_vectors = batch
.column_vectors(&table_name.to_string(), table_schema.clone())
.map_err(BoxedError::new)
.context(QueryExecutionSnafu)?;
// We ignore the insert op.
let output = self
.insert(table_name, column_vectors, query_ctx.clone())
.await?;
let (rows, cost) = output.extract_rows_and_cost();
affected_rows += rows;
insert_cost += cost;
}
Ok((affected_rows, insert_cost))
}
async fn consume_delete_record_batches(
&self,
mut batch_rx: tokio::sync::mpsc::UnboundedReceiver<Result<RecordBatch>>,
table_name: &ResolvedTableReference,
table: &TableRef,
query_ctx: QueryContextRef,
) -> Result<usize> {
let mut affected_rows = 0;
while let Some(batch) = batch_rx.recv().await {
let batch = batch?;
let column_vectors = batch
.column_vectors(&table_name.to_string(), table.schema())
.map_err(BoxedError::new)
.context(QueryExecutionSnafu)?;
affected_rows += self
.delete(table_name, table, column_vectors, query_ctx.clone())
.await?;
}
Ok(affected_rows)
}
#[tracing::instrument(skip_all)]
async fn delete(
&self,
src/query/src/lib.rs
+1 -1
View File
@@ -45,7 +45,7 @@ pub(crate) mod test_util;
#[cfg(test)]
mod tests;
pub use crate::analyze::analyze_plan_metrics_to_json_value;
pub use crate::analyze::{analyze_plan_metrics_to_json_value, dist_analyze_output_schema};
pub use crate::datafusion::DfContextProviderAdapter;
pub use crate::query_engine::{
QueryEngine, QueryEngineContext, QueryEngineFactory, QueryEngineRef,
src/query/src/optimizer.rs
+1
View File
@@ -17,6 +17,7 @@ pub mod constant_term;
pub mod count_nest_aggr;
pub mod count_wildcard;
pub mod global_limit;
pub(crate) mod insert_assignment;
pub(crate) mod json_type_concretize;
pub mod parallelize_scan;
pub mod pass_distribution;
src/query/src/optimizer/insert_assignment.rs
+307
View File
@@ -0,0 +1,307 @@
// 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::sync::Arc;
use common_time::Timezone;
use datafusion::config::ConfigOptions;
use datafusion_common::{DFSchemaRef, Result, ScalarValue};
use datafusion_expr::expr::{Alias, Cast};
use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Values};
use datafusion_optimizer::analyzer::AnalyzerRule;
use datafusion_optimizer::analyzer::type_coercion::TypeCoercion;
use datatypes::arrow::datatypes::{DataType, TimeUnit};
use session::context::QueryContextRef;
use crate::optimizer::type_conversion::cast_string_to_timestamp;
/// Interprets strings assigned to timestamp columns at an `INSERT` boundary
/// using the session timezone. `plan` is the assignment projection under a
/// `WriteOp::Insert`.
///
/// DataFusion plans `INSERT` as a projection casting each source column to its
/// target column type, and that cast reads a naive string as UTC. Arrow does
/// apply a timezone when the cast target carries one, so the assignment is
/// routed through `Timestamp(unit, Some(tz))` and back. Stripping the timezone
/// afterwards is value-preserving — arrow only shifts values in the opposite
/// direction.
///
/// The source query is left untouched. Reinterpreting a value where it is
/// *produced* would change what the source query means: pushing the conversion
/// below a `UNION`'s `DISTINCT`, for instance, moves the dedup key from the raw
/// strings to parsed instants and silently drops rows.
///
/// # Why `TypeCoercion` runs here
///
/// The rewrite reads source types, and those are only settled once a `UNION`'s
/// branch types have been reconciled: before coercion a union carries its loose
/// schema (the first branch's types), so a mixed
/// `SELECT 'string' UNION ALL SELECT CAST(.. AS TIMESTAMP)` still looks like a
/// string. Retargeting that cast would leave `Timestamp(None) ->
/// Timestamp(Some(tz))` behind once coercion retypes the union — the one
/// direction in which arrow shifts the value instead of relabelling it.
///
/// Coercing here rather than deferring to the analyzer is forced by where an
/// INSERT is still identifiable: `exec_dml_statement` strips the `Dml` node and
/// executes its input, so by the time the analyzer runs, an assignment
/// projection is indistinguishable from any other projection.
///
/// Explicit casts stay out of this: the SQL layer turns a user's
/// `CAST(x AS TIMESTAMP)` into an `arrow_cast` call, which only becomes an
/// `Expr::Cast` in the optimizer's `SimplifyExpressions`. Assignment casts are
/// therefore the only `Expr::Cast` reaching a timestamp column here.
///
/// # Reach
///
/// The emitted cast only carries its timezone where the expression is evaluated
/// on this node. Substrait drops the timezone name when a plan is pushed down —
/// it encodes any zoned timestamp as `PrecisionTimestampTz` and decodes it back
/// as UTC — so a source reading from a table falls back to UTC, the behaviour it
/// had before this rule existed. Sources that never leave this node (literals,
/// `VALUES`, and `UNION`s of them) keep the session timezone, and those are what
/// an INSERT's timestamp assignment is in practice.
pub(crate) fn rewrite_insert_assignments(
plan: LogicalPlan,
query_ctx: &QueryContextRef,
config: &ConfigOptions,
) -> Result<LogicalPlan> {
let Some(timezone) = session_timezone(query_ctx) else {
return Ok(plan);
};
let plan = TypeCoercion::new().analyze(plan, config)?;
rewrite_assignment(plan, &timezone)
}
/// Session timezone, in both forms the rewrite needs.
struct SessionTimezone {
/// Parses literals, matching the plain `INSERT ... VALUES` path.
parsed: Timezone,
/// Names the intermediate arrow cast target.
name: Arc<str>,
}
fn session_timezone(query_ctx: &QueryContextRef) -> Option<SessionTimezone> {
let parsed = query_ctx.timezone();
// A UTC session already gets UTC semantics from the plain assignment cast.
if parsed.is_utc() {
return None;
}
Some(SessionTimezone {
name: Arc::from(parsed.to_string()),
parsed,
})
}
fn rewrite_assignment(plan: LogicalPlan, timezone: &SessionTimezone) -> Result<LogicalPlan> {
let LogicalPlan::Projection(assignment) = plan else {
return Ok(plan);
};
let mut exprs = assignment.expr.clone();
let mut changed = false;
for expr in &mut exprs {
changed |= retarget_assignment_cast(
expr,
assignment.input.schema(),
Some(assignment.input.as_ref()),
timezone,
)?;
}
// The planner types `VALUES` against the target table, so the assignment
// cast lands inside the `Values` rows instead of on the projection above.
let mut input = assignment.input.clone();
if let LogicalPlan::Values(values) = assignment.input.as_ref()
&& let Some(rewritten) = rewrite_values(values, timezone)?
{
input = Arc::new(LogicalPlan::Values(rewritten));
changed = true;
}
if !changed {
return Ok(LogicalPlan::Projection(assignment));
}
Projection::try_new(exprs, input).map(LogicalPlan::Projection)
}
fn rewrite_values(values: &Values, timezone: &SessionTimezone) -> Result<Option<Values>> {
let mut rewritten = values.clone();
let mut changed = false;
for row in &mut rewritten.values {
for expr in row.iter_mut() {
changed |= retarget_assignment_cast(expr, &values.schema, None, timezone)?;
}
}
Ok(changed.then_some(rewritten))
}
/// Reinterprets one assignment cast, returning whether it was rewritten.
///
/// `source_plan` is the projection's input, used to resolve a literal behind a
/// column reference; `Values` rows carry their expression inline and pass `None`.
fn retarget_assignment_cast(
expr: &mut Expr,
schema: &DFSchemaRef,
source_plan: Option<&LogicalPlan>,
timezone: &SessionTimezone,
) -> Result<bool> {
let expr = unalias_mut(expr);
let Expr::Cast(Cast {
expr: source,
data_type: DataType::Timestamp(unit, None),
}) = expr
else {
return Ok(false);
};
let unit = *unit;
if !matches!(
source.get_type(schema)?,
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
) {
return Ok(false);
}
// Fold literals with the same parser the plain `INSERT ... VALUES` path
// uses, so a given string means the same thing however it reaches a column.
// The parsers disagree on ambiguous local times: this one resolves them,
// arrow rejects them.
let folded = source_literal(source.as_ref(), source_plan)
.and_then(|literal| convert_literal(&literal, unit, &timezone.parsed));
if let Some(folded) = folded {
*expr = folded;
return Ok(true);
}
let source = source.as_ref().clone();
*expr = Expr::Cast(Cast::new(
Box::new(Expr::Cast(Cast::new(
Box::new(source),
DataType::Timestamp(unit, Some(timezone.name.clone())),
))),
DataType::Timestamp(unit, None),
));
Ok(true)
}
fn source_literal(source: &Expr, source_plan: Option<&LogicalPlan>) -> Option<ScalarValue> {
match source {
Expr::Literal(value, _) => Some(value.clone()),
Expr::Column(column) => {
let plan = source_plan?;
let index = plan.schema().maybe_index_of_column(column)?;
lineage_literal(plan, index).cloned()
}
_ => None,
}
}
/// Resolves a literal when every row carries the same value at `output_idx`.
///
/// Read-only: the literal is folded into the assignment above, so nodes that
/// drop, reorder or deduplicate rows can be traversed — none of them changes
/// the value a surviving row carries, and folding above them leaves their keys
/// on the original strings.
fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> {
if output_idx >= plan.schema().fields().len() {
return None;
}
match plan {
LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) {
Expr::Literal(value, _) => Some(value),
Expr::Column(column) => {
let input_idx = projection.input.schema().maybe_index_of_column(column)?;
lineage_literal(projection.input.as_ref(), input_idx)
}
_ => None,
},
LogicalPlan::Filter(_)
| LogicalPlan::Sort(_)
| LogicalPlan::Limit(_)
| LogicalPlan::SubqueryAlias(_)
| LogicalPlan::Distinct(Distinct::All(_)) => {
let inputs = plan.inputs();
let [input] = inputs.as_slice() else {
return None;
};
lineage_literal(input, output_idx)
}
_ => None,
}
}
fn convert_literal(value: &ScalarValue, unit: TimeUnit, timezone: &Timezone) -> Option<Expr> {
let ScalarValue::Utf8(Some(value)) = value else {
return None;
};
cast_string_to_timestamp(value, &DataType::Timestamp(unit, None), Some(timezone))
.ok()
.filter(|value| !value.is_null())
.map(|value| Expr::Literal(value, None))
}
fn unalias(expr: &Expr) -> &Expr {
match expr {
Expr::Alias(Alias { expr, .. }) => unalias(expr),
expr => expr,
}
}
fn unalias_mut(expr: &mut Expr) -> &mut Expr {
match expr {
Expr::Alias(Alias { expr, .. }) => unalias_mut(expr),
expr => expr,
}
}
#[cfg(test)]
mod tests {
use datafusion_common::DFSchema;
use datafusion_expr::expr::Placeholder;
use super::*;
fn shanghai() -> SessionTimezone {
let parsed = Timezone::from_tz_string("Asia/Shanghai").unwrap();
SessionTimezone {
name: Arc::from(parsed.to_string()),
parsed,
}
}
/// A prepared `INSERT ... VALUES (?)` arrives here as a cast over an untyped
/// placeholder, which must survive for parameter substitution.
#[test]
fn test_untyped_placeholder_assignment_is_left_alone() {
let schema = Arc::new(DFSchema::empty());
let mut expr = Expr::Cast(Cast::new(
Box::new(Expr::Placeholder(Placeholder::new_with_field(
"$1".to_string(),
None,
))),
DataType::Timestamp(TimeUnit::Millisecond, None),
));
let original = expr.clone();
let changed = retarget_assignment_cast(&mut expr, &schema, None, &shanghai()).unwrap();
assert!(!changed);
assert_eq!(expr, original);
}
}
src/query/src/optimizer/type_conversion.rs
+5 -6
View File
@@ -12,8 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
mod insert_assignment;
use std::sync::Arc;
use common_time::Timezone;
@@ -31,7 +29,7 @@ use session::context::QueryContextRef;
use crate::QueryEngineContext;
use crate::optimizer::ExtensionAnalyzerRule;
use crate::optimizer::type_conversion::insert_assignment::rewrite_insert_assignments;
use crate::optimizer::insert_assignment::rewrite_insert_assignments;
use crate::plan::ExtractExpr;
/// TypeConversionRule converts some literal values in logical plan to other types according
@@ -46,7 +44,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
&self,
plan: LogicalPlan,
ctx: &QueryEngineContext,
_config: &ConfigOptions,
config: &ConfigOptions,
) -> Result<LogicalPlan> {
plan.transform_up_with_subqueries(|plan| match plan {
LogicalPlan::Filter(filter) => {
@@ -129,7 +127,8 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
LogicalPlan::Dml(mut dml) if matches!(dml.op, WriteOp::Insert(_)) => {
dml.input = Arc::new(rewrite_insert_assignments(
dml.input.as_ref().clone(),
ctx.query_ctx(),
&ctx.query_ctx(),
config,
)?);
Ok(Transformed::yes(LogicalPlan::Dml(dml)))
}
@@ -324,7 +323,7 @@ fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr {
)
}
fn cast_string_to_timestamp(
pub(crate) fn cast_string_to_timestamp(
string: &str,
target_type: &DataType,
timezone: Option<&Timezone>,
src/query/src/optimizer/type_conversion/insert_assignment.rs
-336
View File
@@ -1,336 +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.
use std::sync::Arc;
use datafusion_common::{Column, Result, ScalarValue};
use datafusion_expr::expr::{Alias, Cast};
use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Union, Values};
use datatypes::arrow::datatypes::DataType;
use session::context::QueryContextRef;
use crate::optimizer::type_conversion::cast_string_to_timestamp;
use crate::plan::ExtractExpr;
/// Rewrites string literals that feed timestamp columns at an INSERT boundary.
/// Constants are folded at the assignment to avoid changing source types;
/// `VALUES` and `UNION` inputs are rewritten per row or branch. Explicit casts
/// stay on DataFusion's existing path.
pub(super) fn rewrite_insert_assignments(
plan: LogicalPlan,
query_ctx: QueryContextRef,
) -> Result<LogicalPlan> {
let LogicalPlan::Projection(assignment) = plan else {
return Ok(plan);
};
let converter = InsertAssignmentConverter { query_ctx };
let mut exprs = assignment.expr.clone();
let mut input = assignment.input.as_ref().clone();
let mut changed = false;
for (output_idx, expr) in assignment.expr.iter().enumerate() {
let target_type = assignment.schema.field(output_idx).data_type();
if !matches!(target_type, DataType::Timestamp(_, _)) {
continue;
}
let Some(column) = assignment_input_column(expr) else {
continue;
};
let Some(input_idx) = input.schema().maybe_index_of_column(column) else {
continue;
};
let literal = lineage_literal(&input, input_idx).cloned();
if let Some(literal) = literal
&& let Some(folded) = converter.convert_literal(&literal, target_type)
{
let (qualifier, field) = assignment.schema.qualified_field(output_idx);
exprs[output_idx] = folded.alias_qualified(qualifier.cloned(), field.name());
changed = true;
continue;
}
// A hand-built DML plan (e.g. from flow) may share one source column
// between targets; rewriting it in place would retype every reader.
if assignment
.expr
.iter()
.enumerate()
.any(|(idx, other)| idx != output_idx && other.column_refs().contains(column))
{
continue;
}
if let Some(rewritten) = converter.rewrite_output_column(&input, input_idx, target_type)? {
input = rewritten;
changed = true;
}
}
if !changed {
return Ok(LogicalPlan::Projection(assignment));
}
Projection::try_new(exprs, Arc::new(input)).map(LogicalPlan::Projection)
}
/// Resolves a literal when every row carries the same value at `output_idx`.
fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> {
if output_idx >= plan.schema().fields().len() {
return None;
}
match plan {
LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) {
Expr::Literal(value, _) => Some(value),
Expr::Column(column) => {
let input_idx = projection.input.schema().maybe_index_of_column(column)?;
lineage_literal(projection.input.as_ref(), input_idx)
}
_ => None,
},
// These nodes drop, reorder or deduplicate rows without touching the
// value a surviving row carries, and their schema stays positional.
LogicalPlan::Filter(_)
| LogicalPlan::Sort(_)
| LogicalPlan::Limit(_)
| LogicalPlan::SubqueryAlias(_)
| LogicalPlan::Distinct(Distinct::All(_)) => {
let inputs = plan.inputs();
let [input] = inputs.as_slice() else {
return None;
};
lineage_literal(input, output_idx)
}
_ => None,
}
}
struct InsertAssignmentConverter {
query_ctx: QueryContextRef,
}
impl InsertAssignmentConverter {
fn rewrite_output_column(
&self,
plan: &LogicalPlan,
output_idx: usize,
target_type: &DataType,
) -> Result<Option<LogicalPlan>> {
if output_idx >= plan.schema().fields().len() {
return Ok(None);
}
match plan {
// DataFusion pushes INSERT assignment casts into Values, so inspect
// them even though the Values schema already has the target type.
LogicalPlan::Values(values) => self.rewrite_values(values, output_idx, target_type),
// Once the source query has produced the target type, its casts and
// coercions belong to the source query rather than INSERT assignment.
_ if plan.schema().field(output_idx).data_type() == target_type => Ok(None),
LogicalPlan::Projection(projection) => {
self.rewrite_projection(projection, output_idx, target_type)
}
LogicalPlan::Union(union) => self.rewrite_union(union, output_idx, target_type),
// Filter and Sort are excluded here: their predicates and keys read
// the retyped column, which would change the source query. Constants
// still reach the assignment through `lineage_literal`. Distinct::All
// holds no expressions; allowing it shifts dedup keys from raw
// strings to parsed instants.
LogicalPlan::Limit(_)
| LogicalPlan::SubqueryAlias(_)
| LogicalPlan::Distinct(Distinct::All(_)) => {
self.rewrite_passthrough(plan, output_idx, target_type)
}
_ => Ok(None),
}
}
fn rewrite_projection(
&self,
projection: &Projection,
output_idx: usize,
target_type: &DataType,
) -> Result<Option<LogicalPlan>> {
match unalias(&projection.expr[output_idx]) {
Expr::Literal(value, _) => {
let Some(converted) = self.convert_literal(value, target_type) else {
return Ok(None);
};
let (qualifier, field) = projection.schema.qualified_field(output_idx);
let mut exprs = projection.expr.clone();
exprs[output_idx] = converted.alias_qualified(qualifier.cloned(), field.name());
Projection::try_new(exprs, projection.input.clone())
.map(LogicalPlan::Projection)
.map(Some)
}
Expr::Column(column) => {
// Retyping the input column affects every output column reading
// it, so only follow lineage with a single consumer.
if projection
.expr
.iter()
.enumerate()
.any(|(idx, other)| idx != output_idx && other.column_refs().contains(column))
{
return Ok(None);
}
let input = projection.input.as_ref();
let Some(input_idx) = input.schema().maybe_index_of_column(column) else {
return Ok(None);
};
let Some(rewritten) = self.rewrite_output_column(input, input_idx, target_type)?
else {
return Ok(None);
};
Projection::try_new(projection.expr.clone(), Arc::new(rewritten))
.map(LogicalPlan::Projection)
.map(Some)
}
_ => Ok(None),
}
}
fn rewrite_values(
&self,
values: &Values,
output_idx: usize,
target_type: &DataType,
) -> Result<Option<LogicalPlan>> {
let mut rewritten = values.clone();
let mut changed = false;
for row in &mut rewritten.values {
let Some(expr) = row.get_mut(output_idx) else {
return Ok(None);
};
if let Expr::Cast(Cast {
expr: inner,
data_type,
}) = expr
&& data_type == target_type
&& let Expr::Literal(value, _) = inner.as_ref()
&& let Some(value) = self.convert_literal(value, target_type)
{
*expr = value;
changed = true;
continue;
}
if expr.get_type(values.schema.as_ref()).ok().as_ref() != Some(target_type) {
return Ok(None);
}
}
Ok(changed.then_some(LogicalPlan::Values(rewritten)))
}
fn rewrite_union(
&self,
union: &Union,
output_idx: usize,
target_type: &DataType,
) -> Result<Option<LogicalPlan>> {
let mut inputs = Vec::with_capacity(union.inputs.len());
// A union has one shared schema, so rewrite it only when the same output
// column is an INSERT-assigned literal in every branch.
for input in &union.inputs {
let Some(rewritten) = self.rewrite_output_column(input, output_idx, target_type)?
else {
return Ok(None);
};
inputs.push(Arc::new(rewritten));
}
// Untouched columns may still have mismatched branch types before
// TypeCoercion runs, so rebuild loosely like the SQL planner does.
Union::try_new_with_loose_types(inputs)
.map(LogicalPlan::Union)
.map(Some)
}
fn rewrite_passthrough(
&self,
plan: &LogicalPlan,
output_idx: usize,
target_type: &DataType,
) -> Result<Option<LogicalPlan>> {
let inputs = plan.inputs();
let [input] = inputs.as_slice() else {
return Ok(None);
};
let Some(rewritten_input) = self.rewrite_output_column(input, output_idx, target_type)?
else {
return Ok(None);
};
plan.with_new_exprs(plan.expressions_consider_join(), vec![rewritten_input])
.map(Some)
}
fn convert_literal(&self, value: &ScalarValue, target_type: &DataType) -> Option<Expr> {
let ScalarValue::Utf8(Some(value)) = value else {
return None;
};
cast_string_to_timestamp(value, target_type, Some(&self.query_ctx.timezone()))
.ok()
.filter(|value| !value.is_null())
.map(|value| Expr::Literal(value, None))
}
}
fn assignment_input_column(expr: &Expr) -> Option<&Column> {
let expr = match unalias(expr) {
Expr::Cast(Cast { expr, .. }) => expr.as_ref(),
expr => expr,
};
let Expr::Column(column) = expr else {
return None;
};
Some(column)
}
fn unalias(expr: &Expr) -> &Expr {
match expr {
Expr::Alias(Alias { expr, .. }) => unalias(expr),
expr => expr,
}
}
#[cfg(test)]
mod tests {
use datafusion_common::ScalarValue;
use datafusion_common::arrow::datatypes::TimeUnit;
use session::context::QueryContext;
use super::*;
#[test]
fn test_convert_literal_falls_back_for_unsupported_literal() {
let converter = InsertAssignmentConverter {
query_ctx: QueryContext::arc(),
};
let target_type = DataType::Timestamp(TimeUnit::Nanosecond, None);
for literal in ["1970-01-01", "-8-01-01 00:00:01.5"] {
assert_eq!(
converter
.convert_literal(&ScalarValue::Utf8(Some(literal.to_string())), &target_type,),
None
);
}
}
}
src/query/src/planner.rs
+64 -62
View File
@@ -866,6 +866,30 @@ mod tests {
engine.planner().plan(&stmt, query_ctx).await.unwrap()
}
/// Plans `sql` and runs the DataFusion analyzer, which is where
/// `InsertAssignmentRule` sits. Planning alone stops short of it, so these
/// assertions would not see the assignment rewrite at all.
async fn analyze_insert(
engine: &QueryEngineRef,
sql: &str,
query_ctx: &QueryContextRef,
) -> String {
let stmt = QueryLanguageParser::parse_sql(sql, query_ctx).unwrap();
let plan = engine
.planner()
.plan(&stmt, query_ctx.clone())
.await
.unwrap();
let context = engine.engine_context(query_ctx.clone());
let state = context.state();
state
.analyzer()
.execute_and_check(plan, state.config_options(), |_, _| {})
.unwrap()
.display_indent_schema()
.to_string()
}
#[tokio::test]
async fn test_insert_timestamp_literals_use_query_timezone() {
let query_ctx = Arc::new(
@@ -898,20 +922,6 @@ mod tests {
) AS source",
&[1_785_902_400_001_i64][..],
),
(
"INSERT INTO timestamps (ts, st) \
SELECT '2026-08-06 12:00:00.001', now() \
UNION ALL \
SELECT '2026-08-07 12:00:00.001', now()",
&[1_785_988_800_001_i64, 1_786_075_200_001_i64][..],
),
(
"INSERT INTO timestamps (ts, st) \
SELECT '2026-08-16 12:00:00.001', now() \
UNION \
SELECT '2026-08-17 12:00:00.001', now()",
&[1_786_852_800_001_i64, 1_786_939_200_001_i64][..],
),
(
"INSERT INTO timestamps (ts, st) \
SELECT '2026-08-18 12:00:00.001', max(st) \
@@ -940,14 +950,7 @@ mod tests {
&[1_786_766_400_001_i64][..],
),
] {
let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap();
let plan = engine
.planner()
.plan(&stmt, query_ctx.clone())
.await
.unwrap()
.display_indent()
.to_string();
let plan = analyze_insert(&engine, sql, &query_ctx).await;
for expected_timestamp in expected_timestamps {
assert!(
@@ -968,19 +971,19 @@ mod tests {
let engine = create_timestamp_test_engine().await;
let sql = "INSERT INTO timestamps (ts, st) \
VALUES (CAST('2026-08-08 12:00:00.001' AS TIMESTAMP), now())";
let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap();
let plan = engine
.planner()
.plan(&stmt, query_ctx)
.await
.unwrap()
.display_indent()
.to_string();
let plan = analyze_insert(&engine, sql, &query_ctx).await;
// An explicit cast reaches the analyzer as an `arrow_cast` call rather
// than an `Expr::Cast`, which is how it stays out of the rewrite.
assert!(
plan.contains("arrow_cast(Utf8(\"2026-08-08 12:00:00.001\")"),
"{plan}"
);
// 12:00:00.001 read as Shanghai local time; the source query keeps UTC.
assert!(
!plan.contains("TimestampMillisecond(1786104000001, None)"),
"{plan}"
);
}
#[tokio::test]
@@ -1017,14 +1020,7 @@ mod tests {
][..],
),
] {
let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap();
let plan = engine
.planner()
.plan(&stmt, query_ctx.clone())
.await
.unwrap()
.display_indent_schema()
.to_string();
let plan = analyze_insert(&engine, sql, &query_ctx).await;
for expected in expected {
assert!(plan.contains(expected), "{plan}");
@@ -1033,30 +1029,37 @@ mod tests {
}
#[tokio::test]
async fn test_insert_union_tolerates_uncoerced_untouched_column() {
async fn test_insert_union_converts_via_assignment_cast() {
let query_ctx = Arc::new(
QueryContextBuilder::default()
.timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap())
.build(),
);
let engine = create_timestamp_test_engine().await;
// `st` stays Timestamp vs Null across branches until TypeCoercion runs.
let sql = "INSERT INTO timestamps (ts, st) \
SELECT '2026-08-06 12:00:00.001', now() \
UNION ALL \
SELECT '2026-08-07 12:00:00.001', NULL";
let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap();
let plan = engine
.planner()
.plan(&stmt, query_ctx)
.await
.unwrap()
.display_indent()
.to_string();
// Branches disagree, so the conversion stays as a cast on the
// assignment instead of folding. One cast covers every branch, which is
// why a NULL branch no longer cancels the conversion for the column and
// why UNION's dedup keys stay on the original strings.
for sql in [
"INSERT INTO timestamps (ts, st) \
SELECT '2026-08-06 12:00:00.001', now() \
UNION ALL \
SELECT '2026-08-07 12:00:00.001', NULL",
"INSERT INTO timestamps (ts, st) \
SELECT '2026-08-16 12:00:00.001', now() \
UNION \
SELECT '2026-08-17 12:00:00.001', now()",
] {
let plan = analyze_insert(&engine, sql, &query_ctx).await;
for expected_timestamp in [1_785_988_800_001_i64, 1_786_075_200_001_i64] {
assert!(
plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")),
plan.contains("AS Timestamp(ms, \"Asia/Shanghai\")"),
"{plan}"
);
// The branches themselves are untouched.
assert!(
plan.contains("Utf8(\"2026-08-07 12:00:00.001\")")
|| plan.contains("Utf8(\"2026-08-17 12:00:00.001\")"),
"{plan}"
);
}
@@ -1074,14 +1077,8 @@ mod tests {
SELECT '2026-08-10 12:00:00.001', now() \
UNION ALL \
SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now()";
let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap();
let plan = engine
.planner()
.plan(&stmt, query_ctx)
.await
.unwrap()
.display_indent_schema()
.to_string();
let plan = analyze_insert(&engine, sql, &query_ctx).await;
assert!(
!plan.contains("TimestampMillisecond(1786334400001, None)"),
"{plan}"
@@ -1090,6 +1087,11 @@ mod tests {
plan.contains("arrow_cast(Utf8(\"2026-08-11 12:00:00.001\")"),
"{plan}"
);
// TypeCoercion has already settled this union to timestamp, so the
// assignment has nothing left to reinterpret. Retargeting the cast here
// would leave a Timestamp(None) -> Timestamp(Some(tz)) step behind,
// which shifts the value instead of relabelling it.
assert!(!plan.contains("Asia/Shanghai"), "{plan}");
}
#[tokio::test]
src/query/src/sql.rs
+206 -43
View File
@@ -41,6 +41,7 @@ use common_recordbatch::RecordBatches;
use common_recordbatch::adapter::RecordBatchStreamAdapter;
use common_time::Timestamp;
use common_time::timezone::get_timezone;
use datafusion::dataframe::DataFrame;
use datafusion::prelude::SessionContext;
use datafusion_expr::{Expr, SortExpr, col, lit};
use datatypes::prelude::*;
@@ -109,7 +110,7 @@ const INDEX_KEY_NAME_COLUMN: &str = "Key_name";
const INDEX_SEQ_IN_INDEX_COLUMN: &str = "Seq_in_index";
const INDEX_COLUMN_NAME_COLUMN: &str = "Column_name";
static DESCRIBE_TABLE_OUTPUT_SCHEMA: Lazy<Arc<Schema>> = Lazy::new(|| {
pub static DESCRIBE_TABLE_OUTPUT_SCHEMA: Lazy<Arc<Schema>> = Lazy::new(|| {
Arc::new(Schema::new(vec![
ColumnSchema::new(
COLUMN_NAME_COLUMN,
@@ -178,6 +179,18 @@ pub async fn show_databases(
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<Output> {
let dataframe =
show_databases_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?;
dataframe_to_output(dataframe).await
}
/// Builds the [`DataFrame`] for `SHOW DATABASES` without executing it.
pub async fn show_databases_dataframe(
stmt: &ShowDatabases,
query_engine: &QueryEngineRef,
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<DataFrame> {
let projects = if stmt.full {
vec![
(schemata::SCHEMA_NAME, SCHEMAS_COLUMN),
@@ -191,7 +204,7 @@ pub async fn show_databases(
let like_field = Some(schemata::SCHEMA_NAME);
let sort = vec![col(schemata::SCHEMA_NAME).sort(true, true)];
query_from_information_schema_table(
query_from_information_schema_dataframe(
query_engine,
catalog_manager,
query_ctx,
@@ -201,7 +214,7 @@ pub async fn show_databases(
filters,
like_field,
sort,
stmt.kind,
&stmt.kind,
)
.await
}
@@ -249,6 +262,45 @@ async fn query_from_information_schema_table(
sort: Vec<SortExpr>,
kind: ShowKind,
) -> Result<Output> {
let dataframe = query_from_information_schema_dataframe(
query_engine,
catalog_manager,
query_ctx,
table_name,
select,
projects,
filters,
like_field,
sort,
&kind,
)
.await?;
dataframe_to_output(dataframe).await
}
async fn dataframe_to_output(dataframe: DataFrame) -> Result<Output> {
let stream = dataframe.execute_stream().await?;
Ok(Output::new_with_stream(Box::pin(
RecordBatchStreamAdapter::try_new(stream).context(error::CreateRecordBatchSnafu)?,
)))
}
/// Builds the [`DataFrame`] for a `SHOW` statement without executing it,
/// so `Describe` handlers can derive the output schema from the same
/// projection the executor uses.
#[allow(clippy::too_many_arguments)]
async fn query_from_information_schema_dataframe(
query_engine: &QueryEngineRef,
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
table_name: &str,
select: Vec<Expr>,
projects: Vec<(&str, &str)>,
filters: Vec<Expr>,
like_field: Option<&str>,
sort: Vec<SortExpr>,
kind: &ShowKind,
) -> Result<DataFrame> {
let table = catalog_manager
.table(
query_ctx.current_catalog(),
@@ -336,7 +388,7 @@ async fn query_from_information_schema_table(
.expect("Must be the datafusion planner");
let filter = planner
.sql_to_expr(filter, dataframe.schema(), false, query_ctx)
.sql_to_expr(filter.clone(), dataframe.schema(), false, query_ctx)
.await?;
// Apply the `where` clause filters
@@ -344,11 +396,7 @@ async fn query_from_information_schema_table(
}
};
let stream = dataframe.execute_stream().await?;
Ok(Output::new_with_stream(Box::pin(
RecordBatchStreamAdapter::try_new(stream).context(error::CreateRecordBatchSnafu)?,
)))
Ok(dataframe)
}
/// Execute `SHOW COLUMNS` statement.
@@ -358,8 +406,19 @@ pub async fn show_columns(
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<Output> {
let schema_name = if let Some(database) = stmt.database {
database
let dataframe = show_columns_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?;
dataframe_to_output(dataframe).await
}
/// Builds the [`DataFrame`] for `SHOW COLUMNS` without executing it.
pub async fn show_columns_dataframe(
stmt: &ShowColumns,
query_engine: &QueryEngineRef,
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<DataFrame> {
let schema_name = if let Some(database) = &stmt.database {
database.clone()
} else {
query_ctx.current_schema()
};
@@ -397,7 +456,7 @@ pub async fn show_columns(
let like_field = Some(columns::COLUMN_NAME);
let sort = vec![col(columns::COLUMN_NAME).sort(true, true)];
query_from_information_schema_table(
query_from_information_schema_dataframe(
query_engine,
catalog_manager,
query_ctx,
@@ -407,7 +466,7 @@ pub async fn show_columns(
filters,
like_field,
sort,
stmt.kind,
&stmt.kind,
)
.await
}
@@ -419,8 +478,19 @@ pub async fn show_index(
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<Output> {
let schema_name = if let Some(database) = stmt.database {
database
let dataframe = show_index_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?;
dataframe_to_output(dataframe).await
}
/// Builds the [`DataFrame`] for `SHOW INDEX` without executing it.
pub async fn show_index_dataframe(
stmt: &ShowIndex,
query_engine: &QueryEngineRef,
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<DataFrame> {
let schema_name = if let Some(database) = &stmt.database {
database.clone()
} else {
query_ctx.current_schema()
};
@@ -471,7 +541,7 @@ pub async fn show_index(
col(statistics::SEQ_IN_INDEX).sort(true, true),
];
query_from_information_schema_table(
query_from_information_schema_dataframe(
query_engine,
catalog_manager,
query_ctx,
@@ -481,7 +551,7 @@ pub async fn show_index(
filters,
like_field,
sort,
stmt.kind,
&stmt.kind,
)
.await
}
@@ -493,8 +563,19 @@ pub async fn show_region(
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<Output> {
let schema_name = if let Some(database) = stmt.database {
database
let dataframe = show_region_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?;
dataframe_to_output(dataframe).await
}
/// Builds the [`DataFrame`] for `SHOW REGION` without executing it.
pub async fn show_region_dataframe(
stmt: &ShowRegion,
query_engine: &QueryEngineRef,
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<DataFrame> {
let schema_name = if let Some(database) = &stmt.database {
database.clone()
} else {
query_ctx.current_schema()
};
@@ -517,7 +598,7 @@ pub async fn show_region(
col(columns::PEER_ID).sort(true, true),
];
query_from_information_schema_table(
query_from_information_schema_dataframe(
query_engine,
catalog_manager,
query_ctx,
@@ -527,7 +608,7 @@ pub async fn show_region(
filters,
like_field,
sort,
stmt.kind,
&stmt.kind,
)
.await
}
@@ -539,8 +620,19 @@ pub async fn show_tables(
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<Output> {
let schema_name = if let Some(database) = stmt.database {
database
let dataframe = show_tables_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?;
dataframe_to_output(dataframe).await
}
/// Builds the [`DataFrame`] for [`ShowTables`] without executing it.
pub async fn show_tables_dataframe(
stmt: &ShowTables,
query_engine: &QueryEngineRef,
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<DataFrame> {
let schema_name = if let Some(database) = &stmt.database {
database.clone()
} else {
query_ctx.current_schema()
};
@@ -564,15 +656,16 @@ pub async fn show_tables(
// Transform the WHERE clause for backward compatibility:
// Replace "Tables" with "Tables_in_{schema}" to support old queries
let kind = match stmt.kind {
ShowKind::Where(mut filter) => {
let rewritten_kind = match &stmt.kind {
ShowKind::Where(filter) => {
let mut filter = filter.clone();
replace_column_in_expr(&mut filter, "Tables", &tables_column);
ShowKind::Where(filter)
}
other => other,
kind => kind.clone(),
};
query_from_information_schema_table(
query_from_information_schema_dataframe(
query_engine,
catalog_manager,
query_ctx,
@@ -582,7 +675,7 @@ pub async fn show_tables(
filters,
like_field,
sort,
kind,
&rewritten_kind,
)
.await
}
@@ -594,8 +687,20 @@ pub async fn show_table_status(
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<Output> {
let schema_name = if let Some(database) = stmt.database {
database
let dataframe =
show_table_status_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?;
dataframe_to_output(dataframe).await
}
/// Builds the [`DataFrame`] for [`ShowTableStatus`] without executing it.
pub async fn show_table_status_dataframe(
stmt: &ShowTableStatus,
query_engine: &QueryEngineRef,
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<DataFrame> {
let schema_name = if let Some(database) = &stmt.database {
database.clone()
} else {
query_ctx.current_schema()
};
@@ -629,7 +734,7 @@ pub async fn show_table_status(
let like_field = Some(tables::TABLE_NAME);
let sort = vec![col(tables::TABLE_NAME).sort(true, true)];
query_from_information_schema_table(
query_from_information_schema_dataframe(
query_engine,
catalog_manager,
query_ctx,
@@ -639,7 +744,7 @@ pub async fn show_table_status(
filters,
like_field,
sort,
stmt.kind,
&stmt.kind,
)
.await
}
@@ -651,6 +756,18 @@ pub async fn show_collations(
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<Output> {
let dataframe =
show_collations_dataframe(&kind, query_engine, catalog_manager, query_ctx).await?;
dataframe_to_output(dataframe).await
}
/// Builds the [`DataFrame`] for `SHOW COLLATION` without executing it.
pub async fn show_collations_dataframe(
kind: &ShowKind,
query_engine: &QueryEngineRef,
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<DataFrame> {
// Refer to https://dev.mysql.com/doc/refman/8.0/en/show-collation.html
let projects = vec![
("collation_name", "Collation"),
@@ -665,7 +782,7 @@ pub async fn show_collations(
let like_field = Some("collation_name");
let sort = vec![];
query_from_information_schema_table(
query_from_information_schema_dataframe(
query_engine,
catalog_manager,
query_ctx,
@@ -687,6 +804,18 @@ pub async fn show_charsets(
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<Output> {
let dataframe =
show_charsets_dataframe(&kind, query_engine, catalog_manager, query_ctx).await?;
dataframe_to_output(dataframe).await
}
/// Builds the [`DataFrame`] for `SHOW CHARSET` without executing it.
pub async fn show_charsets_dataframe(
kind: &ShowKind,
query_engine: &QueryEngineRef,
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<DataFrame> {
// Refer to https://dev.mysql.com/doc/refman/8.0/en/show-character-set.html
let projects = vec![
("character_set_name", "Charset"),
@@ -699,7 +828,7 @@ pub async fn show_charsets(
let like_field = Some("character_set_name");
let sort = vec![];
query_from_information_schema_table(
query_from_information_schema_dataframe(
query_engine,
catalog_manager,
query_ctx,
@@ -916,8 +1045,19 @@ pub async fn show_views(
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<Output> {
let schema_name = if let Some(database) = stmt.database {
database
let dataframe = show_views_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?;
dataframe_to_output(dataframe).await
}
/// Builds the [`DataFrame`] for [`ShowViews`] without executing it.
pub async fn show_views_dataframe(
stmt: &ShowViews,
query_engine: &QueryEngineRef,
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<DataFrame> {
let schema_name = if let Some(database) = &stmt.database {
database.clone()
} else {
query_ctx.current_schema()
};
@@ -930,7 +1070,7 @@ pub async fn show_views(
let like_field = Some(tables::TABLE_NAME);
let sort = vec![col(tables::TABLE_NAME).sort(true, true)];
query_from_information_schema_table(
query_from_information_schema_dataframe(
query_engine,
catalog_manager,
query_ctx,
@@ -940,7 +1080,7 @@ pub async fn show_views(
filters,
like_field,
sort,
stmt.kind,
&stmt.kind,
)
.await
}
@@ -952,12 +1092,23 @@ pub async fn show_flows(
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<Output> {
let dataframe = show_flows_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?;
dataframe_to_output(dataframe).await
}
/// Builds the [`DataFrame`] for [`ShowFlows`] without executing it.
pub async fn show_flows_dataframe(
stmt: &ShowFlows,
query_engine: &QueryEngineRef,
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<DataFrame> {
let projects = vec![(flows::FLOW_NAME, FLOWS_COLUMN)];
let filters = vec![col(flows::TABLE_CATALOG).eq(lit(query_ctx.current_catalog()))];
let like_field = Some(flows::FLOW_NAME);
let sort = vec![col(flows::FLOW_NAME).sort(true, true)];
query_from_information_schema_table(
query_from_information_schema_dataframe(
query_engine,
catalog_manager,
query_ctx,
@@ -967,7 +1118,7 @@ pub async fn show_flows(
filters,
like_field,
sort,
stmt.kind,
&stmt.kind,
)
.await
}
@@ -1340,6 +1491,18 @@ pub async fn show_processlist(
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<Output> {
let dataframe =
show_processlist_dataframe(&stmt, query_engine, catalog_manager, query_ctx).await?;
dataframe_to_output(dataframe).await
}
/// Builds the [`DataFrame`] for `SHOW PROCESSLIST` without executing it.
pub async fn show_processlist_dataframe(
stmt: &ShowProcessList,
query_engine: &QueryEngineRef,
catalog_manager: &CatalogManagerRef,
query_ctx: QueryContextRef,
) -> Result<DataFrame> {
let projects = if stmt.full {
vec![
(process_list::ID, "Id"),
@@ -1367,17 +1530,17 @@ pub async fn show_processlist(
};
let like_field = None;
let sort = vec![col("id").sort(true, true)];
query_from_information_schema_table(
query_from_information_schema_dataframe(
query_engine,
catalog_manager,
query_ctx.clone(),
"process_list",
vec![],
projects.clone(),
projects,
filters,
like_field,
sort,
ShowKind::All,
&ShowKind::All,
)
.await
}
src/servers/Cargo.toml
+2 -2
View File
@@ -167,8 +167,8 @@ servers = { workspace = true, features = ["testing"] }
session = { workspace = true, features = ["testing"] }
table.workspace = true
tempfile = "3.0.0"
tokio-postgres = "0.7"
tokio-postgres-rustls = "0.12"
tokio-postgres.workspace = true
tokio-postgres-rustls = "0.14"
[target.'cfg(unix)'.dev-dependencies]
pprof = { version = "0.14", features = ["criterion", "flamegraph"] }
src/servers/dashboard/VERSION
+1 -1
View File
@@ -1 +1 @@
v0.13.13
v0.13.14
src/servers/src/grpc/flight.rs
+40 -35
View File
@@ -30,8 +30,7 @@ use bytes::{self, Bytes};
use common_error::ext::ErrorExt;
use common_grpc::flight::do_put::{DoPutMetadata, DoPutResponse};
use common_grpc::flight::{
FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightEncoder, FlightMessage,
SNAPSHOT_SEQS_METADATA_KEY,
FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage, SNAPSHOT_SEQS_METADATA_KEY,
};
use common_memory_manager::MemoryGuard;
use common_query::{Output, OutputData};
@@ -46,14 +45,16 @@ use prost::Message;
use query::metrics::terminal_recordbatch_metrics_from_plan_if_requested;
use query::options::FlowQueryExtensions;
use session::context::{Channel, QueryContextRef};
use snafu::{IntoError, ResultExt, ensure};
use snafu::{IntoError, OptionExt, ResultExt, ensure};
use table::table_name::TableName;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tonic::{Request, Response, Status, Streaming};
use crate::error::{InvalidParameterSnafu, Result, ToJsonSnafu};
pub use crate::grpc::flight::stream::{FlightRecordBatchSource, FlightRecordBatchStream};
use crate::error::{InvalidParameterSnafu, InvalidQuerySnafu, Result, ToJsonSnafu};
pub use crate::grpc::flight::stream::{
FlightRecordBatchSource, FlightRecordBatchStream, FlightRecordBatchStreamInput,
};
use crate::grpc::greptime_handler::{
GreptimeRequestHandler, create_query_context, get_request_type,
};
@@ -221,9 +222,12 @@ impl FlightCraft for GreptimeRequestHandler {
);
let flight_compression = self.flight_compression;
async {
let output = self
.handle_request_with_query_ctx(request, query_ctx.clone())
let query = request.request.context(InvalidQuerySnafu {
reason: "Expecting non-empty GreptimeRequest.",
})?;
self.authenticate_request_with_query_ctx(request.header.as_ref(), &query_ctx)
.await?;
let output = self.handle_request_with_query_ctx(query, query_ctx.clone());
let stream = to_flight_data_stream(
output,
TracingContext::from_current_span(),
@@ -573,32 +577,36 @@ fn extract_json_metadata<T: serde::de::DeserializeOwned>(
Ok(Some(parsed))
}
fn to_flight_data_stream(
output: Output,
fn to_flight_data_stream<F>(
output: F,
tracing_context: TracingContext,
flight_compression: FlightCompression,
query_ctx: QueryContextRef,
should_emit_terminal_metrics: bool,
) -> TonicStream<FlightData> {
) -> TonicStream<FlightData>
where
F: std::future::Future<Output = Result<Output>> + Send + 'static,
{
let initializer = async move {
let output = output.await.map_err(Status::from)?;
output_to_flight_record_batch_source(output, should_emit_terminal_metrics)
};
let stream = FlightRecordBatchStream::new(
FlightRecordBatchStreamInput::initializer(initializer),
tracing_context,
flight_compression,
query_ctx,
);
Box::pin(stream) as _
}
fn output_to_flight_record_batch_source(
output: Output,
should_emit_terminal_metrics: bool,
) -> TonicResult<FlightRecordBatchSource> {
match output.data {
OutputData::Stream(stream) => {
let stream = FlightRecordBatchStream::new(
FlightRecordBatchSource::RecordBatches(stream),
tracing_context,
flight_compression,
query_ctx,
);
Box::pin(stream) as _
}
OutputData::RecordBatches(x) => {
let stream = FlightRecordBatchStream::new(
FlightRecordBatchSource::RecordBatches(x.as_stream()),
tracing_context,
flight_compression,
query_ctx,
);
Box::pin(stream) as _
}
OutputData::Stream(stream) => Ok(FlightRecordBatchSource::RecordBatches(stream)),
OutputData::RecordBatches(x) => Ok(FlightRecordBatchSource::RecordBatches(x.as_stream())),
OutputData::AffectedRows(rows) => {
let terminal_metrics = match terminal_recordbatch_metrics_from_plan_if_requested(
output.meta.plan,
@@ -607,20 +615,17 @@ fn to_flight_data_stream(
Some(metrics) => match serde_json::to_string(&metrics) {
Ok(metrics) => Some(metrics),
Err(e) => {
let stream = tokio_stream::once(Err(Status::internal(format!(
return Err(Status::internal(format!(
"Failed to serialize terminal metrics: {e}"
))));
return Box::pin(stream) as _;
)));
}
},
None => None,
};
let affected_rows = FlightEncoder::default().encode(FlightMessage::AffectedRows {
Ok(FlightRecordBatchSource::AffectedRows {
rows,
metrics: terminal_metrics,
});
let stream = tokio_stream::iter(affected_rows.into_iter().map(Ok));
Box::pin(stream) as _
})
}
}
}
src/servers/src/grpc/flight/stream.rs
+106 -44
View File
@@ -41,27 +41,34 @@ use crate::error;
use crate::grpc::FlightCompression;
use crate::grpc::flight::TonicResult;
pub type FlightRecordBatchStreamInitializer = Pin<
Box<
dyn Future<Output = std::result::Result<SendableRecordBatchStream, tonic::Status>>
+ Send
+ 'static,
>,
>;
pub enum FlightRecordBatchSource {
RecordBatches(SendableRecordBatchStream),
Initializer(FlightRecordBatchStreamInitializer),
AffectedRows {
rows: usize,
metrics: Option<String>,
},
}
impl FlightRecordBatchSource {
pub fn initializer<F>(initializer: F) -> Self
where
F: Future<Output = std::result::Result<SendableRecordBatchStream, tonic::Status>>
+ Send
+ 'static,
{
Self::Initializer(Box::pin(initializer))
/// Determines whether a Flight result is ready now or initialized asynchronously.
pub enum FlightRecordBatchStreamInput<F = std::future::Ready<TonicResult<FlightRecordBatchSource>>>
{
Ready(FlightRecordBatchSource),
Initializer(F),
}
impl FlightRecordBatchStreamInput {
/// Creates an input from a source that is already available.
pub fn ready(source: FlightRecordBatchSource) -> Self {
Self::Ready(source)
}
}
impl<F> FlightRecordBatchStreamInput<F> {
/// Creates an input that obtains its source asynchronously.
///
/// Errors from the initializer are returned through the Flight response stream.
pub fn initializer(initializer: F) -> Self {
Self::Initializer(initializer)
}
}
@@ -160,31 +167,32 @@ impl FlightRecordBatchStream {
Self::send_metrics(tx, metrics, metrics_str).await
}
pub fn new(
source: FlightRecordBatchSource,
pub fn new<F>(
input: FlightRecordBatchStreamInput<F>,
tracing_context: TracingContext,
compression: FlightCompression,
query_ctx: QueryContextRef,
) -> Self {
let should_send_partial_metrics = query_ctx.explain_verbose();
let can_send_metrics_before_batch = query_ctx.explain_verbose()
&& query_ctx.live_analyze_metrics_enabled()
&& query_ctx
.remote_query_id()
.zip(query_ctx.extension(SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY))
.is_some_and(|(remote_query_id, capability)| capability == remote_query_id);
) -> Self
where
F: Future<Output = TonicResult<FlightRecordBatchSource>> + Send + 'static,
{
let (mut tx, rx) = mpsc::channel::<TonicResult<FlightMessage>>(1);
let source_type = match &source {
FlightRecordBatchSource::RecordBatches(_) => "record_batches",
FlightRecordBatchSource::Initializer(_) => "initializer",
let source_type = match &input {
FlightRecordBatchStreamInput::Ready(FlightRecordBatchSource::RecordBatches(_)) => {
"record_batches"
}
FlightRecordBatchStreamInput::Ready(FlightRecordBatchSource::AffectedRows {
..
}) => "affected_rows",
FlightRecordBatchStreamInput::Initializer(_) => "initializer",
};
let initializer_tracing_context = tracing_context.clone();
let join_handle = common_runtime::spawn_global(
async move {
let recordbatches = async move {
match source {
FlightRecordBatchSource::RecordBatches(recordbatches) => Ok(recordbatches),
FlightRecordBatchSource::Initializer(initializer) => initializer.await,
let source = async move {
match input {
FlightRecordBatchStreamInput::Ready(source) => Ok(source),
FlightRecordBatchStreamInput::Initializer(initializer) => initializer.await,
}
}
.trace(
@@ -193,8 +201,20 @@ impl FlightRecordBatchStream {
)
.await;
match recordbatches {
Ok(recordbatches) => {
match source {
Ok(FlightRecordBatchSource::RecordBatches(recordbatches)) => {
let should_send_partial_metrics = query_ctx.explain_verbose();
let can_send_metrics_before_batch =
query_ctx.explain_verbose()
&& query_ctx.live_analyze_metrics_enabled()
&& query_ctx
.remote_query_id()
.zip(query_ctx.extension(
SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY,
))
.is_some_and(|(remote_query_id, capability)| {
capability == remote_query_id
});
Self::flight_data_stream(
recordbatches,
tx,
@@ -203,6 +223,11 @@ impl FlightRecordBatchStream {
)
.await;
}
Ok(FlightRecordBatchSource::AffectedRows { rows, metrics }) => {
let _ = tx
.send(Ok(FlightMessage::AffectedRows { rows, metrics }))
.await;
}
Err(status) => {
let _ = tx.send(Err(status)).await;
}
@@ -485,7 +510,9 @@ mod test {
.unwrap()
.as_stream();
let mut stream = FlightRecordBatchStream::new(
FlightRecordBatchSource::RecordBatches(recordbatches),
FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches(
recordbatches,
)),
TracingContext::default(),
FlightCompression::default(),
QueryContext::arc(),
@@ -519,10 +546,34 @@ mod test {
}
}
#[tokio::test]
async fn test_flight_record_batch_stream_encodes_affected_rows() {
let mut stream = FlightRecordBatchStream::new(
FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::AffectedRows {
rows: 42,
metrics: Some(r#"{"region_watermarks":[]}"#.to_string()),
}),
TracingContext::default(),
FlightCompression::default(),
QueryContext::arc(),
);
let data = stream.next().await.unwrap().unwrap();
let message = FlightDecoder::default().try_decode(&data).unwrap().unwrap();
assert!(matches!(
message,
FlightMessage::AffectedRows {
rows: 42,
metrics: Some(_),
}
));
assert!(stream.next().await.is_none());
}
#[tokio::test]
async fn test_flight_record_batch_stream_forwards_initializer_error() {
let mut stream = FlightRecordBatchStream::new(
FlightRecordBatchSource::initializer(async {
FlightRecordBatchStreamInput::initializer(async {
Err(tonic::Status::unavailable(
"remote read initialization failed",
))
@@ -552,9 +603,12 @@ mod test {
metrics,
});
let query_ctx = query_context_with_live_metrics_and_matching_capability();
query_ctx.set_explain_verbose(true);
let initializer_query_ctx = query_ctx.clone();
let mut stream = FlightRecordBatchStream::new(
FlightRecordBatchSource::RecordBatches(recordbatches),
FlightRecordBatchStreamInput::initializer(async move {
initializer_query_ctx.set_explain_verbose(true);
Ok(FlightRecordBatchSource::RecordBatches(recordbatches))
}),
TracingContext::default(),
FlightCompression::default(),
query_ctx,
@@ -609,7 +663,9 @@ mod test {
let query_ctx = query_context_with_live_metrics_and_matching_capability();
query_ctx.set_explain_verbose(true);
let mut stream = FlightRecordBatchStream::new(
FlightRecordBatchSource::RecordBatches(recordbatches),
FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches(
recordbatches,
)),
TracingContext::default(),
FlightCompression::default(),
query_ctx,
@@ -663,7 +719,9 @@ mod test {
let query_ctx = query_context_with_matching_capability();
query_ctx.set_explain_verbose(true);
let mut stream = FlightRecordBatchStream::new(
FlightRecordBatchSource::RecordBatches(recordbatches),
FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches(
recordbatches,
)),
TracingContext::default(),
FlightCompression::default(),
query_ctx,
@@ -706,7 +764,9 @@ mod test {
let query_ctx = Arc::new(query_ctx);
query_ctx.set_explain_verbose(true);
let mut stream = FlightRecordBatchStream::new(
FlightRecordBatchSource::RecordBatches(recordbatches),
FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches(
recordbatches,
)),
TracingContext::default(),
FlightCompression::default(),
query_ctx,
@@ -745,7 +805,9 @@ mod test {
});
let query_ctx = query_context_with_matching_capability();
let mut stream = FlightRecordBatchStream::new(
FlightRecordBatchSource::RecordBatches(recordbatches),
FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches(
recordbatches,
)),
TracingContext::default(),
FlightCompression::default(),
query_ctx,

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