From adffac87ce972555f77ca6b9a6881c488cf2f2f4 Mon Sep 17 00:00:00 2001 From: tprevot Date: Fri, 4 Sep 2026 17:48:24 +0200 Subject: [PATCH 01/21] feat: nats source NATS JetStream source binding to a pre-provisioned durable consumer: messages are acknowledged (server-confirmed) once the split containing them is published, and the pipeline drain settles in-flight messages on planned teardowns. Ships config, docs, tutorial, and broker tests. Co-Authored-By: Claude Fable 5 --- .env.example | 1 + .github/workflows/coverage.yml | 3 + LICENSE-3rdparty.csv | 5 + .../tutorials/stackoverflow/nats-source.yaml | 11 + .../stackoverflow/send_messages_to_nats.py | 19 + docker-compose.yml | 16 + docs/configuration/source-config.md | 45 + docs/ingest-data/nats.md | 247 ++++ quickwit/Cargo.lock | 100 ++ quickwit/Cargo.toml | 5 + quickwit/quickwit-cli/Cargo.toml | 3 + quickwit/quickwit-common/src/metrics.rs | 3 + quickwit/quickwit-config/src/lib.rs | 10 +- .../quickwit-config/src/source_config/mod.rs | 196 ++++ .../src/source_config/serialize.rs | 4 +- .../src/indexing_scheduler/mod.rs | 1 + quickwit/quickwit-indexing/Cargo.toml | 6 + quickwit/quickwit-indexing/src/source/mod.rs | 35 +- .../src/source/nats_source.rs | 1017 +++++++++++++++++ .../quickwit-metastore/src/metastore/mod.rs | 1 + 20 files changed, 1717 insertions(+), 11 deletions(-) create mode 100644 config/tutorials/stackoverflow/nats-source.yaml create mode 100644 config/tutorials/stackoverflow/send_messages_to_nats.py create mode 100644 docs/ingest-data/nats.md create mode 100644 quickwit/quickwit-indexing/src/source/nats_source.rs diff --git a/.env.example b/.env.example index 9556b302b34..2cf7a40877a 100644 --- a/.env.example +++ b/.env.example @@ -19,6 +19,7 @@ MAP_HOST_POSTGRES=0.0.0.0 MAP_HOST_PULSAR=0.0.0.0 MAP_HOST_KAFKA=0.0.0.0 MAP_HOST_ZOOKEEPER=0.0.0.0 +MAP_HOST_NATS=0.0.0.0 MAP_HOST_AZURITE=0.0.0.0 MAP_HOST_GRAFANA=0.0.0.0 MAP_HOST_JAEGER=0.0.0.0 diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 5fc1564e680..2fe43c32f47 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -153,6 +153,9 @@ jobs: - name: Run Pulsar service run: DOCKER_SERVICES=pulsar make docker-compose-up + - name: Run NATS service + run: DOCKER_SERVICES=nats make docker-compose-up + - name: Install Rust run: rustup toolchain install working-directory: ./quickwit diff --git a/LICENSE-3rdparty.csv b/LICENSE-3rdparty.csv index b3c307aa803..20339d32348 100644 --- a/LICENSE-3rdparty.csv +++ b/LICENSE-3rdparty.csv @@ -48,6 +48,7 @@ async-channel,https://github.com/smol-rs/async-channel,Apache-2.0 OR MIT,Stjepan async-compression,https://github.com/Nullus157/async-compression,MIT OR Apache-2.0,"Wim Looman , Allen Bui " async-io,https://github.com/smol-rs/async-io,Apache-2.0 OR MIT,Stjepan Glavina async-lock,https://github.com/smol-rs/async-lock,Apache-2.0 OR MIT,Stjepan Glavina +async-nats,https://github.com/nats-io/nats.rs,Apache-2.0,"Tomasz Pietrek , Casper Beyer " async-process,https://github.com/smol-rs/async-process,Apache-2.0 OR MIT,Stjepan Glavina async-recursion,https://github.com/dcchut/async-recursion,MIT OR Apache-2.0,Robert Usher <266585+dcchut@users.noreply.github.com> async-signal,https://github.com/smol-rs/async-signal,Apache-2.0 OR MIT,John Nunley @@ -746,9 +747,11 @@ serde_derive,https://github.com/serde-rs/serde,MIT OR Apache-2.0,"Erick Tryzelaa serde_derive_internals,https://github.com/serde-rs/serde,MIT OR Apache-2.0,"Erick Tryzelaar , David Tolnay " serde_json,https://github.com/serde-rs/json,MIT OR Apache-2.0,"Erick Tryzelaar , David Tolnay " serde_json_borrow,https://github.com/PSeitz/serde_json_borrow,MIT,Pascal Seitz +serde_nanos,https://github.com/caspervonb/serde_nanos,MIT OR Apache-2.0,Casper Beyer serde_path_to_error,https://github.com/dtolnay/path-to-error,MIT OR Apache-2.0,David Tolnay serde_plain,https://github.com/mitsuhiko/serde-plain,MIT OR Apache-2.0,Armin Ronacher serde_qs,https://github.com/samscott89/serde_qs,MIT OR Apache-2.0,Sam Scott +serde_repr,https://github.com/dtolnay/serde-repr,MIT OR Apache-2.0,David Tolnay serde_spanned,https://github.com/toml-rs/toml,MIT OR Apache-2.0,The serde_spanned Authors serde_tokenstream,https://github.com/oxidecomputer/serde_tokenstream,Apache-2.0,The serde_tokenstream Authors serde_urlencoded,https://github.com/nox/serde_urlencoded,MIT OR Apache-2.0,Anthony Ramine @@ -856,6 +859,7 @@ tokio-retry2,https://github.com/naomijub/tokio-retry,MIT,"Julia Naomi tokio-util,https://github.com/tokio-rs/tokio,MIT,Tokio Contributors +tokio-websockets,https://github.com/Gelbpunkt/tokio-websockets,MIT,The tokio-websockets Authors toml,https://github.com/toml-rs/toml,MIT OR Apache-2.0,The toml Authors toml_datetime,https://github.com/toml-rs/toml,MIT OR Apache-2.0,The toml_datetime Authors toml_edit,https://github.com/toml-rs/toml,MIT OR Apache-2.0,The toml_edit Authors @@ -881,6 +885,7 @@ tracing-serde,https://github.com/tokio-rs/tracing,MIT,Tokio Contributors , David Barsky , Tokio Contributors " triomphe,https://github.com/Manishearth/triomphe,MIT OR Apache-2.0,"Manish Goregaokar , The Servo Project Developers" try-lock,https://github.com/seanmonstar/try-lock,MIT,Sean McArthur +tryhard,https://github.com/EmbarkStudios/tryhard,MIT OR Apache-2.0,Embark ttl_cache,https://github.com/stusmall/ttl_cache,MIT OR Apache-2.0,Stu Small twox-hash,https://github.com/shepmaster/twox-hash,MIT,Jake Goulding typeid,https://github.com/dtolnay/typeid,MIT OR Apache-2.0,David Tolnay diff --git a/config/tutorials/stackoverflow/nats-source.yaml b/config/tutorials/stackoverflow/nats-source.yaml new file mode 100644 index 00000000000..fac37d4cbd2 --- /dev/null +++ b/config/tutorials/stackoverflow/nats-source.yaml @@ -0,0 +1,11 @@ +# +# NATS source config file. +# +version: 0.8 +source_id: nats-source +source_type: nats +params: + uris: + - nats://localhost:4222 + stream: stackoverflow + consumer: quickwit-consumer diff --git a/config/tutorials/stackoverflow/send_messages_to_nats.py b/config/tutorials/stackoverflow/send_messages_to_nats.py new file mode 100644 index 00000000000..7cefd4e3f1f --- /dev/null +++ b/config/tutorials/stackoverflow/send_messages_to_nats.py @@ -0,0 +1,19 @@ +import asyncio + +import nats + + +async def main(): + client = await nats.connect("nats://localhost:4222") + jetstream = client.jetstream() + + with open("stackoverflow.posts.transformed-10000.json", encoding="utf8") as file: + for i, line in enumerate(file): + await jetstream.publish("stackoverflow.posts", line.strip().encode("utf-8")) + if i % 1000 == 0: + print(f"{i}/10000 messages sent.") + + await client.close() + + +asyncio.run(main()) diff --git a/docker-compose.yml b/docker-compose.yml index 24f5ed29b47..9198d13aba6 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -127,6 +127,22 @@ services: timeout: 10s retries: 100 + nats: + # 2.10 is the oldest version supporting multiple consumer filter subjects. + image: nats:${NATS_VERSION:-2.10-alpine} + container_name: nats + command: --jetstream --http_port 8222 + ports: + - "${MAP_HOST_NATS:-127.0.0.1}:4222:4222" + profiles: + - all + - nats + healthcheck: + test: ["CMD", "wget", "-q", "-O-", "http://localhost:8222/healthz"] + interval: 1s + timeout: 5s + retries: 100 + azurite: image: mcr.microsoft.com/azure-storage/azurite:${AZURITE_VERSION:-3.24.0} container_name: azurite diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 30894d3faad..41f5ad26787 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -183,6 +183,51 @@ EOF quickwit source create --index my-index --source-config source-config.yaml ``` +### NATS source + +A NATS source reads data from a [NATS JetStream](https://docs.nats.io/nats-concepts/jetstream) stream through a durable consumer. Each message carries one payload in the source's [input format](#input-format): a single JSON object (`json`, the default), a plain text document (`plain_text`), or an OTLP export request whose log records or spans are each indexed as a separate document (`otlp_*` formats). Payloads must not exceed 1 MiB. + +A tutorial is available [here](/docs/ingest-data/nats.md). + +The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The consumer must use the **explicit ack policy**: the source acknowledges each message once the split containing it is published, and the consumer's ack floor is the resume point. + +Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged — with server-confirmed acknowledgments — before the pipeline stops, so nothing is indexed twice; messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. `ack_wait` must exceed the end-to-end publish latency (roughly the commit timeout plus the upload time), otherwise messages are redelivered while they are still being indexed. `max_ack_pending` bounds the messages in flight across all pipelines of the consumer, and therefore the aggregate throughput. + +**Scaling** + +Multiple indexing pipelines can share the consumer: NATS load-balances the messages across them, so scaling up or down is a plain `num_pipelines` update, and the control plane places the pipelines across the indexers of the cluster. + +**Monitoring** + +Being durable, the consumer is observable through NATS's own monitoring (`nats consumer info`, exporters): `num_pending` is the indexing lag and `num_ack_pending` the in-flight window, both available even while the pipelines are down. The source also reports `num_pending_acks` (messages indexed but whose split is not published yet) in its observable state. + +**NATS source parameters** + +| Property | Description | Default value | +| --- | --- | --- | +| `uris` | List of NATS server URIs (e.g. `nats://localhost:4222`). | required | +| `stream` | Name of the JetStream stream to consume. | required | +| `consumer` | Name of the pre-provisioned durable consumer to bind to. | required | +| `tls` | TLS options: `ca_certificates_path` (PEM file whose root certificates are trusted in addition to the system ones), and `client_certificate_path` + `client_key_path` (PEM files, set together) for mutual TLS. TLS itself is enabled by connecting to `tls://` URIs. The files are read by the indexer nodes when the connection is established. | optional | +| `authentication` | Authentication parameters: either `user_password` (with `user` and `password`) or `token`. | optional | + +*Adding a NATS source to an index with the [CLI](../reference/cli.md#source)* + +```bash +cat << EOF > source-config.yaml +version: 0.8 +source_id: my-nats-source +source_type: nats +num_pipelines: 2 +params: + uris: + - nats://localhost:4222 + stream: my-stream + consumer: my-consumer +EOF +./quickwit source create --index my-index --source-config source-config.yaml +``` + ### Pulsar source A Puslar source reads data from one or several Pulsar topics. Each message in topic(s) must hold a JSON object. diff --git a/docs/ingest-data/nats.md b/docs/ingest-data/nats.md new file mode 100644 index 00000000000..2dc3f78cb42 --- /dev/null +++ b/docs/ingest-data/nats.md @@ -0,0 +1,247 @@ +--- +title: NATS +description: A short tutorial describing how to set up Quickwit to ingest data from NATS JetStream in a few minutes +tags: [nats, integration] +icon_url: /img/tutorials/nats.svg +sidebar_position: 4 +--- + +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +In this tutorial, we will describe how to set up Quickwit to ingest data from [NATS JetStream](https://docs.nats.io/nats-concepts/jetstream) in a few minutes. First, we will create an index and configure a NATS source. Then, we will create a JetStream stream and load some events from the [Stack Overflow dataset](https://www.kaggle.com/stackoverflow/stacksample) into it. Finally, we will execute some searches. + +## Prerequisites + +You will need the following to complete this tutorial: +- A local running [Quickwit instance](/docs/get-started/installation.md) +- A local running NATS server (2.10+) with JetStream enabled +- The [NATS CLI](https://github.com/nats-io/natscli) + +### Quickwit setup + +[Download](/docs/get-started/installation.md) Quickwit and start a server. Then open a new terminal to execute CLI commands with the same binary. + +```bash +./quickwit run +``` + +Test that the cluster is running: + +```bash +./quickwit index list +``` + +### NATS setup + + + + + +Download the [NATS server](https://docs.nats.io/running-a-nats-service/introduction/installation) and start it with JetStream enabled: + +```bash +nats-server --jetstream +``` + + + + + +```bash +docker run -it -p 4222:4222 nats:2.10 --jetstream +``` + +See the details on the [official documentation](https://docs.nats.io/running-a-nats-service/nats_docker). + + + + + +## Prepare Quickwit + +First, let's create a new index. Here is the index config and doc mapping corresponding to the schema of Stack Overflow posts: + +```yaml title="index-config.yaml" +# +# Index config file for Stack Overflow dataset. +# +version: 0.7 + +index_id: stackoverflow + +doc_mapping: + field_mappings: + - name: user + type: text + fast: true + tokenizer: raw + - name: tags + type: array + fast: true + tokenizer: raw + - name: type + type: text + fast: true + tokenizer: raw + - name: title + type: text + tokenizer: default + record: position + stored: true + - name: body + type: text + tokenizer: default + record: position + stored: true + - name: questionId + type: u64 + - name: answerId + type: u64 + - name: acceptedAnswerId + type: u64 + - name: creationDate + type: datetime + fast: true + input_formats: + - rfc3339 + fast_precision: seconds + timestamp_field: creationDate + +search_settings: + default_search_fields: [title, body] + +indexing_settings: + commit_timeout_secs: 10 +``` + +Execute these Bash commands to download the index config and create the `stackoverflow` index. + +```bash +# Download stackoverflow index config. +wget -O stackoverflow.yaml https://raw.githubusercontent.com/quickwit-oss/quickwit/main/config/tutorials/stackoverflow/index-config.yaml + +# Create index. +./quickwit index create --index-config stackoverflow.yaml +``` + +## Create a JetStream stream and a durable consumer + +The NATS source consumes a JetStream stream through a durable consumer, so messages must be published on subjects captured by a stream, and the consumer must be provisioned before creating the source — Quickwit only ever fetches it, and never creates, updates, nor deletes it. Let's create both: + +```bash +nats stream add stackoverflow --subjects "stackoverflow.posts" --defaults +nats consumer add stackoverflow quickwit-consumer --pull --deliver all --ack explicit --wait 5m --defaults +``` + +:::info + +The consumer must use the explicit ack policy: the source acknowledges each message once it is durably indexed, and the consumer's ack floor is the resume point. The subject filters, the deliver policy, and the ack tuning are properties of the consumer, not of the Quickwit source; the ack wait must exceed the indexing commit timeout. See the [NATS source reference](../configuration/source-config.md#nats-source) for details. + +::: + +## Create the NATS source + +A NATS source just needs to define the server URIs, the stream, and the durable consumer to bind to. + +```yaml title="nats-source.yaml" +# +# NATS source config file. +# +version: 0.8 +source_id: nats-source +source_type: nats +params: + uris: + - nats://localhost:4222 + stream: stackoverflow + consumer: quickwit-consumer +``` + +Run these commands to download the source config file and create the source. + +```bash +# Download NATS source config. +wget -O stackoverflow-nats-source.yaml https://raw.githubusercontent.com/quickwit-oss/quickwit/main/config/tutorials/stackoverflow/nats-source.yaml + +# Create source. +./quickwit source create --index stackoverflow --source-config stackoverflow-nats-source.yaml +``` + +As soon as the NATS source is created, the Quickwit control plane will ask an indexer to start a new indexing pipeline. You will see logs like below on the indexer: + +```bash +INFO spawn_pipeline{index=stackoverflow gen=0}: quickwit_indexing::source::nats_source: starting NATS source index_id=stackoverflow source_id=nats-source stream=stackoverflow consumer_name=quickwit-consumer +``` + +The consumer's indexing lag and ack floor are visible at any time — even while the pipeline is down — with `nats consumer info stackoverflow quickwit-consumer`. + +## Populate the stream + +To populate the stream, we will use a python script: + +```python title=send_messages_to_nats.py +import asyncio + +import nats + + +async def main(): + client = await nats.connect("nats://localhost:4222") + jetstream = client.jetstream() + + with open("stackoverflow.posts.transformed-10000.json", encoding="utf8") as file: + for i, line in enumerate(file): + await jetstream.publish("stackoverflow.posts", line.strip().encode("utf-8")) + if i % 1000 == 0: + print(f"{i}/10000 messages sent.") + + await client.close() + + +asyncio.run(main()) +``` + +Install the [python client](https://github.com/nats-io/nats.py) locally and run the script: + +```bash +# Download the first 10_000 Stackoverflow posts articles. +curl -O https://quickwit-datasets-public.s3.amazonaws.com/stackoverflow.posts.transformed-10000.json + +# Install nats python client. +pip3 install nats-py +wget https://raw.githubusercontent.com/quickwit-oss/quickwit/main/config/tutorials/stackoverflow/send_messages_to_nats.py +python3 send_messages_to_nats.py +``` + +## Time to search! + +You can run this command to inspect the properties of the index and check the current number of published splits and documents: + +```bash +# Display some general information about the index. +./quickwit index describe --index stackoverflow +``` + +You will notably see the number of published documents. + +You are now ready to execute some queries. + +```bash +curl 'http://localhost:7280/api/v1/stackoverflow/search?query=search+AND+engine' +``` + +If your Quickwit server is local, you can access the results through the Quickwit UI on [localhost:7280](http://localhost:7280/ui/search?query=&index_id=stackoverflow&max_hits=10). + +## Tear down resources (optional) + +Let's delete the files and resources created for the purpose of this tutorial. + +```bash +# Delete quickwit index. +./quickwit index delete --index stackoverflow --yes +# Delete NATS stream. +nats stream rm -f stackoverflow +``` + +This concludes the tutorial. If you have any questions regarding Quickwit or encounter any issues, don't hesitate to ask a [question](https://github.com/quickwit-oss/quickwit/discussions) or open an [issue](https://github.com/quickwit-oss/quickwit/issues) on [GitHub](https://github.com/quickwit-oss/quickwit) or contact us directly on [Discord](https://discord.com/invite/MT27AG5EVE). diff --git a/quickwit/Cargo.lock b/quickwit/Cargo.lock index caede75b90f..18957554214 100644 --- a/quickwit/Cargo.lock +++ b/quickwit/Cargo.lock @@ -582,6 +582,40 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "async-nats" +version = "0.50.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d83a251fa1a4c9d0fe6e816b7acd60549e473e08d14f27a1d992c2675abff05f" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-util", + "memchr", + "pin-project", + "portable-atomic", + "rand 0.10.2", + "regex", + "ring", + "rustls-native-certs", + "rustls-pki-types", + "rustls-webpki 0.103.13", + "serde", + "serde_json", + "serde_nanos", + "serde_repr", + "thiserror 2.0.18", + "time", + "tokio", + "tokio-rustls 0.26.4", + "tokio-stream", + "tokio-util", + "tokio-websockets", + "tracing", + "tryhard", + "url", +] + [[package]] name = "async-process" version = "2.5.0" @@ -8777,6 +8811,7 @@ dependencies = [ "arc-swap", "arrow", "async-compression", + "async-nats", "async-trait", "aws-sdk-kinesis", "aws-sdk-sqs", @@ -8796,6 +8831,8 @@ dependencies = [ "mockall", "oneshot 0.2.1", "openssl", + "opentelemetry", + "opentelemetry_sdk", "parquet", "percent-encoding", "proptest", @@ -8832,6 +8869,7 @@ dependencies = [ "time", "tokio", "tracing", + "tracing-opentelemetry", "ulid", "unicode-segmentation", "utoipa", @@ -10784,6 +10822,15 @@ dependencies = [ "serde_json", ] +[[package]] +name = "serde_nanos" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a93142f0367a4cc53ae0fead1bcda39e85beccfad3dcd717656cacab94b12985" +dependencies = [ + "serde", +] + [[package]] name = "serde_path_to_error" version = "0.1.20" @@ -10827,6 +10874,17 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_repr" +version = "0.1.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d3b1629de253c70a0508c3899572da79ca359fdab27c7920ff00406df418906" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.4", +] + [[package]] name = "serde_spanned" version = "1.1.1" @@ -11688,6 +11746,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6275cddf4610d1775e6d1fe9469b2e77d0f39fd98fb7450901b821e0c53649f" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "sync_wrapper" version = "1.0.2" @@ -12257,6 +12326,27 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-websockets" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f591660438b3038dd04d16c938271c79e7e06260ad2ea2885a4861bfb238605d" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-core", + "futures-sink", + "http 1.4.2", + "httparse", + "rand 0.8.6", + "ring", + "rustls-pki-types", + "tokio", + "tokio-rustls 0.26.4", + "tokio-util", + "webpki-roots 0.26.11", +] + [[package]] name = "toml" version = "1.1.2+spec-1.1.0" @@ -12626,6 +12716,16 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "tryhard" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fe58ebd5edd976e0fe0f8a14d2a04b7c81ef153ea9a54eebc42e67c2c23b4e5" +dependencies = [ + "pin-project-lite", + "tokio", +] + [[package]] name = "ttl_cache" version = "0.5.1" diff --git a/quickwit/Cargo.toml b/quickwit/Cargo.toml index 8904ed1dd3c..84cbbd71c88 100644 --- a/quickwit/Cargo.toml +++ b/quickwit/Cargo.toml @@ -104,6 +104,11 @@ arc-swap = "1.9" arrow = { version = "58", default-features = false, features = ["ipc"] } assert-json-diff = "2" async-compression = { version = "0.4", features = ["tokio", "gzip"] } +async-nats = { version = "0.50", default-features = false, features = [ + "jetstream", + "ring", + "server_2_10", +] } async-speed-limit = "0.4" async-trait = "0.1" atomic_float = "1.1" diff --git a/quickwit/quickwit-cli/Cargo.toml b/quickwit/quickwit-cli/Cargo.toml index 01857eb221f..6e115ac4023 100644 --- a/quickwit/quickwit-cli/Cargo.toml +++ b/quickwit/quickwit-cli/Cargo.toml @@ -108,6 +108,7 @@ release-feature-set = [ "pprof", "quickwit-indexing/kafka", "quickwit-indexing/kinesis", + "quickwit-indexing/nats", "quickwit-indexing/pulsar", "quickwit-indexing/sqs", "quickwit-indexing/vrl", @@ -122,6 +123,7 @@ release-feature-vendored-set = [ "pprof", "quickwit-indexing/kinesis", "quickwit-indexing/metrics", + "quickwit-indexing/nats", "quickwit-indexing/pulsar", "quickwit-indexing/sqs", "quickwit-indexing/vrl", @@ -136,6 +138,7 @@ release-macos-feature-vendored-set = [ "openssl-support", "quickwit-indexing/kinesis", "quickwit-indexing/metrics", + "quickwit-indexing/nats", "quickwit-indexing/pulsar", "quickwit-indexing/sqs", "quickwit-indexing/vrl", diff --git a/quickwit/quickwit-common/src/metrics.rs b/quickwit/quickwit-common/src/metrics.rs index 7244c3816e4..aa8f6e97adf 100644 --- a/quickwit/quickwit-common/src/metrics.rs +++ b/quickwit/quickwit-common/src/metrics.rs @@ -88,6 +88,9 @@ pub static IN_FLIGHT_KAFKA_SOURCE: LazyGauge = pub static IN_FLIGHT_KINESIS_SOURCE: LazyGauge = LazyGauge::new(|| in_flight_data_gauge("kinesis_source")); +pub static IN_FLIGHT_NATS_SOURCE: LazyGauge = + LazyGauge::new(|| in_flight_data_gauge("nats_source")); + pub static IN_FLIGHT_PUBSUB_SOURCE: LazyGauge = LazyGauge::new(|| in_flight_data_gauge("pubsub_source")); diff --git a/quickwit/quickwit-config/src/lib.rs b/quickwit/quickwit-config/src/lib.rs index 904e65f8180..1a06d42d2b9 100644 --- a/quickwit/quickwit-config/src/lib.rs +++ b/quickwit/quickwit-config/src/lib.rs @@ -61,9 +61,10 @@ use source_config::FileSourceParamsForSerde; pub use source_config::{ CLI_SOURCE_ID, FileSourceMessageType, FileSourceNotification, FileSourceParams, FileSourceSqs, INGEST_API_SOURCE_ID, INGEST_V2_SOURCE_ID, KafkaSourceParams, KinesisSourceParams, - PubSubSourceParams, PulsarSourceAuth, PulsarSourceParams, RegionOrEndpoint, SourceConfig, - SourceInputFormat, SourceParams, TransformConfig, VecSourceParams, VoidSourceParams, - load_source_config_from_user_config, load_source_config_update, + NatsSourceAuth, NatsSourceParams, NatsSourceTls, PubSubSourceParams, PulsarSourceAuth, + PulsarSourceParams, RegionOrEndpoint, SourceConfig, SourceInputFormat, SourceParams, + TransformConfig, VecSourceParams, VoidSourceParams, load_source_config_from_user_config, + load_source_config_update, }; use tracing::warn; @@ -116,6 +117,9 @@ pub fn disable_ingest_v1() -> bool { KafkaSourceParams, KinesisSourceParams, MergePolicyConfig, + NatsSourceAuth, + NatsSourceParams, + NatsSourceTls, ParquetIndexingConfig, ParquetMergePolicyConfig, PubSubSourceParams, diff --git a/quickwit/quickwit-config/src/source_config/mod.rs b/quickwit/quickwit-config/src/source_config/mod.rs index 81811e9f47c..9283f2a1dd5 100644 --- a/quickwit/quickwit-config/src/source_config/mod.rs +++ b/quickwit/quickwit-config/src/source_config/mod.rs @@ -86,6 +86,7 @@ impl SourceConfig { SourceParams::Kafka(params) => serde_json::to_value(params), SourceParams::Kinesis(params) => serde_json::to_value(params), SourceParams::Pulsar(params) => serde_json::to_value(params), + SourceParams::Nats(params) => serde_json::to_value(params), SourceParams::Stdin => serde_json::to_value(()), SourceParams::Vec(params) => serde_json::to_value(params), SourceParams::Void(params) => serde_json::to_value(params), @@ -231,6 +232,7 @@ pub enum SourceParams { #[serde(rename = "pubsub")] PubSub(PubSubSourceParams), Pulsar(PulsarSourceParams), + Nats(NatsSourceParams), Stdin, Vec(VecSourceParams), Void(VoidSourceParams), @@ -263,6 +265,7 @@ impl SourceParams { SourceParams::Kinesis(_) => SourceType::Kinesis, SourceParams::PubSub(_) => SourceType::PubSub, SourceParams::Pulsar(_) => SourceType::Pulsar, + SourceParams::Nats(_) => SourceType::Nats, SourceParams::Stdin => SourceType::Stdin, SourceParams::Vec(_) => SourceType::Vec, SourceParams::Void(_) => SourceType::Void, @@ -287,6 +290,7 @@ impl SourceParams { (SourceParams::Pulsar(current), SourceParams::Pulsar(new)) => { current.validate_update(new) } + (SourceParams::Nats(current), SourceParams::Nats(new)) => current.validate_update(new), (current, new) if current.source_type() != new.source_type() => Err(anyhow::anyhow!( "source type cannot be changed, current type {}", current.source_type(), @@ -620,6 +624,88 @@ fn default_consumer_name() -> String { "quickwit".to_string() } +#[derive( + Clone, Debug, Eq, PartialEq, Hash, serde::Serialize, serde::Deserialize, utoipa::ToSchema, +)] +#[serde(deny_unknown_fields)] +pub struct NatsSourceParams { + /// The URIs of the NATS servers to connect to (e.g. `nats://localhost:4222`). + pub uris: Vec, + /// Name of the NATS JetStream stream that the source consumes. + pub stream: String, + /// Name of the durable consumer to bind to. The consumer is provisioned + /// externally. It must use the explicit ack policy: the source acknowledges + /// each message once the split containing it is published. + pub consumer: String, + /// TLS options: custom root CA and client certificates (mutual TLS). TLS + /// itself is enabled by connecting to `tls://` URIs. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tls: Option, + // Serde yaml has some specific behaviour when deserializing + // enums (see https://github.com/dtolnay/serde-yaml/issues/342) + // and requires explicitly stating `default` in order to make the parameter + // optional on the yaml config. + #[serde(default, with = "serde_yaml::with::singleton_map")] + /// Authentication for NATS. + pub authentication: Option, +} + +impl NatsSourceParams { + pub fn validate(&self) -> anyhow::Result<()> { + ensure!( + !self.uris.is_empty(), + "`uris` must contain at least one NATS server URI" + ); + ensure!(!self.stream.is_empty(), "`stream` must be set"); + ensure!(!self.consumer.is_empty(), "`consumer` must be set"); + if let Some(tls) = &self.tls { + ensure!( + tls.client_certificate_path.is_some() == tls.client_key_path.is_some(), + "`tls.client_certificate_path` and `tls.client_key_path` must be set together" + ); + } + Ok(()) + } + + fn validate_update(&self, _other: &Self) -> anyhow::Result<()> { + // The resume point is the consumer's ack floor, not the metastore + // checkpoint: re-binding to another stream or consumer is a plain + // reconfiguration and cannot corrupt any tracked progress. + Ok(()) + } +} + +/// TLS options for the NATS connection. Certificates are PEM files loaded by +/// the indexer when the connection is established, so the paths must exist on +/// the indexer nodes; validation only checks the options' consistency. +#[derive( + Clone, Debug, Eq, PartialEq, Hash, serde::Serialize, serde::Deserialize, utoipa::ToSchema, +)] +#[serde(deny_unknown_fields)] +pub struct NatsSourceTls { + /// Path to a PEM file whose root certificates are trusted in addition to + /// the system ones (e.g. a private CA). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ca_certificates_path: Option, + /// Path to the client certificate PEM file, for mutual TLS. Requires + /// `client_key_path`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub client_certificate_path: Option, + /// Path to the client private key PEM file, for mutual TLS. Requires + /// `client_certificate_path`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub client_key_path: Option, +} + +#[derive( + Clone, Debug, Eq, PartialEq, Hash, serde::Serialize, serde::Deserialize, utoipa::ToSchema, +)] +#[serde(rename_all = "snake_case")] +pub enum NatsSourceAuth { + UserPassword { user: String, password: String }, + Token(String), +} + #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, utoipa::ToSchema)] #[serde(deny_unknown_fields)] pub struct TransformConfig { @@ -1341,6 +1427,116 @@ mod tests { } } + #[test] + fn test_nats_source_params_deserialization() { + let base_params = NatsSourceParams { + uris: vec!["nats://localhost:4222".to_string()], + stream: "my-stream".to_string(), + consumer: "my-consumer".to_string(), + tls: None, + authentication: None, + }; + { + let yaml = r#" + uris: + - nats://localhost:4222 + stream: my-stream + consumer: my-consumer + "#; + let params = serde_yaml::from_str::(yaml).unwrap(); + assert_eq!(params, base_params); + params.validate().unwrap(); + } + { + let yaml = r#" + uris: + - tls://localhost:4222 + stream: my-stream + consumer: my-consumer + tls: + ca_certificates_path: /etc/ssl/private-ca.pem + client_certificate_path: /etc/ssl/client.pem + client_key_path: /etc/ssl/client.key + authentication: + token: my-token + "#; + let params = serde_yaml::from_str::(yaml).unwrap(); + assert_eq!( + params, + NatsSourceParams { + uris: vec!["tls://localhost:4222".to_string()], + tls: Some(NatsSourceTls { + ca_certificates_path: Some("/etc/ssl/private-ca.pem".to_string()), + client_certificate_path: Some("/etc/ssl/client.pem".to_string()), + client_key_path: Some("/etc/ssl/client.key".to_string()), + }), + authentication: Some(NatsSourceAuth::Token("my-token".to_string())), + ..base_params.clone() + } + ); + params.validate().unwrap(); + } + { + let yaml = r#" + uris: + - nats://localhost:4222 + stream: my-stream + "#; + serde_yaml::from_str::(yaml) + .expect_err("parameters should error on missing consumer"); + } + { + let params = NatsSourceParams { + uris: Vec::new(), + ..base_params.clone() + }; + params + .validate() + .expect_err("validation should reject empty uris"); + } + { + let params = NatsSourceParams { + consumer: String::new(), + ..base_params.clone() + }; + params + .validate() + .expect_err("validation should reject an empty consumer name"); + } + { + let params = NatsSourceParams { + tls: Some(NatsSourceTls { + ca_certificates_path: None, + client_certificate_path: Some("/etc/ssl/client.pem".to_string()), + client_key_path: None, + }), + ..base_params.clone() + }; + params + .validate() + .expect_err("validation should reject a client certificate without its key"); + } + } + + #[test] + fn test_nats_source_supports_multiple_pipelines() { + let source_config_yaml = r#" + version: 0.8 + source_id: my-nats-source + num_pipelines: 2 + source_type: nats + params: + uris: + - nats://localhost:4222 + stream: my-stream + consumer: my-consumer + "#; + let source_config = + load_source_config_from_user_config(ConfigFormat::Yaml, source_config_yaml.as_bytes()) + .unwrap(); + assert_eq!(source_config.num_pipelines.get(), 2); + } + #[cfg(feature = "vrl")] #[tokio::test] async fn test_load_ingest_api_source_config() { diff --git a/quickwit/quickwit-config/src/source_config/serialize.rs b/quickwit/quickwit-config/src/source_config/serialize.rs index 3f580ee8fa3..25e92fe0a28 100644 --- a/quickwit/quickwit-config/src/source_config/serialize.rs +++ b/quickwit/quickwit-config/src/source_config/serialize.rs @@ -116,6 +116,7 @@ impl SourceConfigForSerialization { | SourceParams::Pulsar(_) => { // TODO consider any validation opportunity } + SourceParams::Nats(params) => params.validate()?, SourceParams::PubSub(_) | SourceParams::Ingest | SourceParams::IngestApi @@ -127,10 +128,11 @@ impl SourceConfigForSerialization { SourceParams::PubSub(_) | SourceParams::Kafka(_) | SourceParams::File(FileSourceParams::Notifications(_)) => {} + SourceParams::Nats(_) => {} _ => { if self.num_pipelines > 1 { bail!( - "Quickwit currently supports multiple pipelines only for GCP PubSub or Kafka sources. open an issue https://github.com/quickwit-oss/quickwit/issues if you need the feature for other source types" + "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or NATS sources. open an issue https://github.com/quickwit-oss/quickwit/issues if you need the feature for other source types" ); } } diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs index a656cdcfc5f..40ffe42145a 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs @@ -272,6 +272,7 @@ fn get_sources_to_schedule( } SourceParams::Kafka(_) | SourceParams::Kinesis(_) + | SourceParams::Nats(_) | SourceParams::PubSub(_) | SourceParams::Pulsar(_) | SourceParams::File(FileSourceParams::Notifications(_)) => { diff --git a/quickwit/quickwit-indexing/Cargo.toml b/quickwit/quickwit-indexing/Cargo.toml index 23d81a3e58c..172dfbf0133 100644 --- a/quickwit/quickwit-indexing/Cargo.toml +++ b/quickwit/quickwit-indexing/Cargo.toml @@ -14,6 +14,7 @@ license.workspace = true anyhow = { workspace = true } arc-swap = { workspace = true } async-compression = { workspace = true } +async-nats = { workspace = true, optional = true } async-trait = { workspace = true } aws-sdk-kinesis = { workspace = true, optional = true } aws-sdk-sqs = { workspace = true, optional = true } @@ -31,6 +32,7 @@ itertools = { workspace = true } libz-sys = { workspace = true, optional = true } oneshot = { workspace = true } openssl = { workspace = true, optional = true } +opentelemetry = { workspace = true, optional = true } percent-encoding = { workspace = true } pulsar = { workspace = true, optional = true } quickwit-query = { workspace = true } @@ -45,6 +47,7 @@ thiserror = { workspace = true } time = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } +tracing-opentelemetry = { workspace = true, optional = true } ulid = { workspace = true } utoipa = { workspace = true } vrl = { workspace = true, optional = true } @@ -83,6 +86,8 @@ kinesis = [ "quickwit-aws/kinesis", ] kinesis-localstack-tests = [] +nats = ["dep:async-nats", "dep:opentelemetry", "dep:tracing-opentelemetry"] +nats-broker-tests = ["nats"] pulsar = ["dep:pulsar"] pulsar-broker-tests = [] queue-sources = [] @@ -117,6 +122,7 @@ ci-test = [] bytes = { workspace = true } criterion = { workspace = true, features = ["async_tokio"] } mockall = { workspace = true } +opentelemetry_sdk = { workspace = true } parquet = { workspace = true } proptest = { workspace = true } prost = { workspace = true } diff --git a/quickwit/quickwit-indexing/src/source/mod.rs b/quickwit/quickwit-indexing/src/source/mod.rs index e23f44dc732..d86bf764dec 100644 --- a/quickwit/quickwit-indexing/src/source/mod.rs +++ b/quickwit/quickwit-indexing/src/source/mod.rs @@ -62,6 +62,8 @@ mod ingest_api_source; mod kafka_source; #[cfg(feature = "kinesis")] mod kinesis; +#[cfg(feature = "nats")] +mod nats_source; #[cfg(feature = "pulsar")] mod pulsar_source; #[cfg(feature = "queue-sources")] @@ -87,6 +89,8 @@ pub use gcp_pubsub_source::{GcpPubSubSource, GcpPubSubSourceFactory}; pub use kafka_source::{KafkaSource, KafkaSourceFactory}; #[cfg(feature = "kinesis")] pub use kinesis::kinesis_source::{KinesisSource, KinesisSourceFactory}; +#[cfg(feature = "nats")] +pub use nats_source::{NatsSource, NatsSourceFactory}; #[cfg(feature = "pulsar")] pub use pulsar_source::{PulsarSource, PulsarSourceFactory}; #[cfg(feature = "sqs")] @@ -94,8 +98,8 @@ pub use queue_sources::sqs_queue; use quickwit_actors::{Actor, ActorContext, ActorExitStatus, Handler}; use quickwit_common::metrics::{ IN_FLIGHT_FILE_SOURCE, IN_FLIGHT_INGEST_SOURCE, IN_FLIGHT_KAFKA_SOURCE, - IN_FLIGHT_KINESIS_SOURCE, IN_FLIGHT_OTHER_SOURCE, IN_FLIGHT_PUBSUB_SOURCE, - IN_FLIGHT_PULSAR_SOURCE, + IN_FLIGHT_KINESIS_SOURCE, IN_FLIGHT_NATS_SOURCE, IN_FLIGHT_OTHER_SOURCE, + IN_FLIGHT_PUBSUB_SOURCE, IN_FLIGHT_PULSAR_SOURCE, }; use quickwit_common::pubsub::EventBroker; use quickwit_common::runtimes::RuntimeType; @@ -434,10 +438,8 @@ impl Handler for SourceActor { /// messages, which are then redelivered and indexed again after the broker's /// ack timeout. Checkpoint-based sources resume from their last checkpoint, so /// a drain buys them nothing. -pub(crate) fn source_needs_drain(_source_params: &SourceParams) -> bool { - // No built-in source opts in yet: acknowledgment-based sources add their - // variant here as they land. - false +pub(crate) fn source_needs_drain(source_params: &SourceParams) -> bool { + matches!(source_params, SourceParams::Nats(_)) } #[async_trait] @@ -489,6 +491,8 @@ pub fn quickwit_supported_sources() -> &'static SourceLoader { source_factory.add_source(SourceType::Kafka, KafkaSourceFactory); #[cfg(feature = "kinesis")] source_factory.add_source(SourceType::Kinesis, KinesisSourceFactory); + #[cfg(feature = "nats")] + source_factory.add_source(SourceType::Nats, NatsSourceFactory); #[cfg(feature = "pulsar")] source_factory.add_source(SourceType::Pulsar, PulsarSourceFactory); source_factory.add_source(SourceType::Stdin, StdinSourceFactory); @@ -556,6 +560,17 @@ pub async fn check_source_connectivity( Ok(()) } } + #[allow(unused_variables)] + SourceParams::Nats(params) => { + #[cfg(not(feature = "nats"))] + anyhow::bail!("Quickwit was compiled without the `nats` feature"); + + #[cfg(feature = "nats")] + { + nats_source::check_connectivity(params).await?; + Ok(()) + } + } _ => Ok(()), } } @@ -608,6 +623,7 @@ impl BatchBuilder { SourceType::IngestV2 => &IN_FLIGHT_INGEST_SOURCE, SourceType::Kafka => &IN_FLIGHT_KAFKA_SOURCE, SourceType::Kinesis => &IN_FLIGHT_KINESIS_SOURCE, + SourceType::Nats => &IN_FLIGHT_NATS_SOURCE, SourceType::PubSub => &IN_FLIGHT_PUBSUB_SOURCE, SourceType::Pulsar => &IN_FLIGHT_PULSAR_SOURCE, _ => &IN_FLIGHT_OTHER_SOURCE, @@ -707,7 +723,11 @@ mod tests { #[cfg(all( test, - any(feature = "kafka-broker-tests", feature = "sqs-localstack-tests") + any( + feature = "kafka-broker-tests", + feature = "nats-broker-tests", + feature = "sqs-localstack-tests" + ) ))] pub fn with_metastore(mut self, metastore: MetastoreServiceClient) -> Self { self.metastore_opt = Some(metastore); @@ -865,6 +885,7 @@ mod tests { any( feature = "sqs-localstack-tests", feature = "kafka-broker-tests", + feature = "nats-broker-tests", feature = "pulsar-broker-tests" ) ))] diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs new file mode 100644 index 00000000000..aeb0acc66fd --- /dev/null +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -0,0 +1,1017 @@ +// Copyright 2021-Present Datadog, Inc. +// +// 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. + +//! A source consuming a NATS JetStream stream through a pre-provisioned +//! durable consumer. +//! +//! The consumer is only ever fetched — never created, updated, nor deleted — +//! so its lifecycle, subject filters, deliver policy, and ack tuning +//! (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The +//! consumer must use the explicit ack policy: each message is acknowledged — +//! with a server-confirmed ack — once the split containing it is published. +//! On a planned teardown (node shutdown, pipeline reassignment) the pipeline +//! is drained first: the in-flight messages are published and acknowledged, +//! and the prefetched but unprocessed ones are negatively acknowledged for +//! immediate redelivery — so delivery is exactly-once. A crash degrades it to +//! at-least-once: the unacknowledged messages — including ones whose split +//! ends up published — are redelivered after the consumer's `ack_wait` and +//! indexed again. The consumer's ack floor is the resume point; the metastore +//! checkpoint only carries synthetic positions used to correlate published +//! batches with the acknowledgments they release. +//! +//! Several indexing pipelines can share the consumer: NATS load-balances the +//! messages across them, so scaling is a plain `num_pipelines` update. Being +//! durable, the consumer is observable through NATS's own monitoring even +//! while the pipelines are down. +//! +//! When a message carries a W3C `traceparent` header, the source stitches the +//! processing of the message into the publisher's distributed trace. + +use std::collections::BTreeMap; +use std::fmt; +use std::path::PathBuf; +use std::str::FromStr; +use std::time::{Duration, Instant}; + +use anyhow::{Context as _, anyhow, bail, ensure}; +use async_nats::header::{HeaderMap, HeaderName}; +use async_nats::jetstream::consumer::pull::{ + MessagesError, MessagesErrorKind, Stream as DurableMessageStream, +}; +use async_nats::jetstream::consumer::{AckPolicy, PullConsumer}; +use async_nats::jetstream::message::AckKind; +use async_nats::{ConnectOptions, Subject, jetstream}; +use async_trait::async_trait; +use futures::{FutureExt, StreamExt}; +use opentelemetry::propagation::Extractor; +use opentelemetry::trace::TraceContextExt; +use opentelemetry::{Context as OtelContext, global}; +use quickwit_actors::ActorExitStatus; +use quickwit_common::rand::append_random_suffix; +use quickwit_config::{NatsSourceAuth, NatsSourceParams}; +use quickwit_metastore::checkpoint::{PartitionId, SourceCheckpoint}; +use quickwit_proto::metastore::SourceType; +use quickwit_proto::types::Position; +use serde_json::{Value as JsonValue, json}; +use tokio::time; +use tracing::{Span, debug, info, warn}; +use tracing_opentelemetry::OpenTelemetrySpanExt; + +use crate::source::{ + BATCH_NUM_BYTES_LIMIT, BatchBuilder, EMIT_BATCHES_TIMEOUT, Source, SourceContext, + SourceRuntime, SourceSink, TypedSourceFactory, +}; + +pub struct NatsSourceFactory; + +#[async_trait] +impl TypedSourceFactory for NatsSourceFactory { + type Source = NatsSource; + type Params = NatsSourceParams; + + async fn typed_create_source( + source_runtime: SourceRuntime, + source_params: NatsSourceParams, + ) -> anyhow::Result { + NatsSource::try_new(source_runtime, source_params).await + } +} + +#[derive(Default, Debug)] +pub struct NatsSourceState { + /// Number of bytes processed by the source. + pub num_bytes_processed: u64, + /// Number of messages processed by the source (including invalid messages). + pub num_messages_processed: u64, + /// Number of invalid messages, i.e., that were empty. + pub num_invalid_messages: u64, +} + +/// Source flavor binding to a pre-provisioned durable consumer. +/// +/// The consumer is only ever fetched — never created, updated, nor deleted — +/// so its lifecycle, subject filters, deliver policy, and ack tuning +/// (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. Each +/// message is acknowledged once the split containing it is published +/// ([`Source::suggest_truncate`]); see the module documentation for the +/// delivery semantics. +/// +/// The metastore checkpoint only carries synthetic positions (a per-pipeline +/// partition and a delivery counter) used to correlate published batches with +/// the acknowledgments they release; the resume point is the consumer's ack +/// floor, not the checkpoint. +/// +/// The source acknowledges inline in [`Source::suggest_truncate`], so once a +/// drain reports completion (see [`Source::is_drained`]) every ack has been +/// confirmed by the server. If the source actor is instead killed with batches +/// still in the pipeline, their `suggest_truncate` notifications are lost and +/// the unacknowledged messages — including ones whose split ends up published — +/// are redelivered after the consumer's `ack_wait`. +pub struct NatsSource { + source_runtime: SourceRuntime, + source_params: NatsSourceParams, + nats_client: async_nats::Client, + message_stream: DurableMessageStream, + consumer_name: String, + partition_id: PartitionId, + delivery_counter: u64, + /// Ack subjects of the messages delivered but not published yet, keyed by + /// delivery counter. Bounded by the consumer's `max_ack_pending`: the + /// server stops delivering when too many messages are unacknowledged. + pending_acks: BTreeMap, + state: NatsSourceState, +} + +impl fmt::Debug for NatsSource { + fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter + .debug_struct("NatsSource") + .field("index_uid", self.source_runtime.index_uid()) + .field("source_id", &self.source_runtime.source_id()) + .field("stream", &self.source_params.stream) + .field("consumer_name", &self.consumer_name) + .finish() + } +} + +impl NatsSource { + pub async fn try_new( + source_runtime: SourceRuntime, + source_params: NatsSourceParams, + ) -> anyhow::Result { + let consumer_name = source_params.consumer.clone(); + + info!( + index_id=%source_runtime.index_id(), + source_id=%source_runtime.source_id(), + stream=%source_params.stream, + %consumer_name, + "starting NATS source" + ); + + let (nats_client, consumer) = connect_and_fetch_consumer(&source_params).await?; + let message_stream = consumer + .messages() + .await + .context("failed to subscribe to NATS consumer messages")?; + + // A fresh partition per pipeline: durable-mode positions are delivery + // counters local to this instance, and pipelines sharing the consumer + // must not collide on a partition. + let partition_id = + PartitionId::from(append_random_suffix(&format!("nats-{consumer_name}"))); + + Ok(NatsSource { + source_runtime, + source_params, + nats_client, + message_stream, + consumer_name, + partition_id, + delivery_counter: 0, + pending_acks: BTreeMap::new(), + state: NatsSourceState::default(), + }) + } + + fn process_message( + &mut self, + message: jetstream::Message, + batch: &mut BatchBuilder, + ) -> anyhow::Result<()> { + let stream_sequence = message + .info() + .map_err(|error| anyhow!("failed to parse NATS message metadata: {error}"))? + .stream_sequence; + let _span_guard = remote_parented_span(&message, stream_sequence, &self.source_runtime) + .map(Span::entered); + let Some(ack_subject) = message.message.reply else { + bail!("NATS message carries no reply subject to acknowledge it on"); + }; + let doc = message.message.payload; + let num_bytes = doc.len() as u64; + + if doc.is_empty() { + warn!("message received from NATS was empty"); + self.state.num_invalid_messages += 1; + } else { + batch.add_doc(doc); + } + let from_position = if self.delivery_counter == 0 { + Position::Beginning + } else { + Position::offset(self.delivery_counter) + }; + self.delivery_counter += 1; + batch + .checkpoint_delta + .record_partition_delta( + self.partition_id.clone(), + from_position, + Position::offset(self.delivery_counter), + ) + .context("failed to record partition delta")?; + self.pending_acks.insert(self.delivery_counter, ack_subject); + + self.state.num_bytes_processed += num_bytes; + self.state.num_messages_processed += 1; + + Ok(()) + } +} + +#[async_trait] +impl Source for NatsSource { + async fn emit_batches( + &mut self, + source_sink: &SourceSink, + ctx: &SourceContext, + ) -> Result { + let now = Instant::now(); + let mut batch_builder = BatchBuilder::new(SourceType::Nats); + let deadline = time::sleep(*EMIT_BATCHES_TIMEOUT); + tokio::pin!(deadline); + let mut wait_before_next_batch = Duration::default(); + + loop { + tokio::select! { + message_res_opt = self.message_stream.next() => { + let message_res = message_res_opt + .ok_or_else(|| ActorExitStatus::from(anyhow!("NATS message stream ended unexpectedly")))?; + let message = match message_res { + Ok(message) => message, + Err(error) if is_transient_stream_error(&error) => { + warn!(%error, "transient NATS message stream error"); + wait_before_next_batch = TRANSIENT_STREAM_ERROR_BACKOFF; + break; + } + Err(error) => { + return Err(ActorExitStatus::from(anyhow!( + "failed to pull message from NATS consumer: {error}" + ))); + } + }; + self.process_message(message, &mut batch_builder).map_err(ActorExitStatus::from)?; + + if batch_builder.num_bytes >= BATCH_NUM_BYTES_LIMIT { + break; + } + } + _ = &mut deadline => { + break; + } + } + ctx.record_progress(); + } + + if !batch_builder.checkpoint_delta.is_empty() { + debug!( + num_docs=%batch_builder.docs.len(), + num_bytes=%batch_builder.num_bytes, + num_millis=%now.elapsed().as_millis(), + "sending doc batch to indexer" + ); + let message = batch_builder.build(); + source_sink.send_raw_doc_batch(message, ctx).await?; + } + Ok(wait_before_next_batch) + } + + async fn suggest_truncate( + &mut self, + checkpoint: SourceCheckpoint, + ctx: &SourceContext, + ) -> anyhow::Result<()> { + let Some(position) = checkpoint.position_for_partition(&self.partition_id) else { + return Ok(()); + }; + let Some(published_up_to) = position.as_u64() else { + return Ok(()); + }; + ctx.protect_future(ack_up_to( + &self.nats_client, + &mut self.pending_acks, + published_up_to, + )) + .await; + Ok(()) + } + + fn is_drained(&self) -> bool { + // Every delivered message is pending from `process_message` until the + // server confirmed its ack: an empty map means everything delivered so + // far has been published and acknowledged. + self.pending_acks.is_empty() + } + + fn name(&self) -> String { + format!("{self:?}") + } + + async fn finalize( + &mut self, + _exit_status: &ActorExitStatus, + _ctx: &SourceContext, + ) -> anyhow::Result<()> { + // The durable consumer itself is left untouched. Any ack still pending + // at this point was not settled by a drain: its message is redelivered + // after the consumer's `ack_wait`. + // + // Messages the client prefetched but that were never processed are + // negatively acknowledged so the server redelivers them right after + // this connection is gone, instead of after `ack_wait`. Only what is + // already buffered is NAK'd: the loop stops at the first pending poll, + // so a refill in flight cannot keep it alive. + while let Some(Some(message_res)) = self.message_stream.next().now_or_never() { + let Ok(message) = message_res else { + continue; + }; + let Some(ack_subject) = message.message.reply else { + continue; + }; + let _ = self + .nats_client + .publish(ack_subject, AckKind::Nak(Some(NAK_REDELIVERY_DELAY)).into()) + .await; + } + let _ = self.nats_client.drain().await; + Ok(()) + } + + fn observable_state(&self) -> JsonValue { + let num_pending_acks = self.pending_acks.len(); + json!({ + "index_id": self.source_runtime.index_id(), + "source_id": self.source_runtime.source_id(), + "stream": self.source_params.stream, + "consumer_name": self.consumer_name, + "num_bytes_processed": self.state.num_bytes_processed, + "num_messages_processed": self.state.num_messages_processed, + "num_invalid_messages": self.state.num_invalid_messages, + "num_pending_acks": num_pending_acks, + }) + } +} + +/// Bounds the concurrent server-confirmed ack requests: enough to hide the +/// round-trip latency without flooding the connection. +const MAX_CONCURRENT_ACK_REQUESTS: usize = 128; + +/// Pause before pulling again after a transient stream error, to avoid a hot +/// error loop while the connection recovers. `SuggestTruncate` is still +/// processed in the meantime. +const TRANSIENT_STREAM_ERROR_BACKOFF: Duration = Duration::from_secs(1); + +/// Delay before the server redelivers a negatively acknowledged message: long +/// enough for this connection to be gone, so the redelivery goes to a +/// surviving pipeline instead of bouncing back into this one's outstanding +/// pull request. +const NAK_REDELIVERY_DELAY: Duration = Duration::from_secs(1); + +/// Errors the pull stream recovers from on its own: the subscription stays +/// usable and polling simply resumes. `MissingHeartbeat` in particular also +/// fires when the stream was not polled for a while (e.g. under downstream +/// backpressure), not only when the server is unreachable. +fn is_transient_stream_error(error: &MessagesError) -> bool { + matches!( + error.kind(), + MessagesErrorKind::MissingHeartbeat + | MessagesErrorKind::Pull + | MessagesErrorKind::NoResponders + ) +} + +/// Sends server-confirmed acknowledgments ("double acks") for the messages up +/// to `published_up_to`. +/// +/// Entries only leave the pending map once the server confirmed their ack, so +/// `is_drained` never reports progress that could still be lost: a +/// fire-and-forget ack can be dropped by a reconnection even after a +/// successful client-side flush. A retried ack for a message the server +/// already saw acknowledged is simply confirmed again. Whatever could not be +/// acknowledged is redelivered after the consumer's `ack_wait`. +async fn ack_up_to( + nats_client: &async_nats::Client, + pending_acks: &mut BTreeMap, + published_up_to: u64, +) { + let acks: Vec<(u64, Subject)> = pending_acks + .range(..=published_up_to) + .map(|(delivery_counter, ack_subject)| (*delivery_counter, ack_subject.clone())) + .collect(); + if acks.is_empty() { + return; + } + let mut ack_results = futures::stream::iter(acks) + .map(|(delivery_counter, ack_subject)| async move { + // Mirrors `jetstream::Message::double_ack()`: a request instead of + // a plain publish, so the server's reply confirms the ack. + let ack_result = nats_client.request(ack_subject, AckKind::Ack.into()).await; + (delivery_counter, ack_result) + }) + .buffer_unordered(MAX_CONCURRENT_ACK_REQUESTS); + let mut num_acks = 0usize; + let mut num_failed_acks = 0usize; + let mut last_ack_error = None; + while let Some((delivery_counter, ack_result)) = ack_results.next().await { + match ack_result { + Ok(_ack_reply) => { + pending_acks.remove(&delivery_counter); + num_acks += 1; + } + Err(error) => { + num_failed_acks += 1; + last_ack_error = Some(error); + } + } + } + if let Some(error) = last_ack_error { + warn!(%error, num_failed_acks, "failed to ack NATS messages"); + } + debug!(num_acks, "acked published messages"); +} + +/// Fetches the pre-provisioned durable consumer. Read-only by contract: the +/// consumer is never created nor updated. +async fn fetch_durable_consumer( + jetstream_stream: &jetstream::stream::Stream, + consumer_name: &str, +) -> anyhow::Result { + let consumer: PullConsumer = jetstream_stream + .get_consumer(consumer_name) + .await + .map_err(|error| anyhow!("failed to find NATS consumer `{consumer_name}`: {error}"))?; + let ack_policy = consumer.cached_info().config.ack_policy; + ensure!( + ack_policy == AckPolicy::Explicit, + "NATS consumer `{consumer_name}` must use the explicit ack policy, got `{ack_policy:?}`" + ); + Ok(consumer) +} + +/// W3C trace context extraction from NATS message headers, following the +/// pattern of `quickwit_common::tracing_utils` for gRPC metadata. +struct NatsHeaderExtractor<'a>(&'a HeaderMap); + +impl Extractor for NatsHeaderExtractor<'_> { + fn get(&self, key: &str) -> Option<&str> { + let header_name = HeaderName::from_str(key).ok()?; + self.0.get(header_name).map(|value| value.as_str()) + } + + fn keys(&self) -> Vec<&str> { + self.0.iter().map(|(key, _)| key.as_ref()).collect() + } +} + +fn extract_remote_context(headers: &HeaderMap) -> OtelContext { + global::get_text_map_propagator(|propagator| propagator.extract(&NatsHeaderExtractor(headers))) +} + +/// Builds a span parented on the publisher's trace when the message carries +/// a W3C `traceparent` header, stitching the processing of the message into +/// the publisher's distributed trace. Messages without a propagated context +/// cost nothing: no span is created. +fn remote_parented_span( + message: &jetstream::Message, + stream_sequence: u64, + source_runtime: &SourceRuntime, +) -> Option { + let headers = message.headers.as_ref()?; + let parent_context = extract_remote_context(headers); + if !parent_context.span().span_context().is_valid() { + return None; + } + let span = tracing::info_span!( + "process_nats_message", + index_id = %source_runtime.index_id(), + source_id = %source_runtime.source_id(), + subject = %message.subject, + stream_sequence, + ); + let _ = span.set_parent(parent_context); + Some(span) +} + +async fn connect_nats(params: &NatsSourceParams) -> anyhow::Result { + let mut connect_options = ConnectOptions::new(); + match params.authentication.clone() { + None => {} + Some(NatsSourceAuth::UserPassword { user, password }) => { + connect_options = connect_options.user_and_password(user, password); + } + Some(NatsSourceAuth::Token(token)) => { + connect_options = connect_options.token(token); + } + } + if let Some(tls) = ¶ms.tls { + if let Some(ca_certificates_path) = &tls.ca_certificates_path { + connect_options = + connect_options.add_root_certificates(PathBuf::from(ca_certificates_path)); + } + if let (Some(certificate_path), Some(key_path)) = + (&tls.client_certificate_path, &tls.client_key_path) + { + connect_options = connect_options + .add_client_certificate(PathBuf::from(certificate_path), PathBuf::from(key_path)); + } + } + let client = async_nats::connect_with_options(¶ms.uris, connect_options) + .await + .with_context(|| { + format!( + "failed to connect to NATS servers `{}`", + params.uris.join(", ") + ) + })?; + Ok(client) +} + +/// Connects to the NATS servers and binds to the JetStream stream's +/// pre-provisioned durable consumer. +async fn connect_and_fetch_consumer( + params: &NatsSourceParams, +) -> anyhow::Result<(async_nats::Client, PullConsumer)> { + let nats_client = connect_nats(params).await?; + let jetstream_ctx = jetstream::new(nats_client.clone()); + let jetstream_stream = jetstream_ctx + .get_stream(¶ms.stream) + .await + .with_context(|| format!("failed to find NATS JetStream stream `{}`", params.stream))?; + let consumer = fetch_durable_consumer(&jetstream_stream, ¶ms.consumer).await?; + Ok((nats_client, consumer)) +} + +/// Checks whether we can connect to the NATS servers and find the JetStream +/// stream and the pre-provisioned consumer. +pub(crate) async fn check_connectivity(params: &NatsSourceParams) -> anyhow::Result<()> { + connect_and_fetch_consumer(params).await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_extract_trace_context_from_headers() { + use opentelemetry::propagation::TextMapPropagator; + use opentelemetry_sdk::propagation::TraceContextPropagator; + + // The global propagator is not installed in tests, so the extractor + // is exercised against an explicit W3C propagator. + let propagator = TraceContextPropagator::new(); + + let mut headers = HeaderMap::new(); + headers.insert( + "traceparent", + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + ); + let span_context = propagator + .extract(&NatsHeaderExtractor(&headers)) + .span() + .span_context() + .clone(); + assert!(span_context.is_valid()); + assert_eq!( + span_context.trace_id().to_string(), + "4bf92f3577b34da6a3ce929d0e0e4736" + ); + assert_eq!(span_context.span_id().to_string(), "00f067aa0ba902b7"); + + let empty_headers = HeaderMap::new(); + let span_context = propagator + .extract(&NatsHeaderExtractor(&empty_headers)) + .span() + .span_context() + .clone(); + assert!(!span_context.is_valid()); + } +} + +#[cfg(all(test, feature = "nats-broker-tests"))] +mod nats_broker_tests { + use std::num::NonZeroUsize; + use std::ops::Range; + use std::sync::Arc; + + use bytes::Bytes; + use quickwit_actors::{ActorHandle, Inbox, Universe}; + use quickwit_config::{SourceConfig, SourceInputFormat, SourceParams}; + use quickwit_metastore::metastore_for_test; + use quickwit_proto::metastore::MetastoreServiceClient; + use quickwit_proto::types::IndexUid; + + use super::*; + use crate::actors::DocProcessor; + use crate::models::RawDocBatch; + use crate::source::test_setup_helper::setup_index; + use crate::source::tests::SourceRuntimeBuilder; + use crate::source::{SourceActor, SuggestTruncate, quickwit_supported_sources}; + + static NATS_URI: &str = "nats://localhost:4222"; + + async fn setup_nats_stream(stream_name: &str) -> jetstream::Context { + let client = async_nats::connect(NATS_URI).await.unwrap(); + let jetstream_ctx = jetstream::new(client); + jetstream_ctx + .create_stream(jetstream::stream::Config { + name: stream_name.to_string(), + subjects: vec![format!("{stream_name}.>")], + ..Default::default() + }) + .await + .unwrap(); + jetstream_ctx + } + + /// Publishes one JSON doc per ID on the subject and waits for each + /// publish ack, so stream sequences are assigned in `ids` order. Messages + /// carry a W3C `traceparent` header to exercise the trace propagation + /// path. + async fn publish_docs( + jetstream_ctx: &jetstream::Context, + subject: &str, + ids: Range, + ) -> Vec { + let mut headers = HeaderMap::new(); + headers.insert( + "traceparent", + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + ); + let mut docs = Vec::with_capacity(ids.len()); + for id in ids { + let doc = json!({ "id": id, "subject": subject }).to_string(); + jetstream_ctx + .publish_with_headers(subject.to_string(), headers.clone(), doc.clone().into()) + .await + .unwrap() + .await + .unwrap(); + docs.push(doc); + } + docs + } + + async fn create_source_actor( + universe: &Universe, + metastore: MetastoreServiceClient, + index_uid: IndexUid, + source_config: SourceConfig, + ) -> (ActorHandle, Inbox) { + let source_runtime = SourceRuntimeBuilder::new(index_uid, source_config) + .with_metastore(metastore) + .build(); + let source = quickwit_supported_sources() + .load_source(source_runtime) + .await + .unwrap(); + let (doc_processor_mailbox, doc_processor_inbox) = universe.create_test_mailbox(); + let source_actor = SourceActor::new(source, doc_processor_mailbox); + let (_source_mailbox, source_handle) = universe.spawn_builder().spawn(source_actor); + (source_handle, doc_processor_inbox) + } + + async fn wait_for_processed_messages( + source_handle: &ActorHandle, + num_expected: u64, + ) { + loop { + let observation = source_handle.observe().await; + let num_messages_processed = observation + .state + .get("num_messages_processed") + .unwrap() + .as_u64() + .unwrap(); + if num_messages_processed >= num_expected { + return; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + + fn merge_doc_batches(batches: Vec) -> RawDocBatch { + let mut merged_batch = RawDocBatch::default(); + for batch in batches { + merged_batch.docs.extend(batch.docs); + merged_batch + .checkpoint_delta + .extend(batch.checkpoint_delta) + .unwrap(); + } + merged_batch.docs.sort(); + merged_batch + } + + fn get_durable_source_config(stream: &str, consumer: &str) -> SourceConfig { + let source_id = append_random_suffix("test-nats-source--durable-source"); + SourceConfig { + source_id, + num_pipelines: NonZeroUsize::MIN, + enabled: true, + source_params: SourceParams::Nats(NatsSourceParams { + uris: vec![NATS_URI.to_string()], + stream: stream.to_string(), + consumer: consumer.to_string(), + tls: None, + authentication: None, + }), + transform_config: None, + input_format: SourceInputFormat::Json, + } + } + + /// Provisions the durable consumer the way an operator would: the source + /// itself only ever fetches it. + async fn provision_durable_consumer( + jetstream_ctx: &jetstream::Context, + stream: &str, + consumer_name: &str, + ) { + jetstream_ctx + .create_consumer_on_stream( + jetstream::consumer::pull::Config { + name: Some(consumer_name.to_string()), + durable_name: Some(consumer_name.to_string()), + ack_policy: AckPolicy::Explicit, + // Long enough for the tests to never hit a redelivery. + ack_wait: Duration::from_secs(300), + ..Default::default() + }, + stream, + ) + .await + .unwrap(); + } + + #[tokio::test] + async fn test_durable_mode_ingestion_and_ack() { + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-ack-consumer"; + provision_durable_consumer(&jetstream_ctx, &stream, consumer_name).await; + + let subject = format!("{stream}.logs"); + let expected_docs = publish_docs(&jetstream_ctx, &subject, 0..10).await; + + let index_id = append_random_suffix("test-nats-source--durable--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + let (source_handle, doc_processor_inbox) = + create_source_actor(&universe, metastore, index_uid, source_config).await; + + wait_for_processed_messages(&source_handle, 10).await; + + let batches: Vec = doc_processor_inbox.drain_for_test_typed(); + let batch = merge_doc_batches(batches); + assert_eq!(batch.docs, expected_docs); + // Positions are synthetic delivery counters on a per-pipeline + // partition, not stream sequences. + assert_eq!(batch.checkpoint_delta.num_partitions(), 1); + + // A `SuggestTruncate` simulates the split publication notification: + // it must release the acks of the published messages. + let checkpoint = batch.checkpoint_delta.get_source_checkpoint(); + source_handle + .mailbox() + .send_message(SuggestTruncate(checkpoint)) + .await + .unwrap(); + + let mut consumer: PullConsumer = jetstream_ctx + .get_consumer_from_stream(consumer_name, stream.as_str()) + .await + .unwrap(); + let mut acked = false; + for _ in 0..100 { + let consumer_info = consumer.info().await.unwrap(); + if consumer_info.num_ack_pending == 0 { + assert_eq!(consumer_info.ack_floor.stream_sequence, 10); + acked = true; + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + assert!(acked, "messages should be acked after truncation"); + + source_handle.quit().await; + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + + /// End-to-end drain: a full indexing pipeline is asked to drain, which + /// must publish the in-flight batches and flush their acks BEFORE the + /// drain replies — the exactly-once guarantee on planned teardowns. + #[tokio::test] + async fn test_durable_mode_graceful_drain_acks_before_teardown() { + use quickwit_actors::Universe; + use quickwit_common::temp_dir::TempDirectory; + use quickwit_config::IndexingSettings; + use quickwit_doc_mapper::default_doc_mapper_for_test; + use quickwit_ingest::IngesterPool; + use quickwit_proto::indexing::IndexingPipelineId; + use quickwit_proto::types::{NodeId, PipelineUid}; + use quickwit_storage::{RamStorage, StorageResolver}; + + use crate::actors::pipeline_shared::DrainPipeline; + use crate::merge_policy::default_merge_policy; + use crate::{IndexingPipeline, IndexingPipelineParams, IndexingSplitStore}; + + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-drain--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-drain-consumer"; + provision_durable_consumer(&jetstream_ctx, &stream, consumer_name).await; + + let index_id = append_random_suffix("test-nats-source--durable-drain--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + // Docs valid for the test doc mapper, so the pipeline indexes them. + let subject = format!("{stream}.logs"); + for id in 0..10 { + let doc = json!({"timestamp": 1_700_000_000 + id, "body": format!("drain test {id}")}) + .to_string(); + jetstream_ctx + .publish(subject.clone(), doc.into()) + .await + .unwrap() + .await + .unwrap(); + } + + let pipeline_id = IndexingPipelineId { + node_id: NodeId::from_str("test-node"), + index_uid, + source_id: source_config.source_id.clone(), + pipeline_uid: PipelineUid::for_test(0u128), + }; + let storage = Arc::new(RamStorage::default()); + let split_store = IndexingSplitStore::create_without_local_store_for_test(storage.clone()); + let (merge_planner_mailbox, _merge_planner_inbox) = universe.create_test_mailbox(); + let pipeline_params = IndexingPipelineParams { + pipeline_id, + doc_mapper: Arc::new(default_doc_mapper_for_test()), + source_config, + source_storage_resolver: StorageResolver::for_test(), + indexing_directory: TempDirectory::for_test(), + indexing_settings: IndexingSettings::for_test(), + fingerprinter_opt: None, + ingester_pool: IngesterPool::default(), + metastore, + queues_dir_path: PathBuf::from("./queues"), + storage, + split_store, + merge_policy: default_merge_policy(), + retention_policy: None, + max_concurrent_split_uploads_index: 4, + max_concurrent_split_uploads_merge: 5, + cooperative_indexing_permits: None, + merge_planner_mailbox_opt: Some(merge_planner_mailbox), + event_broker: Default::default(), + params_fingerprint: 42u64, + }; + let pipeline = IndexingPipeline::new(pipeline_params); + let (pipeline_mailbox, pipeline_handle) = universe.spawn_builder().spawn(pipeline); + + // Wait for the source to have delivered all the messages. The commit + // timeout has not elapsed: without the drain, nothing would have been + // published nor acked yet. + loop { + let observation = pipeline_handle.observe().await; + if observation.num_docs >= 10 { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + + pipeline_mailbox.send_message(DrainPipeline).await.unwrap(); + // The pipeline exits on its own once the in-flight batches are + // published and their acks flushed. + let (exit_status, _statistics) = pipeline_handle.join().await; + assert!(matches!(exit_status, ActorExitStatus::Success)); + + // The pipeline exit implies the acks are already flushed: no polling, + // no grace period. + let mut consumer: PullConsumer = jetstream_ctx + .get_consumer_from_stream(consumer_name, stream.as_str()) + .await + .unwrap(); + let consumer_info = consumer.info().await.unwrap(); + assert_eq!(consumer_info.num_ack_pending, 0); + assert_eq!(consumer_info.ack_floor.stream_sequence, 10); + + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + + #[tokio::test] + async fn test_durable_mode_load_balancing() { + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-lb--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-lb-consumer"; + provision_durable_consumer(&jetstream_ctx, &stream, consumer_name).await; + + let index_id = append_random_suffix("test-nats-source--durable-lb--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + let (source_handle_1, doc_processor_inbox_1) = create_source_actor( + &universe, + metastore.clone(), + index_uid.clone(), + source_config.clone(), + ) + .await; + let (source_handle_2, doc_processor_inbox_2) = + create_source_actor(&universe, metastore, index_uid, source_config).await; + + let subject = format!("{stream}.logs"); + let mut expected_docs = publish_docs(&jetstream_ctx, &subject, 0..20).await; + expected_docs.sort(); + + // Work-queue delivery splits the messages between the two pipelines + // in some arbitrary way; together they must cover all of them + // exactly once. + loop { + let num_processed_1 = source_handle_1 + .observe() + .await + .state + .get("num_messages_processed") + .unwrap() + .as_u64() + .unwrap(); + let num_processed_2 = source_handle_2 + .observe() + .await + .state + .get("num_messages_processed") + .unwrap() + .as_u64() + .unwrap(); + if num_processed_1 + num_processed_2 >= 20 { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + source_handle_1.quit().await; + source_handle_2.quit().await; + + let mut all_docs: Vec = Vec::new(); + for batch in doc_processor_inbox_1 + .drain_for_test_typed::() + .into_iter() + .chain(doc_processor_inbox_2.drain_for_test_typed::()) + { + all_docs.extend(batch.docs); + } + all_docs.sort(); + assert_eq!(all_docs, expected_docs); + + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + + #[tokio::test] + async fn test_durable_mode_missing_consumer() { + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-missing--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + + let index_id = append_random_suffix("test-nats-source--durable-missing--index"); + let source_config = get_durable_source_config(&stream, "does-not-exist"); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + let source_runtime = SourceRuntimeBuilder::new(index_uid, source_config) + .with_metastore(metastore) + .build(); + let load_source_result = quickwit_supported_sources() + .load_source(source_runtime) + .await; + assert!( + load_source_result.is_err(), + "binding to a missing durable consumer should fail" + ); + + jetstream_ctx.delete_stream(&stream).await.unwrap(); + } +} diff --git a/quickwit/quickwit-metastore/src/metastore/mod.rs b/quickwit/quickwit-metastore/src/metastore/mod.rs index c12b8b11b61..98b2e14864e 100644 --- a/quickwit/quickwit-metastore/src/metastore/mod.rs +++ b/quickwit/quickwit-metastore/src/metastore/mod.rs @@ -1513,6 +1513,7 @@ fn use_shard_api(params: &SourceParams) -> bool { SourceParams::Kinesis(_) => false, SourceParams::PubSub(_) => false, SourceParams::Pulsar(_) => false, + SourceParams::Nats(_) => false, SourceParams::Stdin => panic!("stdin cannot be checkpointed"), SourceParams::Vec(_) => false, SourceParams::Void(_) => false, From 39a89c8bf51d0b386790feffb125bd11a22a895d Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 14:52:20 +0200 Subject: [PATCH 02/21] fix: config --- docs/configuration/source-config.md | 2 +- .../src/source/nats_source.rs | 51 +++++-------------- 2 files changed, 15 insertions(+), 38 deletions(-) diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 41f5ad26787..82d32c05393 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -191,7 +191,7 @@ A tutorial is available [here](/docs/ingest-data/nats.md). The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The consumer must use the **explicit ack policy**: the source acknowledges each message once the split containing it is published, and the consumer's ack floor is the resume point. -Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged — with server-confirmed acknowledgments — before the pipeline stops, so nothing is indexed twice; messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. `ack_wait` must exceed the end-to-end publish latency (roughly the commit timeout plus the upload time), otherwise messages are redelivered while they are still being indexed. `max_ack_pending` bounds the messages in flight across all pipelines of the consumer, and therefore the aggregate throughput. +Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. `ack_wait` must exceed the end-to-end publish latency (roughly the commit timeout plus the upload time), otherwise messages are redelivered while they are still being indexed. `max_ack_pending` bounds the messages in flight across all pipelines of the consumer, and therefore the aggregate throughput. **Scaling** diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index aeb0acc66fd..4cd165c49bf 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -15,20 +15,9 @@ //! A source consuming a NATS JetStream stream through a pre-provisioned //! durable consumer. //! -//! The consumer is only ever fetched — never created, updated, nor deleted — -//! so its lifecycle, subject filters, deliver policy, and ack tuning -//! (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The -//! consumer must use the explicit ack policy: each message is acknowledged — -//! with a server-confirmed ack — once the split containing it is published. -//! On a planned teardown (node shutdown, pipeline reassignment) the pipeline -//! is drained first: the in-flight messages are published and acknowledged, -//! and the prefetched but unprocessed ones are negatively acknowledged for -//! immediate redelivery — so delivery is exactly-once. A crash degrades it to -//! at-least-once: the unacknowledged messages — including ones whose split -//! ends up published — are redelivered after the consumer's `ack_wait` and -//! indexed again. The consumer's ack floor is the resume point; the metastore -//! checkpoint only carries synthetic positions used to correlate published -//! batches with the acknowledgments they release. +//! The consumer is only ever fetched so its lifecycle, subject filters, +//! deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to +//! whoever provisioned it. //! //! Several indexing pipelines can share the consumer: NATS load-balances the //! messages across them, so scaling is a plain `num_pipelines` update. Being @@ -98,26 +87,6 @@ pub struct NatsSourceState { pub num_invalid_messages: u64, } -/// Source flavor binding to a pre-provisioned durable consumer. -/// -/// The consumer is only ever fetched — never created, updated, nor deleted — -/// so its lifecycle, subject filters, deliver policy, and ack tuning -/// (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. Each -/// message is acknowledged once the split containing it is published -/// ([`Source::suggest_truncate`]); see the module documentation for the -/// delivery semantics. -/// -/// The metastore checkpoint only carries synthetic positions (a per-pipeline -/// partition and a delivery counter) used to correlate published batches with -/// the acknowledgments they release; the resume point is the consumer's ack -/// floor, not the checkpoint. -/// -/// The source acknowledges inline in [`Source::suggest_truncate`], so once a -/// drain reports completion (see [`Source::is_drained`]) every ack has been -/// confirmed by the server. If the source actor is instead killed with batches -/// still in the pipeline, their `suggest_truncate` notifications are lost and -/// the unacknowledged messages — including ones whose split ends up published — -/// are redelivered after the consumer's `ack_wait`. pub struct NatsSource { source_runtime: SourceRuntime, source_params: NatsSourceParams, @@ -308,6 +277,10 @@ impl Source for NatsSource { Ok(()) } + fn should_be_drained(&self) -> bool { + true + } + fn is_drained(&self) -> bool { // Every delivered message is pending from `process_message` until the // server confirmed its ack: an empty map means everything delivered so @@ -442,8 +415,7 @@ async fn ack_up_to( debug!(num_acks, "acked published messages"); } -/// Fetches the pre-provisioned durable consumer. Read-only by contract: the -/// consumer is never created nor updated. +/// Fetches the pre-provisioned durable consumer. async fn fetch_durable_consumer( jetstream_stream: &jetstream::stream::Stream, consumer_name: &str, @@ -901,7 +873,12 @@ mod nats_broker_tests { tokio::time::sleep(Duration::from_millis(100)).await; } - pipeline_mailbox.send_message(DrainPipeline).await.unwrap(); + pipeline_mailbox + .send_message(DrainPipeline { + drain_timeout: Duration::from_secs(60), + }) + .await + .unwrap(); // The pipeline exits on its own once the in-flight batches are // published and their acks flushed. let (exit_status, _statistics) = pipeline_handle.join().await; From 3b5e971749118857ac74496297fa3b596a2675ab Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 14:59:27 +0200 Subject: [PATCH 03/21] fix: config --- .../src/source/nats_source.rs | 23 ++----------------- 1 file changed, 2 insertions(+), 21 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 4cd165c49bf..9c243e70be0 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -297,15 +297,8 @@ impl Source for NatsSource { _exit_status: &ActorExitStatus, _ctx: &SourceContext, ) -> anyhow::Result<()> { - // The durable consumer itself is left untouched. Any ack still pending - // at this point was not settled by a drain: its message is redelivered - // after the consumer's `ack_wait`. - // - // Messages the client prefetched but that were never processed are - // negatively acknowledged so the server redelivers them right after - // this connection is gone, instead of after `ack_wait`. Only what is - // already buffered is NAK'd: the loop stops at the first pending poll, - // so a refill in flight cannot keep it alive. + // Any ack still pending at this point was not completed by a drain: + // its message is redelivered after the consumer's `ack_wait`. while let Some(Some(message_res)) = self.message_stream.next().now_or_never() { let Ok(message) = message_res else { continue; @@ -367,13 +360,6 @@ fn is_transient_stream_error(error: &MessagesError) -> bool { /// Sends server-confirmed acknowledgments ("double acks") for the messages up /// to `published_up_to`. -/// -/// Entries only leave the pending map once the server confirmed their ack, so -/// `is_drained` never reports progress that could still be lost: a -/// fire-and-forget ack can be dropped by a reconnection even after a -/// successful client-side flush. A retried ack for a message the server -/// already saw acknowledged is simply confirmed again. Whatever could not be -/// acknowledged is redelivered after the consumer's `ack_wait`. async fn ack_up_to( nats_client: &async_nats::Client, pending_acks: &mut BTreeMap, @@ -415,7 +401,6 @@ async fn ack_up_to( debug!(num_acks, "acked published messages"); } -/// Fetches the pre-provisioned durable consumer. async fn fetch_durable_consumer( jetstream_stream: &jetstream::stream::Stream, consumer_name: &str, @@ -510,8 +495,6 @@ async fn connect_nats(params: &NatsSourceParams) -> anyhow::Result anyhow::Result<(async_nats::Client, PullConsumer)> { @@ -525,8 +508,6 @@ async fn connect_and_fetch_consumer( Ok((nats_client, consumer)) } -/// Checks whether we can connect to the NATS servers and find the JetStream -/// stream and the pre-provisioned consumer. pub(crate) async fn check_connectivity(params: &NatsSourceParams) -> anyhow::Result<()> { connect_and_fetch_consumer(params).await?; Ok(()) From 2d9c42cc9de9d20dce3441aa984d26917049729e Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 15:31:16 +0200 Subject: [PATCH 04/21] feat: add tests --- .../src/source/nats_source.rs | 352 +++++++++++++++++- 1 file changed, 332 insertions(+), 20 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 9c243e70be0..0cb0c3cd10b 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -518,7 +518,7 @@ mod tests { use super::*; #[test] - fn test_extract_trace_context_from_headers() { + fn extract_trace_context_from_headers() { use opentelemetry::propagation::TextMapPropagator; use opentelemetry_sdk::propagation::TraceContextPropagator; @@ -551,6 +551,32 @@ mod tests { .clone(); assert!(!span_context.is_valid()); } + + #[test] + fn transient_stream_error_classification() { + let transient_kinds = [ + MessagesErrorKind::MissingHeartbeat, + MessagesErrorKind::Pull, + MessagesErrorKind::NoResponders, + ]; + for kind in transient_kinds { + assert!( + is_transient_stream_error(&MessagesError::new(kind)), + "`{kind:?}` should be retried, not kill the pipeline" + ); + } + let terminal_kinds = [ + MessagesErrorKind::ConsumerDeleted, + MessagesErrorKind::PushBasedConsumer, + MessagesErrorKind::Other, + ]; + for kind in terminal_kinds { + assert!( + !is_transient_stream_error(&MessagesError::new(kind)), + "`{kind:?}` should kill the pipeline" + ); + } + } } #[cfg(all(test, feature = "nats-broker-tests"))] @@ -562,6 +588,7 @@ mod nats_broker_tests { use bytes::Bytes; use quickwit_actors::{ActorHandle, Inbox, Universe}; use quickwit_config::{SourceConfig, SourceInputFormat, SourceParams}; + use quickwit_metastore::checkpoint::SourceCheckpointDelta; use quickwit_metastore::metastore_for_test; use quickwit_proto::metastore::MetastoreServiceClient; use quickwit_proto::types::IndexUid; @@ -636,10 +663,15 @@ mod nats_broker_tests { (source_handle, doc_processor_inbox) } + /// Waits until the source reports at least `num_expected` processed + /// messages, panicking past `timeout` — kept tight where the delivery + /// delay itself is the assertion (NAK vs `ack_wait` redelivery). async fn wait_for_processed_messages( source_handle: &ActorHandle, num_expected: u64, + timeout: Duration, ) { + let deadline = Instant::now() + timeout; loop { let observation = source_handle.observe().await; let num_messages_processed = observation @@ -651,8 +683,34 @@ mod nats_broker_tests { if num_messages_processed >= num_expected { return; } + assert!( + Instant::now() < deadline, + "source did not process {num_expected} messages within {timeout:?}" + ); + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + + /// Waits until the consumer reports the expected ack state, panicking + /// after ~10s. + async fn wait_for_consumer_ack_state( + consumer: &mut PullConsumer, + expected_num_ack_pending: usize, + expected_ack_floor: u64, + ) { + for _ in 0..100 { + let consumer_info = consumer.info().await.unwrap(); + if consumer_info.num_ack_pending == expected_num_ack_pending + && consumer_info.ack_floor.stream_sequence == expected_ack_floor + { + return; + } tokio::time::sleep(Duration::from_millis(100)).await; } + panic!( + "consumer never reached {expected_num_ack_pending} pending acks with ack floor \ + {expected_ack_floor}" + ); } fn merge_doc_batches(batches: Vec) -> RawDocBatch { @@ -688,10 +746,11 @@ mod nats_broker_tests { /// Provisions the durable consumer the way an operator would: the source /// itself only ever fetches it. - async fn provision_durable_consumer( + async fn provision_durable_consumer_with_ack_wait( jetstream_ctx: &jetstream::Context, stream: &str, consumer_name: &str, + ack_wait: Duration, ) { jetstream_ctx .create_consumer_on_stream( @@ -699,8 +758,7 @@ mod nats_broker_tests { name: Some(consumer_name.to_string()), durable_name: Some(consumer_name.to_string()), ack_policy: AckPolicy::Explicit, - // Long enough for the tests to never hit a redelivery. - ack_wait: Duration::from_secs(300), + ack_wait, ..Default::default() }, stream, @@ -709,8 +767,23 @@ mod nats_broker_tests { .unwrap(); } + async fn provision_durable_consumer( + jetstream_ctx: &jetstream::Context, + stream: &str, + consumer_name: &str, + ) { + // An `ack_wait` long enough for the test to never hit a redelivery. + provision_durable_consumer_with_ack_wait( + jetstream_ctx, + stream, + consumer_name, + Duration::from_secs(300), + ) + .await; + } + #[tokio::test] - async fn test_durable_mode_ingestion_and_ack() { + async fn durable_mode_ingestion_and_ack() { let universe = Universe::with_accelerated_time(); let metastore = metastore_for_test(); let stream = append_random_suffix("test-nats-source--durable--stream"); @@ -728,7 +801,7 @@ mod nats_broker_tests { let (source_handle, doc_processor_inbox) = create_source_actor(&universe, metastore, index_uid, source_config).await; - wait_for_processed_messages(&source_handle, 10).await; + wait_for_processed_messages(&source_handle, 10, Duration::from_secs(60)).await; let batches: Vec = doc_processor_inbox.drain_for_test_typed(); let batch = merge_doc_batches(batches); @@ -750,17 +823,7 @@ mod nats_broker_tests { .get_consumer_from_stream(consumer_name, stream.as_str()) .await .unwrap(); - let mut acked = false; - for _ in 0..100 { - let consumer_info = consumer.info().await.unwrap(); - if consumer_info.num_ack_pending == 0 { - assert_eq!(consumer_info.ack_floor.stream_sequence, 10); - acked = true; - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - assert!(acked, "messages should be acked after truncation"); + wait_for_consumer_ack_state(&mut consumer, 0, 10).await; source_handle.quit().await; jetstream_ctx.delete_stream(&stream).await.unwrap(); @@ -771,7 +834,7 @@ mod nats_broker_tests { /// must publish the in-flight batches and flush their acks BEFORE the /// drain replies — the exactly-once guarantee on planned teardowns. #[tokio::test] - async fn test_durable_mode_graceful_drain_acks_before_teardown() { + async fn durable_mode_graceful_drain_acks_before_teardown() { use quickwit_actors::Universe; use quickwit_common::temp_dir::TempDirectory; use quickwit_config::IndexingSettings; @@ -880,7 +943,7 @@ mod nats_broker_tests { } #[tokio::test] - async fn test_durable_mode_load_balancing() { + async fn durable_mode_load_balancing() { let universe = Universe::with_accelerated_time(); let metastore = metastore_for_test(); let stream = append_random_suffix("test-nats-source--durable-lb--stream"); @@ -950,7 +1013,256 @@ mod nats_broker_tests { } #[tokio::test] - async fn test_durable_mode_missing_consumer() { + async fn durable_mode_crash_redelivery_after_ack_wait() { + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-crash--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-crash-consumer"; + provision_durable_consumer_with_ack_wait( + &jetstream_ctx, + &stream, + consumer_name, + Duration::from_secs(3), + ) + .await; + + let subject = format!("{stream}.logs"); + let expected_docs = publish_docs(&jetstream_ctx, &subject, 0..10).await; + + let index_id = append_random_suffix("test-nats-source--durable-crash--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + let (source_handle_1, _doc_processor_inbox_1) = create_source_actor( + &universe, + metastore.clone(), + index_uid.clone(), + source_config.clone(), + ) + .await; + wait_for_processed_messages(&source_handle_1, 10, Duration::from_secs(60)).await; + // No `SuggestTruncate` was sent: quitting here loses the acks of the + // processed messages, like a crash would. + source_handle_1.quit().await; + + let (source_handle_2, doc_processor_inbox_2) = + create_source_actor(&universe, metastore, index_uid, source_config).await; + wait_for_processed_messages(&source_handle_2, 10, Duration::from_secs(30)).await; + + let batches: Vec = doc_processor_inbox_2.drain_for_test_typed(); + // A message can be redelivered more than once if the test straddles + // several `ack_wait` windows. + let mut redelivered_docs = merge_doc_batches(batches).docs; + redelivered_docs.dedup(); + assert_eq!(redelivered_docs, expected_docs); + + source_handle_2.quit().await; + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + + #[tokio::test] + async fn durable_mode_nak_redelivers_prefetched_messages_promptly() { + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-nak--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-nak-consumer"; + // The `ack_wait` is much longer than the test: only a NAK can explain + // a prompt redelivery. + provision_durable_consumer(&jetstream_ctx, &stream, consumer_name).await; + + let index_id = append_random_suffix("test-nats-source--durable-nak--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + // The source is driven by hand rather than through a `SourceActor`: + // an idle actor keeps polling, so it would process the messages + // instead of leaving them prefetched. + let source_runtime = SourceRuntimeBuilder::new(index_uid.clone(), source_config.clone()) + .with_metastore(metastore.clone()) + .build(); + let mut source = quickwit_supported_sources() + .load_source(source_runtime) + .await + .unwrap(); + let (source_mailbox, _source_inbox) = universe.create_test_mailbox::(); + let (doc_processor_mailbox, _doc_processor_inbox) = + universe.create_test_mailbox::(); + let source_sink = SourceSink::from(doc_processor_mailbox); + let (observable_state_tx, _observable_state_rx) = + tokio::sync::watch::channel(JsonValue::Null); + let ctx: SourceContext = + quickwit_actors::ActorContext::for_test(&universe, source_mailbox, observable_state_tx); + + // A first empty emit issues the pull request: the messages published + // next are prefetched into the client buffer but never processed. + source.emit_batches(&source_sink, &ctx).await.unwrap(); + let subject = format!("{stream}.logs"); + let expected_docs = publish_docs(&jetstream_ctx, &subject, 0..10).await; + tokio::time::sleep(Duration::from_secs(1)).await; + assert_eq!( + source + .observable_state() + .get("num_messages_processed") + .unwrap(), + &json!(0), + "the messages must be prefetched, not processed" + ); + + source.finalize(&ActorExitStatus::Quit, &ctx).await.unwrap(); + drop(source); + + let (source_handle_2, doc_processor_inbox_2) = + create_source_actor(&universe, metastore, index_uid, source_config).await; + wait_for_processed_messages(&source_handle_2, 10, Duration::from_secs(20)).await; + + let batches: Vec = doc_processor_inbox_2.drain_for_test_typed(); + assert_eq!(merge_doc_batches(batches).docs, expected_docs); + + source_handle_2.quit().await; + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + + /// A truncate covering only part of the delivered messages must ack + /// exactly up to its position and leave the rest pending. + #[tokio::test] + async fn durable_mode_partial_truncate_acks_up_to_position() { + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-partial--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-partial-consumer"; + provision_durable_consumer(&jetstream_ctx, &stream, consumer_name).await; + + let subject = format!("{stream}.logs"); + publish_docs(&jetstream_ctx, &subject, 0..10).await; + + let index_id = append_random_suffix("test-nats-source--durable-partial--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + let (source_handle, doc_processor_inbox) = + create_source_actor(&universe, metastore, index_uid, source_config).await; + wait_for_processed_messages(&source_handle, 10, Duration::from_secs(60)).await; + + let batches: Vec = doc_processor_inbox.drain_for_test_typed(); + let batch = merge_doc_batches(batches); + let partition_id = batch.checkpoint_delta.partitions().next().unwrap().clone(); + + // Sequential publishes on a single pipeline: delivery counters map + // one-to-one to stream sequences, so truncating up to position 5 must + // ack stream sequences 1 to 5. + let partial_checkpoint = SourceCheckpointDelta::from_partition_delta( + partition_id, + Position::Beginning, + Position::offset(5u64), + ) + .unwrap() + .get_source_checkpoint(); + source_handle + .mailbox() + .send_message(SuggestTruncate(partial_checkpoint)) + .await + .unwrap(); + + let mut consumer: PullConsumer = jetstream_ctx + .get_consumer_from_stream(consumer_name, stream.as_str()) + .await + .unwrap(); + wait_for_consumer_ack_state(&mut consumer, 5, 5).await; + let observation = source_handle.observe().await; + assert_eq!( + observation.state.get("num_pending_acks").unwrap(), + &json!(5) + ); + + // The remainder is released by the next truncate. + let full_checkpoint = batch.checkpoint_delta.get_source_checkpoint(); + source_handle + .mailbox() + .send_message(SuggestTruncate(full_checkpoint)) + .await + .unwrap(); + wait_for_consumer_ack_state(&mut consumer, 0, 10).await; + + source_handle.quit().await; + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + + /// An empty payload is counted invalid and skipped from the batch, but + /// its message must still be acknowledged on truncate: a poison message + /// must not wedge the consumer nor a drain. + #[tokio::test] + async fn durable_mode_empty_message_acked_and_counted_invalid() { + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-empty--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-empty-consumer"; + provision_durable_consumer(&jetstream_ctx, &stream, consumer_name).await; + + let subject = format!("{stream}.logs"); + let valid_docs = [ + json!({ "id": 0, "subject": subject }).to_string(), + json!({ "id": 2, "subject": subject }).to_string(), + ]; + for payload in [ + Bytes::from(valid_docs[0].clone()), + Bytes::new(), + Bytes::from(valid_docs[1].clone()), + ] { + jetstream_ctx + .publish(subject.clone(), payload) + .await + .unwrap() + .await + .unwrap(); + } + + let index_id = append_random_suffix("test-nats-source--durable-empty--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; + + let (source_handle, doc_processor_inbox) = + create_source_actor(&universe, metastore, index_uid, source_config).await; + wait_for_processed_messages(&source_handle, 3, Duration::from_secs(60)).await; + + let observation = source_handle.observe().await; + assert_eq!( + observation.state.get("num_invalid_messages").unwrap(), + &json!(1) + ); + + let batches: Vec = doc_processor_inbox.drain_for_test_typed(); + let batch = merge_doc_batches(batches); + assert_eq!(batch.docs, valid_docs); + + // The checkpoint covers the empty message too, and truncating acks it + // along with the valid ones. + source_handle + .mailbox() + .send_message(SuggestTruncate( + batch.checkpoint_delta.get_source_checkpoint(), + )) + .await + .unwrap(); + let mut consumer: PullConsumer = jetstream_ctx + .get_consumer_from_stream(consumer_name, stream.as_str()) + .await + .unwrap(); + wait_for_consumer_ack_state(&mut consumer, 0, 3).await; + + source_handle.quit().await; + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + + #[tokio::test] + async fn durable_mode_missing_consumer() { let metastore = metastore_for_test(); let stream = append_random_suffix("test-nats-source--durable-missing--stream"); let jetstream_ctx = setup_nats_stream(&stream).await; From 8ac7c8d6f7a2f4580c0d062fd80394cd56935bb0 Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 15:59:59 +0200 Subject: [PATCH 05/21] feat: add tests --- quickwit/quickwit-config/src/source_config/serialize.rs | 3 ++- quickwit/quickwit-serve/src/index_api/rest_handler.rs | 3 +-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/quickwit/quickwit-config/src/source_config/serialize.rs b/quickwit/quickwit-config/src/source_config/serialize.rs index 25e92fe0a28..3b412a9ab19 100644 --- a/quickwit/quickwit-config/src/source_config/serialize.rs +++ b/quickwit/quickwit-config/src/source_config/serialize.rs @@ -132,7 +132,8 @@ impl SourceConfigForSerialization { _ => { if self.num_pipelines > 1 { bail!( - "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or NATS sources. open an issue https://github.com/quickwit-oss/quickwit/issues if you need the feature for other source types" + "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or NATS sources.\ + open an issue https://github.com/quickwit-oss/quickwit/issues if you need the feature for other source types" ); } } diff --git a/quickwit/quickwit-serve/src/index_api/rest_handler.rs b/quickwit/quickwit-serve/src/index_api/rest_handler.rs index bc4c8d9b105..09cd6eb09ea 100644 --- a/quickwit/quickwit-serve/src/index_api/rest_handler.rs +++ b/quickwit/quickwit-serve/src/index_api/rest_handler.rs @@ -1100,8 +1100,7 @@ mod tests { assert_eq!(resp.status(), 400); let body = std::str::from_utf8(resp.body()).unwrap(); assert!(body.contains( - "Quickwit currently supports multiple pipelines only for GCP PubSub or Kafka \ - sources" + "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or NATS sources" )); } { From 928f4a9a5c02d9d1b870b1ee54387eb5a6d63346 Mon Sep 17 00:00:00 2001 From: tprevot Date: Mon, 7 Sep 2026 16:34:46 +0200 Subject: [PATCH 06/21] fix: fmt --- quickwit/quickwit-serve/src/index_api/rest_handler.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/quickwit/quickwit-serve/src/index_api/rest_handler.rs b/quickwit/quickwit-serve/src/index_api/rest_handler.rs index 09cd6eb09ea..1324508ff68 100644 --- a/quickwit/quickwit-serve/src/index_api/rest_handler.rs +++ b/quickwit/quickwit-serve/src/index_api/rest_handler.rs @@ -1100,7 +1100,8 @@ mod tests { assert_eq!(resp.status(), 400); let body = std::str::from_utf8(resp.body()).unwrap(); assert!(body.contains( - "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or NATS sources" + "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or \ + NATS sources" )); } { From a19e18f8439d769bb170642ff0ce8129f91cdf39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9odore=20Pr=C3=A9vot?= Date: Tue, 8 Sep 2026 14:11:18 +0000 Subject: [PATCH 07/21] fix: nats --- .../src/source/nats_source.rs | 124 +++++++++++++++--- 1 file changed, 105 insertions(+), 19 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 0cb0c3cd10b..61f2cfbc598 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -126,11 +126,22 @@ impl NatsSource { source_id=%source_runtime.source_id(), stream=%source_params.stream, %consumer_name, + pull_max_bytes_per_batch=%pull_max_bytes_per_batch(), + pull_max_messages_per_batch=%pull_max_messages_per_batch(), "starting NATS source" ); let (nats_client, consumer) = connect_and_fetch_consumer(&source_params).await?; let message_stream = consumer + .stream() + .max_messages_per_batch(pull_max_messages_per_batch()) + .max_bytes_per_batch(pull_max_bytes_per_batch()) + // `stream()` defaults the heartbeat to zero, i.e. disabled, unlike + // `messages()`. It has to be set back, or a stalled subscription is + // never reported and `is_transient_stream_error` has nothing to act + // on. + .heartbeat(PULL_IDLE_HEARTBEAT) + .expires(PULL_EXPIRES) .messages() .await .context("failed to subscribe to NATS consumer messages")?; @@ -299,19 +310,40 @@ impl Source for NatsSource { ) -> anyhow::Result<()> { // Any ack still pending at this point was not completed by a drain: // its message is redelivered after the consumer's `ack_wait`. + let mut num_naks = 0usize; while let Some(Some(message_res)) = self.message_stream.next().now_or_never() { - let Ok(message) = message_res else { - continue; + let message = match message_res { + Ok(message) => message, + Err(error) => { + warn!(%error, "failed to pull a prefetched NATS message to negatively acknowledge it"); + continue; + } }; let Some(ack_subject) = message.message.reply else { + warn!("prefetched NATS message carries no reply subject to negatively acknowledge it on"); continue; }; - let _ = self + if let Err(error) = self .nats_client .publish(ack_subject, AckKind::Nak(Some(NAK_REDELIVERY_DELAY)).into()) - .await; + .await + { + warn!(%error, "failed to negatively acknowledge a prefetched NATS message"); + continue; + } + num_naks += 1; + } + // The negative acknowledgments above are plain publishes, and `drain` + // returns as soon as its command is queued rather than once the + // connection has been written out. Without this flush they can be lost + // with the connection, and their messages then wait a full `ack_wait` + // instead of being redelivered to a surviving pipeline immediately. + if let Err(error) = self.nats_client.flush().await { + warn!(%error, num_naks, "failed to flush NATS negative acknowledgments"); + } + if let Err(error) = self.nats_client.drain().await { + warn!(%error, "failed to drain the NATS connection"); } - let _ = self.nats_client.drain().await; Ok(()) } @@ -332,7 +364,55 @@ impl Source for NatsSource { /// Bounds the concurrent server-confirmed ack requests: enough to hide the /// round-trip latency without flooding the connection. -const MAX_CONCURRENT_ACK_REQUESTS: usize = 128; +/// Caps a pull batch in bytes, which the message-count cap alone cannot do. +/// +/// The server buffers a pull response on the connection's outbound queue, and a +/// connection that exceeds `max_pending` (64 MiB by default) is declared a slow +/// consumer and closed. Messages already written to it are lost while still +/// counted as delivered, so they only come back after `ack_wait` expires: with +/// 1 MiB messages, an uncapped batch of 200 asks the server for 200 MiB and the +/// source stalls for a whole `ack_wait` at a time. +/// +/// The value is bounded on both sides. It must stay well below the server's +/// `max_pending`, since several pull requests can be outstanding at once, and it +/// must stay *above* the server's `max_payload`: the server answers a pull whose +/// `max_bytes` is smaller than the next message with `409 Message Size Exceeds +/// MaxBytes` and delivers nothing, which turns an intermittent stall into a +/// permanent one. 10 MiB clears the 8 MiB `max_payload` ceiling a server can be +/// configured with in practice, and leaves a six-fold margin against the default +/// `max_pending`. +/// Overridable while the value is being tuned against real message sizes. +const DEFAULT_PULL_MAX_BYTES_PER_BATCH: usize = 10 * 1024 * 1024; + +fn pull_max_bytes_per_batch() -> usize { + quickwit_common::get_from_env_cached!( + usize, + "QW_NATS_PULL_MAX_BYTES_PER_BATCH", + DEFAULT_PULL_MAX_BYTES_PER_BATCH, + false + ) +} + +/// Messages per pull batch. The byte cap above is what actually bounds a batch; +/// this only binds for small messages, where it caps a pull far below the byte +/// budget (200 messages of 1 KiB is under 200 KiB against a 10 MiB budget) and +/// costs a round-trip per 200 messages. Left high so the byte cap is the single +/// thing deciding batch size, and overridable for the same reason as above. +const DEFAULT_PULL_MAX_MESSAGES_PER_BATCH: usize = 100_000; + +fn pull_max_messages_per_batch() -> usize { + quickwit_common::get_from_env_cached!( + usize, + "QW_NATS_PULL_MAX_MESSAGES_PER_BATCH", + DEFAULT_PULL_MAX_MESSAGES_PER_BATCH, + false + ) +} + +/// Both kept at the `Consumer::messages` defaults, so switching to the builder +/// changes the batch bounds and nothing else. +const PULL_IDLE_HEARTBEAT: Duration = Duration::from_secs(15); +const PULL_EXPIRES: Duration = Duration::from_secs(30); /// Pause before pulling again after a transient stream error, to avoid a hot /// error loop while the connection recovers. `SuggestTruncate` is still @@ -358,8 +438,20 @@ fn is_transient_stream_error(error: &MessagesError) -> bool { ) } -/// Sends server-confirmed acknowledgments ("double acks") for the messages up -/// to `published_up_to`. +/// Acknowledges the messages up to `published_up_to`. +/// +/// The acknowledgments are plain publishes, not requests: a request per message +/// would put a server round-trip on the source actor's own task, which is the +/// task that also pulls, and the actor handles one message at a time. What makes +/// that safe is that the connection is only ever torn down through `finalize`, +/// which flushes before draining it, so an acknowledgment queued here is written +/// out before the connection goes away. +/// +/// The consequence for [`Source::is_drained`] is worth being explicit about: a +/// message leaves `pending_acks` once its acknowledgment is queued on the +/// connection rather than once the server has answered for it. A drain therefore +/// declares itself finished on "queued and about to be flushed" instead of +/// "server has it", and it is `finalize`'s flush that closes that gap. async fn ack_up_to( nats_client: &async_nats::Client, pending_acks: &mut BTreeMap, @@ -372,20 +464,12 @@ async fn ack_up_to( if acks.is_empty() { return; } - let mut ack_results = futures::stream::iter(acks) - .map(|(delivery_counter, ack_subject)| async move { - // Mirrors `jetstream::Message::double_ack()`: a request instead of - // a plain publish, so the server's reply confirms the ack. - let ack_result = nats_client.request(ack_subject, AckKind::Ack.into()).await; - (delivery_counter, ack_result) - }) - .buffer_unordered(MAX_CONCURRENT_ACK_REQUESTS); let mut num_acks = 0usize; let mut num_failed_acks = 0usize; let mut last_ack_error = None; - while let Some((delivery_counter, ack_result)) = ack_results.next().await { - match ack_result { - Ok(_ack_reply) => { + for (delivery_counter, ack_subject) in acks { + match nats_client.publish(ack_subject, AckKind::Ack.into()).await { + Ok(()) => { pending_acks.remove(&delivery_counter); num_acks += 1; } @@ -396,6 +480,8 @@ async fn ack_up_to( } } if let Some(error) = last_ack_error { + // The messages stay in `pending_acks`, so the next truncate retries them + // and a drain keeps waiting rather than declaring itself finished. warn!(%error, num_failed_acks, "failed to ack NATS messages"); } debug!(num_acks, "acked published messages"); From 945370252b8c6ce94a63505703de8a3b53b710f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9odore=20Pr=C3=A9vot?= Date: Tue, 8 Sep 2026 14:39:30 +0000 Subject: [PATCH 08/21] fix: nat source --- .../src/source/nats_source.rs | 45 +++++++++++-------- 1 file changed, 27 insertions(+), 18 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 61f2cfbc598..4612f04260a 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -438,20 +438,23 @@ fn is_transient_stream_error(error: &MessagesError) -> bool { ) } -/// Acknowledges the messages up to `published_up_to`. -/// -/// The acknowledgments are plain publishes, not requests: a request per message -/// would put a server round-trip on the source actor's own task, which is the -/// task that also pulls, and the actor handles one message at a time. What makes -/// that safe is that the connection is only ever torn down through `finalize`, -/// which flushes before draining it, so an acknowledgment queued here is written -/// out before the connection goes away. +/// Bounds the concurrent server-confirmed ack requests: enough to hide the +/// round-trip latency without flooding the connection. Measured against 1 KiB +/// messages: 1024 is indistinguishable from 128, and 8192 is 20 % slower. +const MAX_CONCURRENT_ACK_REQUESTS: usize = 128; + +/// Sends server-confirmed acknowledgments ("double acks") for the messages up +/// to `published_up_to`. /// -/// The consequence for [`Source::is_drained`] is worth being explicit about: a -/// message leaves `pending_acks` once its acknowledgment is queued on the -/// connection rather than once the server has answered for it. A drain therefore -/// declares itself finished on "queued and about to be flushed" instead of -/// "server has it", and it is `finalize`'s flush that closes that gap. +/// The confirmation is not for error reporting: it is what holds a message in +/// `pending_acks` long enough for [`Source::is_drained`] to still be false when +/// `Drain` arrives. `Drain` queues its force-commit batch and then checks +/// `is_drained` immediately, so a source that reports itself drained at that +/// instant exits before the in-flight documents are published, and they are +/// redelivered as duplicates. Replacing these requests with plain publishes was +/// measured at 37 % less broker CPU and 40 % more throughput at four pipelines, +/// and it broke exactly-once on `num_pipelines` reassignment (+91 581 duplicates +/// on a 2 000 000 document corpus), which is why they are requests. async fn ack_up_to( nats_client: &async_nats::Client, pending_acks: &mut BTreeMap, @@ -464,12 +467,20 @@ async fn ack_up_to( if acks.is_empty() { return; } + let mut ack_results = futures::stream::iter(acks) + .map(|(delivery_counter, ack_subject)| async move { + // Mirrors `jetstream::Message::double_ack()`, but through the + // client's muxed inbox rather than a subscription per message. + let ack_result = nats_client.request(ack_subject, AckKind::Ack.into()).await; + (delivery_counter, ack_result) + }) + .buffer_unordered(MAX_CONCURRENT_ACK_REQUESTS); let mut num_acks = 0usize; let mut num_failed_acks = 0usize; let mut last_ack_error = None; - for (delivery_counter, ack_subject) in acks { - match nats_client.publish(ack_subject, AckKind::Ack.into()).await { - Ok(()) => { + while let Some((delivery_counter, ack_result)) = ack_results.next().await { + match ack_result { + Ok(_ack_reply) => { pending_acks.remove(&delivery_counter); num_acks += 1; } @@ -480,8 +491,6 @@ async fn ack_up_to( } } if let Some(error) = last_ack_error { - // The messages stay in `pending_acks`, so the next truncate retries them - // and a drain keeps waiting rather than declaring itself finished. warn!(%error, num_failed_acks, "failed to ack NATS messages"); } debug!(num_acks, "acked published messages"); From 17f226cca353c6b6859c955deb3bb65b2628b4c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9odore=20Pr=C3=A9vot?= Date: Tue, 8 Sep 2026 15:37:55 +0000 Subject: [PATCH 09/21] feat: nats source --- docs/configuration/source-config.md | 33 +++- docs/ingest-data/nats.md | 6 +- .../src/source/nats_source.rs | 145 ++++++++++++++++-- 3 files changed, 161 insertions(+), 23 deletions(-) diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 82d32c05393..a9dd816ba10 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -189,15 +189,36 @@ A NATS source reads data from a [NATS JetStream](https://docs.nats.io/nats-conce A tutorial is available [here](/docs/ingest-data/nats.md). -The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The consumer must use the **explicit ack policy**: the source acknowledges each message once the split containing it is published, and the consumer's ack floor is the resume point. +The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The consumer must use either the **explicit** or the **all** ack policy: the source acknowledges messages once the split containing them is published, and the consumer's ack floor is the resume point. Which of the two to pick is a throughput decision, covered in [Ack policy](#ack-policy) below. -Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. `ack_wait` must exceed the end-to-end publish latency (roughly the commit timeout plus the upload time), otherwise messages are redelivered while they are still being indexed. `max_ack_pending` bounds the messages in flight across all pipelines of the consumer, and therefore the aggregate throughput. +Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. `ack_wait` must exceed the end-to-end publish latency (roughly the commit timeout plus the upload time), otherwise messages are redelivered while they are still being indexed. -**Scaling** +#### Ack policy -Multiple indexing pipelines can share the consumer: NATS load-balances the messages across them, so scaling up or down is a plain `num_pipelines` update, and the control plane places the pipelines across the indexers of the cluster. +The ack policy decides how many acknowledgments a corpus costs, and at small message sizes that is what bounds the source's throughput. JetStream delivery is cheap; per-message acknowledgment is not. -**Monitoring** +- **`explicit`** — one acknowledgment per message. Required whenever several pipelines share the consumer, because acknowledging a message under the `all` policy would also acknowledge the messages a sibling pipeline is still indexing. +- **`all`** — one acknowledgment covers every message up to it, so the source sends a single acknowledgment per commit instead of one per message, the same shape as a Kafka offset commit. Quickwit rejects a source configured with `num_pipelines` greater than 1 on such a consumer, since it is only sound when the consumer has a **single puller**. + +Prefer `all`, and get parallelism from partitioning rather than from several pipelines on one consumer. On a 16 vCPU host indexing 1 KiB messages, moving one pipeline from `explicit` to `all` cut broker CPU by 56 % and total CPU per indexed GiB by 31 %; the gain is larger the smaller the messages, and negligible above a few hundred KiB, where the message rate is low enough that acknowledgments no longer matter. + +#### Scaling + +There are two ways to spread a stream over several pipelines, and they do not perform alike. + +**One stream per pipeline** (recommended). Partition the data into one stream per pipeline — the equivalent of Kafka partitions — give each stream its own consumer, and add one Quickwit source with `num_pipelines: 1` per stream. Each consumer then has a single puller, which is what makes the `all` ack policy usable, and no consumer reads a message it does not own. This is the only shape that scales: on the host above, sixteen streams indexed the same 1 KiB corpus 2.7 times faster than sixteen pipelines sharing one consumer. + +**One shared consumer** (simplest). Several pipelines bind to the same consumer and NATS load-balances the messages across them, so scaling is a plain `num_pipelines` update and the control plane places the pipelines across the indexers of the cluster. This requires the `explicit` policy, and the resulting per-message acknowledgments cap the aggregate throughput regardless of how many pipelines are added — adding pipelines beyond a handful buys close to nothing at small message sizes. + +:::warning + +Do not partition a single stream by giving each pipeline a subject-filtered consumer. Unless a subject's messages happen to be contiguous in the stream, every consumer scans past the messages it does not own, and the broker saturates: sixteen subject-filtered consumers over sixteen interleaved subjects were 5.5 times *slower* than one shared consumer, with the broker pinned and the indexers idle. + +::: + +`max_ack_pending` bounds the messages in flight and therefore the achievable throughput. Leave it unlimited (`-1`) unless there is a reason to bound the source's memory: nothing is acknowledged until a split is published, so a whole commit window is always ack-pending, and a low value stalls the pipelines waiting on commits. Under the `all` policy the window is acknowledged by a single message, so the bound has to cover the entire window rather than a sliding fraction of it. + +#### Monitoring Being durable, the consumer is observable through NATS's own monitoring (`nats consumer info`, exporters): `num_pending` is the indexing lag and `num_ack_pending` the in-flight window, both available even while the pipelines are down. The source also reports `num_pending_acks` (messages indexed but whose split is not published yet) in its observable state. @@ -261,7 +282,7 @@ EOF ## Number of pipelines -The `num_pipelines` parameter is only available for distributed sources like Kafka, GCP PubSub, and Pulsar. +The `num_pipelines` parameter is only available for distributed sources like Kafka, GCP PubSub, NATS, and Pulsar. It defines the number of pipelines to run on a cluster for the source. The actual placement of these pipelines on the different indexer will be decided by the control plane. diff --git a/docs/ingest-data/nats.md b/docs/ingest-data/nats.md index 2dc3f78cb42..39f7d8a44d9 100644 --- a/docs/ingest-data/nats.md +++ b/docs/ingest-data/nats.md @@ -131,12 +131,14 @@ The NATS source consumes a JetStream stream through a durable consumer, so messa ```bash nats stream add stackoverflow --subjects "stackoverflow.posts" --defaults -nats consumer add stackoverflow quickwit-consumer --pull --deliver all --ack explicit --wait 5m --defaults +nats consumer add stackoverflow quickwit-consumer --pull --deliver all --ack all --max-pending=-1 --wait 5m --defaults ``` :::info -The consumer must use the explicit ack policy: the source acknowledges each message once it is durably indexed, and the consumer's ack floor is the resume point. The subject filters, the deliver policy, and the ack tuning are properties of the consumer, not of the Quickwit source; the ack wait must exceed the indexing commit timeout. See the [NATS source reference](../configuration/source-config.md#nats-source) for details. +The consumer must use the `all` or the `explicit` ack policy: the source acknowledges messages once they are durably indexed, and the consumer's ack floor is the resume point. The subject filters, the deliver policy, and the ack tuning are properties of the consumer, not of the Quickwit source; the ack wait must exceed the indexing commit timeout. + +`--ack all` is used here because this source runs a single pipeline. One acknowledgment then covers every message up to it, so a commit costs one acknowledgment instead of one per message — the same shape as a Kafka offset commit, and the difference between the two is what bounds throughput on small messages. It is only safe with a single puller, so a consumer shared by several pipelines needs `--ack explicit` instead; see [Ack policy](../configuration/source-config.md#ack-policy) and [Scaling](../configuration/source-config.md#scaling) for how to grow past one pipeline. ::: diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 4612f04260a..e35e8e58d17 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -20,9 +20,14 @@ //! whoever provisioned it. //! //! Several indexing pipelines can share the consumer: NATS load-balances the -//! messages across them, so scaling is a plain `num_pipelines` update. Being -//! durable, the consumer is observable through NATS's own monitoring even -//! while the pipelines are down. +//! messages across them, so scaling is a plain `num_pipelines` update. That +//! shape requires `AckPolicy::Explicit`, which costs one acknowledgment per +//! message and is what bounds throughput on small messages; a consumer with a +//! single puller can use `AckPolicy::All` instead and spend one acknowledgment +//! per commit. See `ack_up_to` and `docs/configuration/source-config.md`. +//! +//! Being durable, the consumer is observable through NATS's own monitoring +//! even while the pipelines are down. //! //! When a message carries a W3C `traceparent` header, the source stitches the //! processing of the message into the publisher's distributed trace. @@ -95,13 +100,28 @@ pub struct NatsSource { consumer_name: String, partition_id: PartitionId, delivery_counter: u64, - /// Ack subjects of the messages delivered but not published yet, keyed by - /// delivery counter. Bounded by the consumer's `max_ack_pending`: the - /// server stops delivering when too many messages are unacknowledged. - pending_acks: BTreeMap, + /// Messages delivered but not published yet, keyed by delivery counter. + /// Bounded by the consumer's `max_ack_pending`: the server stops delivering + /// when too many messages are unacknowledged. + pending_acks: BTreeMap, + /// How the consumer wants its messages acknowledged, read off the consumer + /// at startup. `All` collapses a whole checkpoint range into a single ack, + /// `Explicit` costs one ack per message. See `ack_up_to`. + ack_policy: AckPolicy, state: NatsSourceState, } +/// A delivered message waiting to be published, and what it takes to ack it. +/// +/// The stream sequence is kept because `AckPolicy::All` acknowledges *by stream +/// sequence*: acking one message acks every message at or below its sequence. +/// Delivery order and stream order only coincide while nothing is redelivered, +/// so the sequence has to be carried rather than inferred from the map's key. +struct PendingAck { + ack_subject: Subject, + stream_sequence: u64, +} + impl fmt::Debug for NatsSource { fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { formatter @@ -131,7 +151,12 @@ impl NatsSource { "starting NATS source" ); - let (nats_client, consumer) = connect_and_fetch_consumer(&source_params).await?; + let (nats_client, consumer) = connect_and_fetch_consumer( + &source_params, + source_runtime.source_config.num_pipelines.get(), + ) + .await?; + let ack_policy = consumer.cached_info().config.ack_policy; let message_stream = consumer .stream() .max_messages_per_batch(pull_max_messages_per_batch()) @@ -161,6 +186,7 @@ impl NatsSource { partition_id, delivery_counter: 0, pending_acks: BTreeMap::new(), + ack_policy, state: NatsSourceState::default(), }) } @@ -202,7 +228,13 @@ impl NatsSource { Position::offset(self.delivery_counter), ) .context("failed to record partition delta")?; - self.pending_acks.insert(self.delivery_counter, ack_subject); + self.pending_acks.insert( + self.delivery_counter, + PendingAck { + ack_subject, + stream_sequence, + }, + ); self.state.num_bytes_processed += num_bytes; self.state.num_messages_processed += 1; @@ -283,6 +315,7 @@ impl Source for NatsSource { &self.nats_client, &mut self.pending_acks, published_up_to, + self.ack_policy, )) .await; Ok(()) @@ -457,12 +490,19 @@ const MAX_CONCURRENT_ACK_REQUESTS: usize = 128; /// on a 2 000 000 document corpus), which is why they are requests. async fn ack_up_to( nats_client: &async_nats::Client, - pending_acks: &mut BTreeMap, + pending_acks: &mut BTreeMap, published_up_to: u64, + ack_policy: AckPolicy, ) { + if ack_policy == AckPolicy::All { + ack_all_up_to(nats_client, pending_acks, published_up_to).await; + return; + } let acks: Vec<(u64, Subject)> = pending_acks .range(..=published_up_to) - .map(|(delivery_counter, ack_subject)| (*delivery_counter, ack_subject.clone())) + .map(|(delivery_counter, pending_ack)| { + (*delivery_counter, pending_ack.ack_subject.clone()) + }) .collect(); if acks.is_empty() { return; @@ -496,9 +536,69 @@ async fn ack_up_to( debug!(num_acks, "acked published messages"); } +/// Acks the whole published range with a single request. +/// +/// `AckPolicy::All` makes the server treat one ack as an ack of every message at +/// or below that message's stream sequence, so the ack of the highest published +/// sequence stands in for the entire range. That is the same shape as a Kafka +/// offset commit, and it takes the per-message ack off the hot path: a broker +/// that serves ~400 000 messages/s of pure delivery only serves ~60 000/s once +/// every message needs its own acknowledgment. +/// +/// The hidden contract is that the consumer must have a single puller. Acking by +/// sequence also acks the messages a sibling pipeline holds but has not +/// published, so `fetch_durable_consumer` refuses this policy for a source with +/// more than one pipeline. +/// +/// Nothing is dropped from `pending_acks` unless the server confirmed the ack, +/// so a failure leaves the range pending and the next `suggest_truncate` retries +/// it. `is_drained` therefore still reports the truth. +async fn ack_all_up_to( + nats_client: &async_nats::Client, + pending_acks: &mut BTreeMap, + published_up_to: u64, +) { + // A redelivery is delivered late but carries its original, lower stream + // sequence, so delivery order and stream order can disagree. Acking the + // highest published sequence would then also ack a redelivered message that + // is still waiting to be published, and a crash would lose it. The highest + // sequence that is safe to ack is therefore the highest published one that + // is still below every unpublished one. + let first_unpublished_sequence = pending_acks + .range(published_up_to.saturating_add(1)..) + .map(|(_, pending_ack)| pending_ack.stream_sequence) + .min() + .unwrap_or(u64::MAX); + let Some((_, highest_safe_ack)) = pending_acks + .range(..=published_up_to) + .filter(|(_, pending_ack)| pending_ack.stream_sequence < first_unpublished_sequence) + .max_by_key(|(_, pending_ack)| pending_ack.stream_sequence) + else { + return; + }; + let ack_subject = highest_safe_ack.ack_subject.clone(); + let stream_sequence = highest_safe_ack.stream_sequence; + if let Err(error) = nats_client.request(ack_subject, AckKind::Ack.into()).await { + warn!(%error, stream_sequence, "failed to ack NATS messages up to a stream sequence"); + return; + } + // Only the messages the server actually acked are dropped. A published + // message above `stream_sequence` stays pending and is covered by the next + // call, which keeps `is_drained` honest either way. + let num_pending_before = pending_acks.len(); + pending_acks.retain(|delivery_counter, pending_ack| { + *delivery_counter > published_up_to || pending_ack.stream_sequence > stream_sequence + }); + debug!( + num_acked = num_pending_before - pending_acks.len(), + stream_sequence, "acked published messages up to a stream sequence" + ); +} + async fn fetch_durable_consumer( jetstream_stream: &jetstream::stream::Stream, consumer_name: &str, + num_pipelines: usize, ) -> anyhow::Result { let consumer: PullConsumer = jetstream_stream .get_consumer(consumer_name) @@ -506,8 +606,19 @@ async fn fetch_durable_consumer( .map_err(|error| anyhow!("failed to find NATS consumer `{consumer_name}`: {error}"))?; let ack_policy = consumer.cached_info().config.ack_policy; ensure!( - ack_policy == AckPolicy::Explicit, - "NATS consumer `{consumer_name}` must use the explicit ack policy, got `{ack_policy:?}`" + matches!(ack_policy, AckPolicy::Explicit | AckPolicy::All), + "NATS consumer `{consumer_name}` must use the explicit or all ack policy, got \ + `{ack_policy:?}`" + ); + // `AckPolicy::All` acks by stream sequence, so one pipeline's ack would also + // ack the lower-sequenced messages its siblings are still indexing, and a + // crash would lose them. Sharing a consumer requires per-message acks. + ensure!( + ack_policy == AckPolicy::Explicit || num_pipelines == 1, + "NATS consumer `{consumer_name}` uses the all ack policy, which acknowledges every \ + message up to the acked one and is only safe with a single puller, but the source is \ + configured with {num_pipelines} pipelines: use the explicit ack policy, or one consumer \ + per pipeline" ); Ok(consumer) } @@ -592,6 +703,7 @@ async fn connect_nats(params: &NatsSourceParams) -> anyhow::Result anyhow::Result<(async_nats::Client, PullConsumer)> { let nats_client = connect_nats(params).await?; let jetstream_ctx = jetstream::new(nats_client.clone()); @@ -599,12 +711,15 @@ async fn connect_and_fetch_consumer( .get_stream(¶ms.stream) .await .with_context(|| format!("failed to find NATS JetStream stream `{}`", params.stream))?; - let consumer = fetch_durable_consumer(&jetstream_stream, ¶ms.consumer).await?; + let consumer = + fetch_durable_consumer(&jetstream_stream, ¶ms.consumer, num_pipelines).await?; Ok((nats_client, consumer)) } pub(crate) async fn check_connectivity(params: &NatsSourceParams) -> anyhow::Result<()> { - connect_and_fetch_consumer(params).await?; + // Connectivity only: a single pipeline is the permissive case, so this never + // rejects a consumer the source itself would accept. + connect_and_fetch_consumer(params, 1).await?; Ok(()) } From 67839a1bbeae1fd158fe45c52743545b77e85518 Mon Sep 17 00:00:00 2001 From: tprevot Date: Tue, 8 Sep 2026 18:19:06 +0200 Subject: [PATCH 10/21] fix: comments --- .../src/source/nats_source.rs | 67 ++++++------------- 1 file changed, 20 insertions(+), 47 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index e35e8e58d17..dd86b7cd150 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -47,6 +47,7 @@ use async_nats::jetstream::consumer::{AckPolicy, PullConsumer}; use async_nats::jetstream::message::AckKind; use async_nats::{ConnectOptions, Subject, jetstream}; use async_trait::async_trait; +use bytesize::ByteSize; use futures::{FutureExt, StreamExt}; use opentelemetry::propagation::Extractor; use opentelemetry::trace::TraceContextExt; @@ -160,11 +161,9 @@ impl NatsSource { let message_stream = consumer .stream() .max_messages_per_batch(pull_max_messages_per_batch()) - .max_bytes_per_batch(pull_max_bytes_per_batch()) - // `stream()` defaults the heartbeat to zero, i.e. disabled, unlike - // `messages()`. It has to be set back, or a stalled subscription is - // never reported and `is_transient_stream_error` has nothing to act - // on. + .max_bytes_per_batch(pull_max_bytes_per_batch().0 as usize) + // default heartbeat that is disabled when using manual builder, + // required for transient conn errors. .heartbeat(PULL_IDLE_HEARTBEAT) .expires(PULL_EXPIRES) .messages() @@ -245,6 +244,7 @@ impl NatsSource { #[async_trait] impl Source for NatsSource { + #[tracing::instrument(skip(source_sink, ctx))] async fn emit_batches( &mut self, source_sink: &SourceSink, @@ -300,6 +300,7 @@ impl Source for NatsSource { Ok(wait_before_next_batch) } + #[tracing::instrument(skip(checkpoint, ctx))] async fn suggest_truncate( &mut self, checkpoint: SourceCheckpoint, @@ -353,7 +354,9 @@ impl Source for NatsSource { } }; let Some(ack_subject) = message.message.reply else { - warn!("prefetched NATS message carries no reply subject to negatively acknowledge it on"); + warn!( + "prefetched NATS message carries no reply subject to negatively acknowledge it on" + ); continue; }; if let Err(error) = self @@ -366,11 +369,7 @@ impl Source for NatsSource { } num_naks += 1; } - // The negative acknowledgments above are plain publishes, and `drain` - // returns as soon as its command is queued rather than once the - // connection has been written out. Without this flush they can be lost - // with the connection, and their messages then wait a full `ack_wait` - // instead of being redelivered to a surviving pipeline immediately. + if let Err(error) = self.nats_client.flush().await { warn!(%error, num_naks, "failed to flush NATS negative acknowledgments"); } @@ -397,7 +396,6 @@ impl Source for NatsSource { /// Bounds the concurrent server-confirmed ack requests: enough to hide the /// round-trip latency without flooding the connection. -/// Caps a pull batch in bytes, which the message-count cap alone cannot do. /// /// The server buffers a pull response on the connection's outbound queue, and a /// connection that exceeds `max_pending` (64 MiB by default) is declared a slow @@ -408,18 +406,12 @@ impl Source for NatsSource { /// /// The value is bounded on both sides. It must stay well below the server's /// `max_pending`, since several pull requests can be outstanding at once, and it -/// must stay *above* the server's `max_payload`: the server answers a pull whose -/// `max_bytes` is smaller than the next message with `409 Message Size Exceeds -/// MaxBytes` and delivers nothing, which turns an intermittent stall into a -/// permanent one. 10 MiB clears the 8 MiB `max_payload` ceiling a server can be -/// configured with in practice, and leaves a six-fold margin against the default -/// `max_pending`. -/// Overridable while the value is being tuned against real message sizes. -const DEFAULT_PULL_MAX_BYTES_PER_BATCH: usize = 10 * 1024 * 1024; - -fn pull_max_bytes_per_batch() -> usize { +/// must stay *above* the server's `max_payload`. +const DEFAULT_PULL_MAX_BYTES_PER_BATCH: ByteSize = ByteSize::mib(10); + +fn pull_max_bytes_per_batch() -> ByteSize { quickwit_common::get_from_env_cached!( - usize, + ByteSize, "QW_NATS_PULL_MAX_BYTES_PER_BATCH", DEFAULT_PULL_MAX_BYTES_PER_BATCH, false @@ -442,8 +434,8 @@ fn pull_max_messages_per_batch() -> usize { ) } -/// Both kept at the `Consumer::messages` defaults, so switching to the builder -/// changes the batch bounds and nothing else. +// Both kept at the `Consumer::messages` defaults, so switching to the builder +// changes the batch bounds and nothing else. const PULL_IDLE_HEARTBEAT: Duration = Duration::from_secs(15); const PULL_EXPIRES: Duration = Duration::from_secs(30); @@ -478,16 +470,6 @@ const MAX_CONCURRENT_ACK_REQUESTS: usize = 128; /// Sends server-confirmed acknowledgments ("double acks") for the messages up /// to `published_up_to`. -/// -/// The confirmation is not for error reporting: it is what holds a message in -/// `pending_acks` long enough for [`Source::is_drained`] to still be false when -/// `Drain` arrives. `Drain` queues its force-commit batch and then checks -/// `is_drained` immediately, so a source that reports itself drained at that -/// instant exits before the in-flight documents are published, and they are -/// redelivered as duplicates. Replacing these requests with plain publishes was -/// measured at 37 % less broker CPU and 40 % more throughput at four pipelines, -/// and it broke exactly-once on `num_pipelines` reassignment (+91 581 duplicates -/// on a 2 000 000 document corpus), which is why they are requests. async fn ack_up_to( nats_client: &async_nats::Client, pending_acks: &mut BTreeMap, @@ -500,9 +482,7 @@ async fn ack_up_to( } let acks: Vec<(u64, Subject)> = pending_acks .range(..=published_up_to) - .map(|(delivery_counter, pending_ack)| { - (*delivery_counter, pending_ack.ack_subject.clone()) - }) + .map(|(delivery_counter, pending_ack)| (*delivery_counter, pending_ack.ack_subject.clone())) .collect(); if acks.is_empty() { return; @@ -540,15 +520,8 @@ async fn ack_up_to( /// /// `AckPolicy::All` makes the server treat one ack as an ack of every message at /// or below that message's stream sequence, so the ack of the highest published -/// sequence stands in for the entire range. That is the same shape as a Kafka -/// offset commit, and it takes the per-message ack off the hot path: a broker -/// that serves ~400 000 messages/s of pure delivery only serves ~60 000/s once -/// every message needs its own acknowledgment. -/// -/// The hidden contract is that the consumer must have a single puller. Acking by -/// sequence also acks the messages a sibling pipeline holds but has not -/// published, so `fetch_durable_consumer` refuses this policy for a source with -/// more than one pipeline. +/// sequence stands in for the entire range. It requires that num_pipeline = 1, +/// otherwise this is unsound. /// /// Nothing is dropped from `pending_acks` unless the server confirmed the ack, /// so a failure leaves the range pending and the next `suggest_truncate` retries From 96efbf1a2459c96d9e7b5cd135b07404a6af40f1 Mon Sep 17 00:00:00 2001 From: tprevot Date: Tue, 8 Sep 2026 18:19:17 +0200 Subject: [PATCH 11/21] fix: comments --- quickwit/quickwit-indexing/src/source/nats_source.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index dd86b7cd150..03ac29d2868 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -355,7 +355,8 @@ impl Source for NatsSource { }; let Some(ack_subject) = message.message.reply else { warn!( - "prefetched NATS message carries no reply subject to negatively acknowledge it on" + "prefetched NATS message carries no reply subject to negatively acknowledge \ + it on" ); continue; }; From fcfcca902c01852e778e74b05ea66680fb527d9c Mon Sep 17 00:00:00 2001 From: tprevot Date: Tue, 8 Sep 2026 18:32:26 +0200 Subject: [PATCH 12/21] fix: remove the ack policy all --- .../src/source/nats_source.rs | 213 +++++------------- 1 file changed, 54 insertions(+), 159 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 03ac29d2868..c0b0dbb7b14 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -22,9 +22,7 @@ //! Several indexing pipelines can share the consumer: NATS load-balances the //! messages across them, so scaling is a plain `num_pipelines` update. That //! shape requires `AckPolicy::Explicit`, which costs one acknowledgment per -//! message and is what bounds throughput on small messages; a consumer with a -//! single puller can use `AckPolicy::All` instead and spend one acknowledgment -//! per commit. See `ack_up_to` and `docs/configuration/source-config.md`. +//! message and is what bounds throughput on small messages. //! //! Being durable, the consumer is observable through NATS's own monitoring //! even while the pipelines are down. @@ -104,23 +102,9 @@ pub struct NatsSource { /// Messages delivered but not published yet, keyed by delivery counter. /// Bounded by the consumer's `max_ack_pending`: the server stops delivering /// when too many messages are unacknowledged. - pending_acks: BTreeMap, - /// How the consumer wants its messages acknowledged, read off the consumer - /// at startup. `All` collapses a whole checkpoint range into a single ack, - /// `Explicit` costs one ack per message. See `ack_up_to`. - ack_policy: AckPolicy, - state: NatsSourceState, -} + pending_acks: BTreeMap, -/// A delivered message waiting to be published, and what it takes to ack it. -/// -/// The stream sequence is kept because `AckPolicy::All` acknowledges *by stream -/// sequence*: acking one message acks every message at or below its sequence. -/// Delivery order and stream order only coincide while nothing is redelivered, -/// so the sequence has to be carried rather than inferred from the map's key. -struct PendingAck { - ack_subject: Subject, - stream_sequence: u64, + state: NatsSourceState, } impl fmt::Debug for NatsSource { @@ -152,12 +136,7 @@ impl NatsSource { "starting NATS source" ); - let (nats_client, consumer) = connect_and_fetch_consumer( - &source_params, - source_runtime.source_config.num_pipelines.get(), - ) - .await?; - let ack_policy = consumer.cached_info().config.ack_policy; + let (nats_client, consumer) = connect_and_fetch_consumer(&source_params).await?; let message_stream = consumer .stream() .max_messages_per_batch(pull_max_messages_per_batch()) @@ -185,7 +164,6 @@ impl NatsSource { partition_id, delivery_counter: 0, pending_acks: BTreeMap::new(), - ack_policy, state: NatsSourceState::default(), }) } @@ -227,19 +205,57 @@ impl NatsSource { Position::offset(self.delivery_counter), ) .context("failed to record partition delta")?; - self.pending_acks.insert( - self.delivery_counter, - PendingAck { - ack_subject, - stream_sequence, - }, - ); + self.pending_acks.insert(self.delivery_counter, ack_subject); self.state.num_bytes_processed += num_bytes; self.state.num_messages_processed += 1; Ok(()) } + + /// Sends server-confirmed acknowledgments ("double acks") for the messages up + /// to `published_up_to`. + async fn ack_up_to(&mut self, published_up_to: u64) { + let acks: Vec<(u64, Subject)> = self + .pending_acks + .range(..=published_up_to) + .map(|(delivery_counter, pending_ack)| (*delivery_counter, pending_ack.clone())) + .collect(); + if acks.is_empty() { + return; + } + let mut ack_results = futures::stream::iter(acks) + .map(async |(delivery_counter, ack_subject)| { + // Mirrors `jetstream::Message::double_ack()`, but through the + // client's muxed inbox rather than a subscription per message. + let ack_result = self + .nats_client + .request(ack_subject, AckKind::Ack.into()) + .await; + (delivery_counter, ack_result) + }) + .buffer_unordered(MAX_CONCURRENT_ACK_REQUESTS); + + let mut num_acks = 0usize; + let mut num_failed_acks = 0usize; + let mut last_ack_error = None; + while let Some((delivery_counter, ack_result)) = ack_results.next().await { + match ack_result { + Ok(_ack_reply) => { + self.pending_acks.remove(&delivery_counter); + num_acks += 1; + } + Err(error) => { + num_failed_acks += 1; + last_ack_error = Some(error); + } + } + } + if let Some(error) = last_ack_error { + warn!(%error, num_failed_acks, "failed to ack NATS messages"); + } + debug!(num_acks, "acked published messages"); + } } #[async_trait] @@ -312,13 +328,7 @@ impl Source for NatsSource { let Some(published_up_to) = position.as_u64() else { return Ok(()); }; - ctx.protect_future(ack_up_to( - &self.nats_client, - &mut self.pending_acks, - published_up_to, - self.ack_policy, - )) - .await; + ctx.protect_future(self.ack_up_to(published_up_to)).await; Ok(()) } @@ -469,110 +479,9 @@ fn is_transient_stream_error(error: &MessagesError) -> bool { /// messages: 1024 is indistinguishable from 128, and 8192 is 20 % slower. const MAX_CONCURRENT_ACK_REQUESTS: usize = 128; -/// Sends server-confirmed acknowledgments ("double acks") for the messages up -/// to `published_up_to`. -async fn ack_up_to( - nats_client: &async_nats::Client, - pending_acks: &mut BTreeMap, - published_up_to: u64, - ack_policy: AckPolicy, -) { - if ack_policy == AckPolicy::All { - ack_all_up_to(nats_client, pending_acks, published_up_to).await; - return; - } - let acks: Vec<(u64, Subject)> = pending_acks - .range(..=published_up_to) - .map(|(delivery_counter, pending_ack)| (*delivery_counter, pending_ack.ack_subject.clone())) - .collect(); - if acks.is_empty() { - return; - } - let mut ack_results = futures::stream::iter(acks) - .map(|(delivery_counter, ack_subject)| async move { - // Mirrors `jetstream::Message::double_ack()`, but through the - // client's muxed inbox rather than a subscription per message. - let ack_result = nats_client.request(ack_subject, AckKind::Ack.into()).await; - (delivery_counter, ack_result) - }) - .buffer_unordered(MAX_CONCURRENT_ACK_REQUESTS); - let mut num_acks = 0usize; - let mut num_failed_acks = 0usize; - let mut last_ack_error = None; - while let Some((delivery_counter, ack_result)) = ack_results.next().await { - match ack_result { - Ok(_ack_reply) => { - pending_acks.remove(&delivery_counter); - num_acks += 1; - } - Err(error) => { - num_failed_acks += 1; - last_ack_error = Some(error); - } - } - } - if let Some(error) = last_ack_error { - warn!(%error, num_failed_acks, "failed to ack NATS messages"); - } - debug!(num_acks, "acked published messages"); -} - -/// Acks the whole published range with a single request. -/// -/// `AckPolicy::All` makes the server treat one ack as an ack of every message at -/// or below that message's stream sequence, so the ack of the highest published -/// sequence stands in for the entire range. It requires that num_pipeline = 1, -/// otherwise this is unsound. -/// -/// Nothing is dropped from `pending_acks` unless the server confirmed the ack, -/// so a failure leaves the range pending and the next `suggest_truncate` retries -/// it. `is_drained` therefore still reports the truth. -async fn ack_all_up_to( - nats_client: &async_nats::Client, - pending_acks: &mut BTreeMap, - published_up_to: u64, -) { - // A redelivery is delivered late but carries its original, lower stream - // sequence, so delivery order and stream order can disagree. Acking the - // highest published sequence would then also ack a redelivered message that - // is still waiting to be published, and a crash would lose it. The highest - // sequence that is safe to ack is therefore the highest published one that - // is still below every unpublished one. - let first_unpublished_sequence = pending_acks - .range(published_up_to.saturating_add(1)..) - .map(|(_, pending_ack)| pending_ack.stream_sequence) - .min() - .unwrap_or(u64::MAX); - let Some((_, highest_safe_ack)) = pending_acks - .range(..=published_up_to) - .filter(|(_, pending_ack)| pending_ack.stream_sequence < first_unpublished_sequence) - .max_by_key(|(_, pending_ack)| pending_ack.stream_sequence) - else { - return; - }; - let ack_subject = highest_safe_ack.ack_subject.clone(); - let stream_sequence = highest_safe_ack.stream_sequence; - if let Err(error) = nats_client.request(ack_subject, AckKind::Ack.into()).await { - warn!(%error, stream_sequence, "failed to ack NATS messages up to a stream sequence"); - return; - } - // Only the messages the server actually acked are dropped. A published - // message above `stream_sequence` stays pending and is covered by the next - // call, which keeps `is_drained` honest either way. - let num_pending_before = pending_acks.len(); - pending_acks.retain(|delivery_counter, pending_ack| { - *delivery_counter > published_up_to || pending_ack.stream_sequence > stream_sequence - }); - debug!( - num_acked = num_pending_before - pending_acks.len(), - stream_sequence, "acked published messages up to a stream sequence" - ); -} - async fn fetch_durable_consumer( jetstream_stream: &jetstream::stream::Stream, consumer_name: &str, - num_pipelines: usize, ) -> anyhow::Result { let consumer: PullConsumer = jetstream_stream .get_consumer(consumer_name) @@ -580,20 +489,10 @@ async fn fetch_durable_consumer( .map_err(|error| anyhow!("failed to find NATS consumer `{consumer_name}`: {error}"))?; let ack_policy = consumer.cached_info().config.ack_policy; ensure!( - matches!(ack_policy, AckPolicy::Explicit | AckPolicy::All), - "NATS consumer `{consumer_name}` must use the explicit or all ack policy, got \ + ack_policy == AckPolicy::Explicit, + "NATS consumer `{consumer_name}` must use the explicit ack policy, got \ `{ack_policy:?}`" ); - // `AckPolicy::All` acks by stream sequence, so one pipeline's ack would also - // ack the lower-sequenced messages its siblings are still indexing, and a - // crash would lose them. Sharing a consumer requires per-message acks. - ensure!( - ack_policy == AckPolicy::Explicit || num_pipelines == 1, - "NATS consumer `{consumer_name}` uses the all ack policy, which acknowledges every \ - message up to the acked one and is only safe with a single puller, but the source is \ - configured with {num_pipelines} pipelines: use the explicit ack policy, or one consumer \ - per pipeline" - ); Ok(consumer) } @@ -677,7 +576,6 @@ async fn connect_nats(params: &NatsSourceParams) -> anyhow::Result anyhow::Result<(async_nats::Client, PullConsumer)> { let nats_client = connect_nats(params).await?; let jetstream_ctx = jetstream::new(nats_client.clone()); @@ -685,15 +583,12 @@ async fn connect_and_fetch_consumer( .get_stream(¶ms.stream) .await .with_context(|| format!("failed to find NATS JetStream stream `{}`", params.stream))?; - let consumer = - fetch_durable_consumer(&jetstream_stream, ¶ms.consumer, num_pipelines).await?; + let consumer = fetch_durable_consumer(&jetstream_stream, ¶ms.consumer).await?; Ok((nats_client, consumer)) } pub(crate) async fn check_connectivity(params: &NatsSourceParams) -> anyhow::Result<()> { - // Connectivity only: a single pipeline is the permissive case, so this never - // rejects a consumer the source itself would accept. - connect_and_fetch_consumer(params, 1).await?; + connect_and_fetch_consumer(params).await?; Ok(()) } From 55b5e53ff2cda5a7fc2cb2276662c5c601578f5f Mon Sep 17 00:00:00 2001 From: tprevot Date: Tue, 8 Sep 2026 18:58:47 +0200 Subject: [PATCH 13/21] fix: remove the ack policy all --- quickwit/quickwit-indexing/src/source/nats_source.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index c0b0dbb7b14..f79c8ba68c3 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -24,9 +24,6 @@ //! shape requires `AckPolicy::Explicit`, which costs one acknowledgment per //! message and is what bounds throughput on small messages. //! -//! Being durable, the consumer is observable through NATS's own monitoring -//! even while the pipelines are down. -//! //! When a message carries a W3C `traceparent` header, the source stitches the //! processing of the message into the publisher's distributed trace. @@ -490,8 +487,7 @@ async fn fetch_durable_consumer( let ack_policy = consumer.cached_info().config.ack_policy; ensure!( ack_policy == AckPolicy::Explicit, - "NATS consumer `{consumer_name}` must use the explicit ack policy, got \ - `{ack_policy:?}`" + "NATS consumer `{consumer_name}` must use the explicit ack policy, got `{ack_policy:?}`" ); Ok(consumer) } From 3804790581b0606787111d6a114e61f60fc5a80d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9odore=20Pr=C3=A9vot?= Date: Wed, 9 Sep 2026 09:26:31 +0000 Subject: [PATCH 14/21] feat: nats source --- docs/configuration/source-config.md | 19 ++++++++----------- docs/ingest-data/nats.md | 6 +++--- 2 files changed, 11 insertions(+), 14 deletions(-) diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index a9dd816ba10..7351e35927d 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -189,34 +189,31 @@ A NATS source reads data from a [NATS JetStream](https://docs.nats.io/nats-conce A tutorial is available [here](/docs/ingest-data/nats.md). -The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The consumer must use either the **explicit** or the **all** ack policy: the source acknowledges messages once the split containing them is published, and the consumer's ack floor is the resume point. Which of the two to pick is a throughput decision, covered in [Ack policy](#ack-policy) below. +The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The consumer must use the **explicit ack policy**: the source acknowledges each message once the split containing it is published, and the consumer's ack floor is the resume point. Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. `ack_wait` must exceed the end-to-end publish latency (roughly the commit timeout plus the upload time), otherwise messages are redelivered while they are still being indexed. -#### Ack policy +#### Acknowledgment cost -The ack policy decides how many acknowledgments a corpus costs, and at small message sizes that is what bounds the source's throughput. JetStream delivery is cheap; per-message acknowledgment is not. +The source acknowledges every message individually and waits for the server to confirm each one — that confirmation is what tells the drain the pipeline is empty — so indexing a corpus costs one confirmed acknowledgment per document. At small message sizes that, rather than the indexer, is what bounds throughput. Measured on a 16 vCPU host across sixteen per-tenant consumers of 1 KiB documents: 72 MiB/s, with the broker spending 5.1 of its cores on the acknowledgment path against Quickwit's 4.7 on indexing, roughly ten times the broker CPU per document that an equivalent Kafka source spends on periodic offset commits. -- **`explicit`** — one acknowledgment per message. Required whenever several pipelines share the consumer, because acknowledging a message under the `all` policy would also acknowledge the messages a sibling pipeline is still indexing. -- **`all`** — one acknowledgment covers every message up to it, so the source sends a single acknowledgment per commit instead of one per message, the same shape as a Kafka offset commit. Quickwit rejects a source configured with `num_pipelines` greater than 1 on such a consumer, since it is only sound when the consumer has a **single puller**. - -Prefer `all`, and get parallelism from partitioning rather than from several pipelines on one consumer. On a 16 vCPU host indexing 1 KiB messages, moving one pipeline from `explicit` to `all` cut broker CPU by 56 % and total CPU per indexed GiB by 31 %; the gain is larger the smaller the messages, and negligible above a few hundred KiB, where the message rate is low enough that acknowledgments no longer matter. +The cost is per *message*, so it amortises as messages get larger and is negligible above a few hundred KiB, where the message rate is only a few hundred per second. Batching many log records into one message and letting an `otlp_*` input format expand them is the shape that gets that benefit. #### Scaling There are two ways to spread a stream over several pipelines, and they do not perform alike. -**One stream per pipeline** (recommended). Partition the data into one stream per pipeline — the equivalent of Kafka partitions — give each stream its own consumer, and add one Quickwit source with `num_pipelines: 1` per stream. Each consumer then has a single puller, which is what makes the `all` ack policy usable, and no consumer reads a message it does not own. This is the only shape that scales: on the host above, sixteen streams indexed the same 1 KiB corpus 2.7 times faster than sixteen pipelines sharing one consumer. +**One consumer per partition of the data** (recommended). Give each partition — a tenant, a subject, whatever the natural unit is — its own durable consumer with a single `filter_subject`, and add one Quickwit source with `num_pipelines: 1` per consumer. The partitions then scale and fail independently, and each gets its own checkpoint. Splitting the same corpus over one shared stream with a consumer per partition and over a stream per partition measured 21.28 s and 21.26 s, so partition whichever way suits the producers. Adding pipelines *within* a consumer does not help: sixteen consumers with two pipelines each measured 23.5 s against 21.3 s with one. -**One shared consumer** (simplest). Several pipelines bind to the same consumer and NATS load-balances the messages across them, so scaling is a plain `num_pipelines` update and the control plane places the pipelines across the indexers of the cluster. This requires the `explicit` policy, and the resulting per-message acknowledgments cap the aggregate throughput regardless of how many pipelines are added — adding pipelines beyond a handful buys close to nothing at small message sizes. +**One shared consumer** (simplest). Several pipelines bind to the same consumer and NATS load-balances the messages across them, so scaling is a plain `num_pipelines` update and the control plane places the pipelines across the indexers of the cluster. Because the acknowledgment cost is per message rather than per pipeline, adding pipelines to one consumer does not raise the aggregate throughput much at small message sizes: on the host above it was worth 1.15× between one and sixteen pipelines, against 1.4× for splitting the same work over sixteen consumers. :::warning -Do not partition a single stream by giving each pipeline a subject-filtered consumer. Unless a subject's messages happen to be contiguous in the stream, every consumer scans past the messages it does not own, and the broker saturates: sixteen subject-filtered consumers over sixteen interleaved subjects were 5.5 times *slower* than one shared consumer, with the broker pinned and the indexers idle. +Give each consumer **one** subject filter, not several. A consumer configured with `filter_subjects` (NATS 2.10's multi-filter) uses a different server path than one configured with a single `filter_subject`, and on a stream whose subjects are interleaved it is dramatically slower: sixteen multi-filter consumers over sixteen interleaved subjects took 7.5 times longer than the same sixteen consumers with one filter each, with the broker saturated and the indexers idle. ::: -`max_ack_pending` bounds the messages in flight and therefore the achievable throughput. Leave it unlimited (`-1`) unless there is a reason to bound the source's memory: nothing is acknowledged until a split is published, so a whole commit window is always ack-pending, and a low value stalls the pipelines waiting on commits. Under the `all` policy the window is acknowledged by a single message, so the bound has to cover the entire window rather than a sliding fraction of it. +`max_ack_pending` bounds the messages in flight and therefore the achievable throughput: nothing is acknowledged until a split is published, so a whole commit window is always ack-pending and the achievable rate is `max_ack_pending / commit_timeout` documents per second. Measured at four pipelines with a 10 s commit timeout, 1,000 gave 97 documents per second and 20,000 gave 2,115. Leave it unlimited (`-1`) unless there is a reason to bound the source's memory, and note the trade: unlimited also makes the replay window after a crash the whole delivered span above the ack floor rather than a bounded slice of it. #### Monitoring diff --git a/docs/ingest-data/nats.md b/docs/ingest-data/nats.md index 39f7d8a44d9..3c63c766f92 100644 --- a/docs/ingest-data/nats.md +++ b/docs/ingest-data/nats.md @@ -131,14 +131,14 @@ The NATS source consumes a JetStream stream through a durable consumer, so messa ```bash nats stream add stackoverflow --subjects "stackoverflow.posts" --defaults -nats consumer add stackoverflow quickwit-consumer --pull --deliver all --ack all --max-pending=-1 --wait 5m --defaults +nats consumer add stackoverflow quickwit-consumer --pull --deliver all --ack explicit --max-pending=-1 --wait 5m --defaults ``` :::info -The consumer must use the `all` or the `explicit` ack policy: the source acknowledges messages once they are durably indexed, and the consumer's ack floor is the resume point. The subject filters, the deliver policy, and the ack tuning are properties of the consumer, not of the Quickwit source; the ack wait must exceed the indexing commit timeout. +The consumer must use the explicit ack policy: the source acknowledges each message once it is durably indexed, and the consumer's ack floor is the resume point. The subject filters, the deliver policy, and the ack tuning are properties of the consumer, not of the Quickwit source; the ack wait must exceed the indexing commit timeout. -`--ack all` is used here because this source runs a single pipeline. One acknowledgment then covers every message up to it, so a commit costs one acknowledgment instead of one per message — the same shape as a Kafka offset commit, and the difference between the two is what bounds throughput on small messages. It is only safe with a single puller, so a consumer shared by several pipelines needs `--ack explicit` instead; see [Ack policy](../configuration/source-config.md#ack-policy) and [Scaling](../configuration/source-config.md#scaling) for how to grow past one pipeline. +`--max-pending=-1` leaves the in-flight window unbounded: the source acknowledges only once a split is published, so a whole commit window is always ack-pending and a small bound throttles indexing. See the [NATS source reference](../configuration/source-config.md#nats-source) for the rest, and [Scaling](../configuration/source-config.md#scaling) for how to grow past one pipeline. ::: From 2667e4c9de19ea4024c21c4f6dbf92008f01f920 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9odore=20Pr=C3=A9vot?= Date: Wed, 9 Sep 2026 12:14:37 +0000 Subject: [PATCH 15/21] feat: nats docs --- docs/configuration/source-config.md | 77 ++++++++++++++++++++++++----- docs/ingest-data/nats.md | 13 +++-- 2 files changed, 75 insertions(+), 15 deletions(-) diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 7351e35927d..1ff6a607f53 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -189,31 +189,84 @@ A NATS source reads data from a [NATS JetStream](https://docs.nats.io/nats-conce A tutorial is available [here](/docs/ingest-data/nats.md). -The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning (`ack_wait`, `max_ack_pending`) belong to whoever provisioned it. The consumer must use the **explicit ack policy**: the source acknowledges each message once the split containing it is published, and the consumer's ack floor is the resume point. +The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning belong to whoever provisioned it. -Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. `ack_wait` must exceed the end-to-end publish latency (roughly the commit timeout plus the upload time), otherwise messages are redelivered while they are still being indexed. +Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. -#### Acknowledgment cost +#### Consumer invariants + +These are properties of the consumer, not of the source. Only the ack policy is enforced when the source is created; the rest are not checked, and getting them wrong looks like Quickwit being slow or duplicating rather than like a consumer misconfiguration. + +**`ack_policy` must be `explicit`.** The source acknowledges each message individually once the split containing it is published, and waits for the server to confirm the acknowledgment — the confirmation is what tells the drain that the pipeline is empty. The consumer's ack floor is therefore the resume point, and it is the only progress state that matters. Any other policy is rejected when the source is created. + +**`ack_wait` must exceed the end-to-end publish latency.** The timer starts at delivery and has to outlast three terms: + +``` +ack_wait > commit_timeout + split upload and publish + ack round trip +``` + +Shorter than their sum, and NATS redelivers messages that are still being indexed; they are then indexed twice. 5 minutes with a 60 s `commit_timeout` leaves a wide margin. + +It is also the *recovery* time from a lost delivery, so it should not be arbitrarily large either. If a message is delivered but never arrives — a dropped connection, a slow-consumer disconnect — nothing brings it back until the timer expires and the pipeline simply idles. The same injected fault idled a run for 278 s at `ack_wait=300s` and 23.2 s at `30s`. + +**`max_ack_pending` should be `-1`.** Nothing is acknowledged until a split is published, so a whole commit window is always ack-pending. That makes the setting a hard throughput cap rather than a safety valve: + +``` +achievable rate ≈ max_ack_pending / commit_timeout documents per second +``` + +Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents per second where the formula predicts 100, and `20000` gave 2,115 where it predicts 2,000. A bound that looks generous for an ordinary queue consumer throttles indexing to a trickle here. If the in-flight window has to be bounded, size it above `throughput × commit_timeout` rather than at a small absolute number, and note the trade: unlimited also makes the crash-replay window the whole delivered span above the ack floor rather than a bounded slice. -The source acknowledges every message individually and waits for the server to confirm each one — that confirmation is what tells the drain the pipeline is empty — so indexing a corpus costs one confirmed acknowledgment per document. At small message sizes that, rather than the indexer, is what bounds throughput. Measured on a 16 vCPU host across sixteen per-tenant consumers of 1 KiB documents: 72 MiB/s, with the broker spending 5.1 of its cores on the acknowledgment path against Quickwit's 4.7 on indexing, roughly ten times the broker CPU per document that an equivalent Kafka source spends on periodic offset commits. +**`deliver_policy`** is usually `all`, so a new consumer indexes the stream from the start. -The cost is per *message*, so it amortises as messages get larger and is negligible above a few hundred KiB, where the message rate is only a few hundred per second. Batching many log records into one message and letting an `otlp_*` input format expand them is the shape that gets that benefit. +**Give each consumer a single `filter_subject`.** A consumer configured with several filters through NATS 2.10's multi-filter `filter_subjects` appears to take a much slower path in the broker: on a stream whose subjects are interleaved, sixteen multi-filter consumers took 7.5 times longer than the same sixteen consumers with one filter each, with the broker saturated and the indexers idle. + +#### What can be tuned on the Quickwit side + +| Setting | Where | Effect | +| --- | --- | --- | +| `num_pipelines` | source config | Pipelines bound to the consumer. Scales cleanly on large messages: 52 → 147 MiB/s from one to four pipelines at 512 KiB. On small messages the acknowledgment path caps the aggregate, and it is worth only about 1.15× from one to sixteen — there, parallelism has to come from more consumers. | +| [`commit_timeout_secs`](index-config.md#indexing-settings) | index config | Sets the size of the always-ack-pending window, so it interacts with `max_ack_pending` and with `ack_wait`. It is also the length of the cooperative-indexing cycle. | +| [`docstore_compression_level`](index-config.md#indexing-settings) | index config | Defaults to zstd 8, which is what limits a single pipeline on small documents: the dedicated docstore thread saturates and back-pressures the indexer. Level 1 or 3 measured 21.1 s against 31.2 s on a 1.5 GiB corpus of 1 KiB documents. Not NATS-specific. | +| [`enable_cooperative_indexing`](node-config.md#indexer-configuration) | node config | See below. | +| `QW_NATS_PULL_MAX_BYTES_PER_BATCH` | env, default 10 MiB | Bounds the bytes the server may push per pull request. It must stay well under the server's per-connection `max_pending` (64 MiB by default), or the server declares a slow consumer and closes the connection; the messages are already counted as delivered, so the pipeline then idles for a full `ack_wait`. 10 MiB and 20 MiB measure identically. | +| `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` | env, default 100,000 | Messages per pull request. Does not bind at these message sizes — 200 and 100,000 measured identically — because the byte cap is what bounds a batch. | #### Scaling -There are two ways to spread a stream over several pipelines, and they do not perform alike. +**One consumer per partition of the data** (recommended). Give each partition — a tenant, a subject, whatever the natural unit is — its own durable consumer with a single `filter_subject`, and add one Quickwit source with `num_pipelines: 1` per consumer. The partitions then scale and fail independently, and each gets its own checkpoint. -**One consumer per partition of the data** (recommended). Give each partition — a tenant, a subject, whatever the natural unit is — its own durable consumer with a single `filter_subject`, and add one Quickwit source with `num_pipelines: 1` per consumer. The partitions then scale and fail independently, and each gets its own checkpoint. Splitting the same corpus over one shared stream with a consumer per partition and over a stream per partition measured 21.28 s and 21.26 s, so partition whichever way suits the producers. Adding pipelines *within* a consumer does not help: sixteen consumers with two pipelines each measured 23.5 s against 21.3 s with one. +**One shared consumer** (simplest). Several pipelines bind to the same consumer and NATS load-balances the messages across them, so scaling is a plain `num_pipelines` update and the control plane places the pipelines across the indexers of the cluster. -**One shared consumer** (simplest). Several pipelines bind to the same consumer and NATS load-balances the messages across them, so scaling is a plain `num_pipelines` update and the control plane places the pipelines across the indexers of the cluster. Because the acknowledgment cost is per message rather than per pipeline, adding pipelines to one consumer does not raise the aggregate throughput much at small message sizes: on the host above it was worth 1.15× between one and sixteen pipelines, against 1.4× for splitting the same work over sixteen consumers. +At 512 KiB per message the two shapes are equivalent — 52.1, 91.3 and 146.6 MiB/s at one, two and four units either way — and so is partitioning across separate streams instead of subjects. At 1 KiB they diverge: a shared consumer stops gaining past a couple of pipelines while per-tenant consumers keep scaling, because the acknowledgment cost is per message rather than per pipeline. The per-consumer shape is therefore the safe default: never worse, and better on small documents. -:::warning +#### Acknowledgment cost -Give each consumer **one** subject filter, not several. A consumer configured with `filter_subjects` (NATS 2.10's multi-filter) uses a different server path than one configured with a single `filter_subject`, and on a stream whose subjects are interleaved it is dramatically slower: sixteen multi-filter consumers over sixteen interleaved subjects took 7.5 times longer than the same sixteen consumers with one filter each, with the broker saturated and the indexers idle. +One confirmed acknowledgment per message is the source's dominant broker cost, and it is what decides whether the source or the indexer is the bottleneck. -::: +- At **512 KiB** per message it is invisible. The NATS and Kafka sources measured within a percent of each other across the whole pipeline sweep, and broker CPU stayed at 0.5–0.9 cores for both, which is the host's floor. Priced on the broker alone, a confirmed acknowledgment per message still leaves 1,138 MiB/s — several times what one indexer consumes. +- At **1 KiB** it is the ceiling. Sixteen per-tenant consumers reached 72 MiB/s while an equivalent Kafka source reached 150, and the broker spent 5.1 of the host's cores on acknowledgments against Quickwit's 4.7 on indexing — about ten times the broker CPU per document that a periodic offset commit costs. Priced on the broker alone, acknowledgment divides throughput by 6.5 at this size. + +The cost is per *message*, so the message rate is what matters. Batching many log records into one message and letting an `otlp_*` input format expand them into separate documents is what moves a workload from the second case to the first: one acknowledgment then covers the whole batch. + +#### Cooperative indexing + +[`enable_cooperative_indexing`](node-config.md#indexer-configuration) is a **node-level** switch, off by default. It cannot be set per index or per source: one semaphore per indexer, sized to the node's blocking-thread count, is shared by every pipeline that node runs. + +With it off, an indexer cuts a split when the commit timeout fires or the split hits its size limits, and every pipeline indexes whenever it has work. With it on, a pipeline first sleeps a phase offset derived from a hash of its pipeline id, spread over `[0, commit_timeout)`; it then takes a semaphore permit, cuts its split as soon as *its own mailbox drains* rather than waiting for the commit timeout, releases the permit, and sleeps out the rest of a `commit_timeout`-long cycle. So a bounded number of pipelines hold an `IndexWriter` at any instant, and their CPU, disk and network spikes are spread across the commit window instead of landing together. + +What it costs and buys, measured on a 4.75 GiB corpus with a 10 s commit timeout: + +| | throughput, coop off → on | first published split, off → on | +| --- | --- | --- | +| 512 KiB messages, 1 pipeline | 52.1 → 52.1 MiB/s | — | +| 512 KiB messages, 16 pipelines | 177.8 → 209.5 MiB/s | 13.2 s → 3.0 s | +| ~1 KiB messages, 1 pipeline | 33.6 → 24.2 MiB/s | — | +| ~1 KiB messages, 16 pipelines | 38.6 → 36.7 MiB/s | 13.2 s → 3.0 s | + +At 512 KiB it costs nothing measurable and still makes the first split searchable 4.4× sooner. On small documents it costs up to 28 % for a single busy pipeline, shrinking to a few percent by sixteen — the penalty is worst for one pipeline with a lot to do, which is the opposite of what the setting is for. Enable it when a node runs many pipelines that are each mostly idle, which is also when its bound on concurrent `IndexWriter`s is what keeps the node inside its memory budget. -`max_ack_pending` bounds the messages in flight and therefore the achievable throughput: nothing is acknowledged until a split is published, so a whole commit window is always ack-pending and the achievable rate is `max_ack_pending / commit_timeout` documents per second. Measured at four pipelines with a 10 s commit timeout, 1,000 gave 97 documents per second and 20,000 gave 2,115. Leave it unlimited (`-1`) unless there is a reason to bound the source's memory, and note the trade: unlimited also makes the replay window after a crash the whole delivered span above the ack floor rather than a bounded slice of it. +One caveat: the initial phase sleep is up to one full `commit_timeout`, so with a 60 s commit timeout a pipeline can sit idle for up to a minute after a restart or reassignment before it indexes anything. That is a ramp after each deploy rather than an ongoing cost; shorten the commit timeout if it matters. #### Monitoring diff --git a/docs/ingest-data/nats.md b/docs/ingest-data/nats.md index 3c63c766f92..54b669faf58 100644 --- a/docs/ingest-data/nats.md +++ b/docs/ingest-data/nats.md @@ -136,9 +136,14 @@ nats consumer add stackoverflow quickwit-consumer --pull --deliver all --ack exp :::info -The consumer must use the explicit ack policy: the source acknowledges each message once it is durably indexed, and the consumer's ack floor is the resume point. The subject filters, the deliver policy, and the ack tuning are properties of the consumer, not of the Quickwit source; the ack wait must exceed the indexing commit timeout. +The consumer must use the explicit ack policy: the source acknowledges each message once it is durably indexed, and the consumer's ack floor is the resume point. The subject filters, the deliver policy, and the ack tuning are properties of the consumer, not of the Quickwit source. -`--max-pending=-1` leaves the in-flight window unbounded: the source acknowledges only once a split is published, so a whole commit window is always ack-pending and a small bound throttles indexing. See the [NATS source reference](../configuration/source-config.md#nats-source) for the rest, and [Scaling](../configuration/source-config.md#scaling) for how to grow past one pipeline. +Two of those settings decide whether indexing works at all, and neither is checked when the source is created: + +- `--wait` (`ack_wait`) must outlast the whole path from delivery to acknowledgment — the commit timeout, plus the split upload and publish, plus the acknowledgment round trip. Shorter than that and NATS redelivers messages that are still being indexed, which duplicates them. `5m` against this tutorial's 10 s commit timeout is a wide margin. It is also the recovery time from a lost delivery, so it should not be arbitrarily large either. +- `--max-pending=-1` leaves the in-flight window unbounded. The source acknowledges only once a split is published, so a whole commit window is always ack-pending and a bounded value is a hard throughput cap of roughly `max_ack_pending / commit_timeout` documents per second. + +See [Consumer invariants](../configuration/source-config.md#consumer-invariants) for the details, and [Scaling](../configuration/source-config.md#scaling) for how to grow past one pipeline. ::: @@ -242,8 +247,10 @@ Let's delete the files and resources created for the purpose of this tutorial. ```bash # Delete quickwit index. ./quickwit index delete --index stackoverflow --yes -# Delete NATS stream. +# Delete NATS stream, and with it the consumer it holds. nats stream rm -f stackoverflow ``` +Note that deleting a Quickwit source does **not** delete its consumer: the consumer is provisioned outside of Quickwit and outlives the source. Here the stream is deleted too, which removes it. Deleting only the source would leave the consumer behind, still holding its position. + This concludes the tutorial. If you have any questions regarding Quickwit or encounter any issues, don't hesitate to ask a [question](https://github.com/quickwit-oss/quickwit/discussions) or open an [issue](https://github.com/quickwit-oss/quickwit/issues) on [GitHub](https://github.com/quickwit-oss/quickwit) or contact us directly on [Discord](https://discord.com/invite/MT27AG5EVE). From 4fddd3ae7819664aaa0f7dec2a2c3ca2ab381853 Mon Sep 17 00:00:00 2001 From: tprevot Date: Wed, 9 Sep 2026 14:38:59 +0200 Subject: [PATCH 16/21] fix: source config --- docs/configuration/source-config.md | 45 ++++++----------------------- 1 file changed, 9 insertions(+), 36 deletions(-) diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 1ff6a607f53..380e9095104 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -189,7 +189,7 @@ A NATS source reads data from a [NATS JetStream](https://docs.nats.io/nats-conce A tutorial is available [here](/docs/ingest-data/nats.md). -The durable consumer is provisioned externally — Quickwit only ever fetches it, and never creates, updates, nor deletes it — so its lifecycle, subject filters, deliver policy, and ack tuning belong to whoever provisioned it. +The durable consumer is provisioned externally so its lifecycle, subject filters, deliver policy, and ack tuning belong to whoever provisioned it. Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. @@ -205,17 +205,17 @@ These are properties of the consumer, not of the source. Only the ack policy is ack_wait > commit_timeout + split upload and publish + ack round trip ``` -Shorter than their sum, and NATS redelivers messages that are still being indexed; they are then indexed twice. 5 minutes with a 60 s `commit_timeout` leaves a wide margin. +If `ack_wait` is shorter, NATS redelivers messages that are still being indexed; they are then indexed twice. 5 minutes with a 60 s `commit_timeout` leaves a wide margin. -It is also the *recovery* time from a lost delivery, so it should not be arbitrarily large either. If a message is delivered but never arrives — a dropped connection, a slow-consumer disconnect — nothing brings it back until the timer expires and the pipeline simply idles. The same injected fault idled a run for 278 s at `ack_wait=300s` and 23.2 s at `30s`. +It is also the *recovery* time from a lost delivery, so it should not be arbitrarily large either. If a message is delivered but never arrives (a dropped connection, a slow-consumer disconnect) nothing brings it back until the timer expires and the pipeline simply idles. -**`max_ack_pending` should be `-1`.** Nothing is acknowledged until a split is published, so a whole commit window is always ack-pending. That makes the setting a hard throughput cap rather than a safety valve: +**`max_ack_pending` should be `-1`.** Nothing is acknowledged until a split is published, so a whole commit window is always ack-pending. This setting is a hard throughput cap rather than a safety valve: ``` achievable rate ≈ max_ack_pending / commit_timeout documents per second ``` -Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents per second where the formula predicts 100, and `20000` gave 2,115 where it predicts 2,000. A bound that looks generous for an ordinary queue consumer throttles indexing to a trickle here. If the in-flight window has to be bounded, size it above `throughput × commit_timeout` rather than at a small absolute number, and note the trade: unlimited also makes the crash-replay window the whole delivered span above the ack floor rather than a bounded slice. +Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents per second where the formula predicts 100, and `20000` gave 2,115 where it predicts 2,000. A bound that looks generous for an ordinary queue consumer throttles indexing. If the in-flight window has to be bounded, size it above `throughput × commit_timeout` rather than at a small absolute number, and note the trade: unlimited also makes the crash-replay window the whole delivered span above the ack floor rather than a bounded slice. **`deliver_policy`** is usually `all`, so a new consumer indexes the stream from the start. @@ -226,48 +226,21 @@ Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents | Setting | Where | Effect | | --- | --- | --- | | `num_pipelines` | source config | Pipelines bound to the consumer. Scales cleanly on large messages: 52 → 147 MiB/s from one to four pipelines at 512 KiB. On small messages the acknowledgment path caps the aggregate, and it is worth only about 1.15× from one to sixteen — there, parallelism has to come from more consumers. | -| [`commit_timeout_secs`](index-config.md#indexing-settings) | index config | Sets the size of the always-ack-pending window, so it interacts with `max_ack_pending` and with `ack_wait`. It is also the length of the cooperative-indexing cycle. | -| [`docstore_compression_level`](index-config.md#indexing-settings) | index config | Defaults to zstd 8, which is what limits a single pipeline on small documents: the dedicated docstore thread saturates and back-pressures the indexer. Level 1 or 3 measured 21.1 s against 31.2 s on a 1.5 GiB corpus of 1 KiB documents. Not NATS-specific. | -| [`enable_cooperative_indexing`](node-config.md#indexer-configuration) | node config | See below. | +| [`commit_timeout_secs`](index-config.md#indexing-settings) | index config | Sets the size of the always-ack-pending window, so it interacts with `max_ack_pending` and with `ack_wait`. | | `QW_NATS_PULL_MAX_BYTES_PER_BATCH` | env, default 10 MiB | Bounds the bytes the server may push per pull request. It must stay well under the server's per-connection `max_pending` (64 MiB by default), or the server declares a slow consumer and closes the connection; the messages are already counted as delivered, so the pipeline then idles for a full `ack_wait`. 10 MiB and 20 MiB measure identically. | -| `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` | env, default 100,000 | Messages per pull request. Does not bind at these message sizes — 200 and 100,000 measured identically — because the byte cap is what bounds a batch. | +| `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` | env, default 100,000 | Messages per pull request. Can be used along `QW_NATS_PULL_MAX_BYTES_PER_BATCH` to limit a batch if there is too much contention around acknowledgements | #### Scaling -**One consumer per partition of the data** (recommended). Give each partition — a tenant, a subject, whatever the natural unit is — its own durable consumer with a single `filter_subject`, and add one Quickwit source with `num_pipelines: 1` per consumer. The partitions then scale and fail independently, and each gets its own checkpoint. - -**One shared consumer** (simplest). Several pipelines bind to the same consumer and NATS load-balances the messages across them, so scaling is a plain `num_pipelines` update and the control plane places the pipelines across the indexers of the cluster. - -At 512 KiB per message the two shapes are equivalent — 52.1, 91.3 and 146.6 MiB/s at one, two and four units either way — and so is partitioning across separate streams instead of subjects. At 1 KiB they diverge: a shared consumer stops gaining past a couple of pipelines while per-tenant consumers keep scaling, because the acknowledgment cost is per message rather than per pipeline. The per-consumer shape is therefore the safe default: never worse, and better on small documents. +Several pipelines can bind to the same consumer and NATS load-balances the messages across them, so scaling is a plain `num_pipelines` update and the control plane places the pipelines across the indexers of the cluster. #### Acknowledgment cost One confirmed acknowledgment per message is the source's dominant broker cost, and it is what decides whether the source or the indexer is the bottleneck. -- At **512 KiB** per message it is invisible. The NATS and Kafka sources measured within a percent of each other across the whole pipeline sweep, and broker CPU stayed at 0.5–0.9 cores for both, which is the host's floor. Priced on the broker alone, a confirmed acknowledgment per message still leaves 1,138 MiB/s — several times what one indexer consumes. +- At **512 KiB** per message it is invisible. The NATS and Kafka sources measured within a percent of each other across the whole pipeline sweep. - At **1 KiB** it is the ceiling. Sixteen per-tenant consumers reached 72 MiB/s while an equivalent Kafka source reached 150, and the broker spent 5.1 of the host's cores on acknowledgments against Quickwit's 4.7 on indexing — about ten times the broker CPU per document that a periodic offset commit costs. Priced on the broker alone, acknowledgment divides throughput by 6.5 at this size. -The cost is per *message*, so the message rate is what matters. Batching many log records into one message and letting an `otlp_*` input format expand them into separate documents is what moves a workload from the second case to the first: one acknowledgment then covers the whole batch. - -#### Cooperative indexing - -[`enable_cooperative_indexing`](node-config.md#indexer-configuration) is a **node-level** switch, off by default. It cannot be set per index or per source: one semaphore per indexer, sized to the node's blocking-thread count, is shared by every pipeline that node runs. - -With it off, an indexer cuts a split when the commit timeout fires or the split hits its size limits, and every pipeline indexes whenever it has work. With it on, a pipeline first sleeps a phase offset derived from a hash of its pipeline id, spread over `[0, commit_timeout)`; it then takes a semaphore permit, cuts its split as soon as *its own mailbox drains* rather than waiting for the commit timeout, releases the permit, and sleeps out the rest of a `commit_timeout`-long cycle. So a bounded number of pipelines hold an `IndexWriter` at any instant, and their CPU, disk and network spikes are spread across the commit window instead of landing together. - -What it costs and buys, measured on a 4.75 GiB corpus with a 10 s commit timeout: - -| | throughput, coop off → on | first published split, off → on | -| --- | --- | --- | -| 512 KiB messages, 1 pipeline | 52.1 → 52.1 MiB/s | — | -| 512 KiB messages, 16 pipelines | 177.8 → 209.5 MiB/s | 13.2 s → 3.0 s | -| ~1 KiB messages, 1 pipeline | 33.6 → 24.2 MiB/s | — | -| ~1 KiB messages, 16 pipelines | 38.6 → 36.7 MiB/s | 13.2 s → 3.0 s | - -At 512 KiB it costs nothing measurable and still makes the first split searchable 4.4× sooner. On small documents it costs up to 28 % for a single busy pipeline, shrinking to a few percent by sixteen — the penalty is worst for one pipeline with a lot to do, which is the opposite of what the setting is for. Enable it when a node runs many pipelines that are each mostly idle, which is also when its bound on concurrent `IndexWriter`s is what keeps the node inside its memory budget. - -One caveat: the initial phase sleep is up to one full `commit_timeout`, so with a 60 s commit timeout a pipeline can sit idle for up to a minute after a restart or reassignment before it indexes anything. That is a ramp after each deploy rather than an ongoing cost; shorten the commit timeout if it matters. - #### Monitoring Being durable, the consumer is observable through NATS's own monitoring (`nats consumer info`, exporters): `num_pending` is the indexing lag and `num_ack_pending` the in-flight window, both available even while the pipelines are down. The source also reports `num_pending_acks` (messages indexed but whose split is not published yet) in its observable state. From 29621c975173dfaaf3500a89069ea1840febf59d Mon Sep 17 00:00:00 2001 From: tprevot Date: Wed, 9 Sep 2026 16:52:11 +0200 Subject: [PATCH 17/21] feat(indexers): fixes --- docs/configuration/source-config.md | 18 +- docs/ingest-data/nats.md | 17 +- docs/overview/introduction.md | 2 +- docs/reference/rest-api.md | 4 +- quickwit/Cargo.lock | 1 - quickwit/quickwit-common/src/tracing_utils.rs | 51 ++- .../quickwit-config/src/source_config/mod.rs | 4 +- .../src/source_config/serialize.rs | 2 +- quickwit/quickwit-indexing/Cargo.toml | 5 +- .../src/source/nats_source.rs | 330 ++++++++++++++---- 10 files changed, 344 insertions(+), 90 deletions(-) diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 380e9095104..de5251d90bf 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -23,7 +23,7 @@ The source ID is a string that uniquely identifies the source within an index. I ## Source type -The source type designates the kind of source being configured. As of version 0.5, available source types are `ingest-api`, `kafka`, `kinesis`, and `pulsar`. The `file` type is also supported but only for local ingestion from [the CLI](/docs/reference/cli.md#tool-local-ingest). +The source type designates the kind of source being configured. Available source types are `ingest-api`, `kafka`, `kinesis`, `nats`, `pubsub`, and `pulsar`. The `file` type is also supported but only for local ingestion from [the CLI](/docs/reference/cli.md#tool-local-ingest). ## Source parameters @@ -193,9 +193,11 @@ The durable consumer is provisioned externally so its lifecycle, subject filters Delivery is **exactly-once on planned teardowns and at-least-once on crashes**. On a planned teardown (node shutdown, pipeline reassignment on a `num_pipelines` change), the pipeline is drained first: the source stops pulling, the in-flight messages are committed, published, and acknowledged before the pipeline stops, so nothing is indexed twice. Messages the pipeline had prefetched but not processed are negatively acknowledged so the remaining pipelines pick them up immediately. The drain runs under a time budget, the indexer's [`shutdown_drain_timeout`](node-config.md#indexer-configuration). A drain that cannot finish within it (e.g. the object storage or the metastore is unavailable) is abandoned and delivery degrades to at-least-once, as on a crash. After a crash, the unacknowledged messages are redelivered after `ack_wait` and indexed again, as duplicates. +The source never waits on unreachable NATS servers: acknowledgments and the negative acknowledgments sent at teardown are bounded by a 10 s timeout, unconfirmed acknowledgments stay pending and are retried at the next split publication, and a pipeline can still be stopped or reassigned during an outage. + #### Consumer invariants -These are properties of the consumer, not of the source. Only the ack policy is enforced when the source is created; the rest are not checked, and getting them wrong looks like Quickwit being slow or duplicating rather than like a consumer misconfiguration. +These are properties of the consumer, not of the source. Only the ack policy is enforced when the source is created. The pull limits and `max_deliver` are checked when a pipeline starts and only produce a warning in the indexer logs. The rest are not checked at all, and getting them wrong looks like Quickwit being slow or duplicating rather than like a consumer misconfiguration. **`ack_policy` must be `explicit`.** The source acknowledges each message individually once the split containing it is published, and waits for the server to confirm the acknowledgment — the confirmation is what tells the drain that the pipeline is empty. The consumer's ack floor is therefore the resume point, and it is the only progress state that matters. Any other policy is rejected when the source is created. @@ -217,6 +219,12 @@ achievable rate ≈ max_ack_pending / commit_timeout documents per second Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents per second where the formula predicts 100, and `20000` gave 2,115 where it predicts 2,000. A bound that looks generous for an ordinary queue consumer throttles indexing. If the in-flight window has to be bounded, size it above `throughput × commit_timeout` rather than at a small absolute number, and note the trade: unlimited also makes the crash-replay window the whole delivered span above the ack floor rather than a bounded slice. +**`max_deliver` should stay unlimited (`-1`).** Redelivery is what recovers from a crash or a lost delivery. A finite value turns it into data loss: a message that exhausts it is never delivered again and is silently skipped. The source warns at start when it is finite. + +**`max_batch`, `max_bytes` and `max_expires` must not be below what the source pulls.** Each pull request asks for up to `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` messages and `QW_NATS_PULL_MAX_BYTES_PER_BATCH` bytes with a 30 s expiry. A consumer limit below any of those makes the server reject every pull request and the pipeline idles. The source warns at start. + +**The consumer must outlive the source.** Deleting it, or letting an `inactive_threshold` expire, while a source is bound to it does not stop the pipelines: the server answers their pull requests with "no responders", which is also what a JetStream outage produces, so the source retries indefinitely and warns. Recreate the consumer under the same name to resume, or disable the source. + **`deliver_policy`** is usually `all`, so a new consumer indexes the stream from the start. **Give each consumer a single `filter_subject`.** A consumer configured with several filters through NATS 2.10's multi-filter `filter_subjects` appears to take a much slower path in the broker: on a stream whose subjects are interleaved, sixteen multi-filter consumers took 7.5 times longer than the same sixteen consumers with one filter each, with the broker saturated and the indexers idle. @@ -227,8 +235,8 @@ Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents | --- | --- | --- | | `num_pipelines` | source config | Pipelines bound to the consumer. Scales cleanly on large messages: 52 → 147 MiB/s from one to four pipelines at 512 KiB. On small messages the acknowledgment path caps the aggregate, and it is worth only about 1.15× from one to sixteen — there, parallelism has to come from more consumers. | | [`commit_timeout_secs`](index-config.md#indexing-settings) | index config | Sets the size of the always-ack-pending window, so it interacts with `max_ack_pending` and with `ack_wait`. | -| `QW_NATS_PULL_MAX_BYTES_PER_BATCH` | env, default 10 MiB | Bounds the bytes the server may push per pull request. It must stay well under the server's per-connection `max_pending` (64 MiB by default), or the server declares a slow consumer and closes the connection; the messages are already counted as delivered, so the pipeline then idles for a full `ack_wait`. 10 MiB and 20 MiB measure identically. | -| `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` | env, default 100,000 | Messages per pull request. Can be used along `QW_NATS_PULL_MAX_BYTES_PER_BATCH` to limit a batch if there is too much contention around acknowledgements | +| `QW_NATS_PULL_MAX_BYTES_PER_BATCH` | env, default 10 MiB | Bounds the bytes the server may push per pull request. It must stay well under the server's per-connection `max_pending` (64 MiB by default), or the server declares a slow consumer and closes the connection; the messages are already counted as delivered, so the pipeline then idles for a full `ack_wait`. It must also stay above the server's `max_payload`, or larger messages are never delivered, and at or below the consumer's `max_bytes`; the source warns at start otherwise. 10 MiB and 20 MiB measure identically. | +| `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` | env, default 100,000 | Messages per pull request. Can be used along `QW_NATS_PULL_MAX_BYTES_PER_BATCH` to limit a batch if there is too much contention around acknowledgements. The client buffers eight such batches per subscription; messages beyond that are dropped client-side and come back after `ack_wait`. | #### Scaling @@ -252,7 +260,7 @@ Being durable, the consumer is observable through NATS's own monitoring (`nats c | `uris` | List of NATS server URIs (e.g. `nats://localhost:4222`). | required | | `stream` | Name of the JetStream stream to consume. | required | | `consumer` | Name of the pre-provisioned durable consumer to bind to. | required | -| `tls` | TLS options: `ca_certificates_path` (PEM file whose root certificates are trusted in addition to the system ones), and `client_certificate_path` + `client_key_path` (PEM files, set together) for mutual TLS. TLS itself is enabled by connecting to `tls://` URIs. The files are read by the indexer nodes when the connection is established. | optional | +| `tls` | TLS options: `ca_certificates_path` (PEM file whose root certificates are trusted instead of the system ones), and `client_certificate_path` + `client_key_path` (PEM files, set together) for mutual TLS. TLS itself is enabled by connecting to `tls://` URIs. The files are read when a connection is established: by the indexer nodes running the source, and by the node serving the source creation or update, which checks connectivity. | optional | | `authentication` | Authentication parameters: either `user_password` (with `user` and `password`) or `token`. | optional | *Adding a NATS source to an index with the [CLI](../reference/cli.md#source)* diff --git a/docs/ingest-data/nats.md b/docs/ingest-data/nats.md index 54b669faf58..42e71e75691 100644 --- a/docs/ingest-data/nats.md +++ b/docs/ingest-data/nats.md @@ -131,17 +131,28 @@ The NATS source consumes a JetStream stream through a durable consumer, so messa ```bash nats stream add stackoverflow --subjects "stackoverflow.posts" --defaults -nats consumer add stackoverflow quickwit-consumer --pull --deliver all --ack explicit --max-pending=-1 --wait 5m --defaults +cat > quickwit-consumer.json < Context { + global::get_text_map_propagator(|propagator| propagator.extract(extractor)) +} + /// Extracts an OpenTelemetry context from incoming gRPC request metadata. /// Returns the empty context when no propagator is installed or no headers /// are present. pub fn extract_context(metadata: &MetadataMap) -> Context { - let extractor = MetadataExtractor(metadata); - global::get_text_map_propagator(|propagator| propagator.extract(&extractor)) + extract_context_from(&MetadataExtractor(metadata)) +} + +/// Extracts the context propagated through `extractor` (e.g. a W3C +/// `traceparent` header), or `None` when it carries no valid span context, so +/// callers can skip creating a span for it. +pub fn extract_remote_context(extractor: &impl Extractor) -> Option { + let context = extract_context_from(extractor); + if !context.span().span_context().is_valid() { + return None; + } + Some(context) +} + +/// Parents `span` on `parent_context`, typically one returned by +/// [`extract_remote_context`]. +pub fn set_span_parent(span: &Span, parent_context: Context) { + let _ = span.set_parent(parent_context); } /// Extracts a W3C trace context from incoming gRPC request metadata and @@ -87,7 +109,7 @@ pub fn extract_context(metadata: &MetadataMap) -> Context { /// caller's trace. pub fn set_current_span_parent_from_metadata(metadata: &MetadataMap) { let parent_context = extract_context(metadata); - let _ = Span::current().set_parent(parent_context); + set_span_parent(&Span::current(), parent_context); } /// Tonic interceptor that injects the active span's W3C trace context into @@ -149,6 +171,27 @@ mod tests { ); } + #[test] + fn extract_remote_context_requires_a_valid_span_context() { + // The helper reads the process-wide propagator, a no-op by default. + global::set_text_map_propagator(TraceContextPropagator::new()); + + let mut metadata = MetadataMap::new(); + assert!(extract_remote_context(&MetadataExtractor(&metadata)).is_none()); + + metadata.insert( + "traceparent", + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" + .parse() + .unwrap(), + ); + let context = extract_remote_context(&MetadataExtractor(&metadata)).unwrap(); + assert_eq!( + context.span().span_context().trace_id(), + known_span_context().trace_id() + ); + } + #[test] fn extract_returns_invalid_span_context_when_no_traceparent() { let propagator = TraceContextPropagator::new(); diff --git a/quickwit/quickwit-config/src/source_config/mod.rs b/quickwit/quickwit-config/src/source_config/mod.rs index 9283f2a1dd5..afcedc3c747 100644 --- a/quickwit/quickwit-config/src/source_config/mod.rs +++ b/quickwit/quickwit-config/src/source_config/mod.rs @@ -683,8 +683,8 @@ impl NatsSourceParams { )] #[serde(deny_unknown_fields)] pub struct NatsSourceTls { - /// Path to a PEM file whose root certificates are trusted in addition to - /// the system ones (e.g. a private CA). + /// Path to a PEM file whose root certificates are trusted instead of the + /// system ones (e.g. a private CA). #[serde(default, skip_serializing_if = "Option::is_none")] pub ca_certificates_path: Option, /// Path to the client certificate PEM file, for mutual TLS. Requires diff --git a/quickwit/quickwit-config/src/source_config/serialize.rs b/quickwit/quickwit-config/src/source_config/serialize.rs index 3b412a9ab19..e00402aafcd 100644 --- a/quickwit/quickwit-config/src/source_config/serialize.rs +++ b/quickwit/quickwit-config/src/source_config/serialize.rs @@ -132,7 +132,7 @@ impl SourceConfigForSerialization { _ => { if self.num_pipelines > 1 { bail!( - "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or NATS sources.\ + "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or NATS sources. \ open an issue https://github.com/quickwit-oss/quickwit/issues if you need the feature for other source types" ); } diff --git a/quickwit/quickwit-indexing/Cargo.toml b/quickwit/quickwit-indexing/Cargo.toml index 172dfbf0133..bf8aaaffc60 100644 --- a/quickwit/quickwit-indexing/Cargo.toml +++ b/quickwit/quickwit-indexing/Cargo.toml @@ -32,7 +32,6 @@ itertools = { workspace = true } libz-sys = { workspace = true, optional = true } oneshot = { workspace = true } openssl = { workspace = true, optional = true } -opentelemetry = { workspace = true, optional = true } percent-encoding = { workspace = true } pulsar = { workspace = true, optional = true } quickwit-query = { workspace = true } @@ -47,7 +46,6 @@ thiserror = { workspace = true } time = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } -tracing-opentelemetry = { workspace = true, optional = true } ulid = { workspace = true } utoipa = { workspace = true } vrl = { workspace = true, optional = true } @@ -86,7 +84,7 @@ kinesis = [ "quickwit-aws/kinesis", ] kinesis-localstack-tests = [] -nats = ["dep:async-nats", "dep:opentelemetry", "dep:tracing-opentelemetry"] +nats = ["dep:async-nats"] nats-broker-tests = ["nats"] pulsar = ["dep:pulsar"] pulsar-broker-tests = [] @@ -122,6 +120,7 @@ ci-test = [] bytes = { workspace = true } criterion = { workspace = true, features = ["async_tokio"] } mockall = { workspace = true } +opentelemetry = { workspace = true } opentelemetry_sdk = { workspace = true } parquet = { workspace = true } proptest = { workspace = true } diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index f79c8ba68c3..58091ba7dea 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -34,21 +34,19 @@ use std::str::FromStr; use std::time::{Duration, Instant}; use anyhow::{Context as _, anyhow, bail, ensure}; -use async_nats::header::{HeaderMap, HeaderName}; +use async_nats::connection::State; +use async_nats::header::HeaderMap; use async_nats::jetstream::consumer::pull::{ MessagesError, MessagesErrorKind, Stream as DurableMessageStream, }; use async_nats::jetstream::consumer::{AckPolicy, PullConsumer}; use async_nats::jetstream::message::AckKind; -use async_nats::{ConnectOptions, Subject, jetstream}; +use async_nats::{ConnectOptions, HeaderName, Subject, jetstream}; use async_trait::async_trait; use bytesize::ByteSize; use futures::{FutureExt, StreamExt}; -use opentelemetry::propagation::Extractor; -use opentelemetry::trace::TraceContextExt; -use opentelemetry::{Context as OtelContext, global}; use quickwit_actors::ActorExitStatus; -use quickwit_common::rand::append_random_suffix; +use quickwit_common::tracing_utils::{self, Extractor}; use quickwit_config::{NatsSourceAuth, NatsSourceParams}; use quickwit_metastore::checkpoint::{PartitionId, SourceCheckpoint}; use quickwit_proto::metastore::SourceType; @@ -56,7 +54,6 @@ use quickwit_proto::types::Position; use serde_json::{Value as JsonValue, json}; use tokio::time; use tracing::{Span, debug, info, warn}; -use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::source::{ BATCH_NUM_BYTES_LIMIT, BatchBuilder, EMIT_BATCHES_TIMEOUT, Source, SourceContext, @@ -122,22 +119,30 @@ impl NatsSource { source_params: NatsSourceParams, ) -> anyhow::Result { let consumer_name = source_params.consumer.clone(); + let pull_max_bytes_per_batch = pull_max_bytes_per_batch(); info!( index_id=%source_runtime.index_id(), source_id=%source_runtime.source_id(), stream=%source_params.stream, %consumer_name, - pull_max_bytes_per_batch=%pull_max_bytes_per_batch(), + %pull_max_bytes_per_batch, pull_max_messages_per_batch=%pull_max_messages_per_batch(), "starting NATS source" ); let (nats_client, consumer) = connect_and_fetch_consumer(&source_params).await?; + warn_on_incompatible_pull_settings( + &nats_client, + &consumer.cached_info().config, + pull_max_bytes_per_batch, + pull_max_messages_per_batch(), + ); + let message_stream = consumer .stream() .max_messages_per_batch(pull_max_messages_per_batch()) - .max_bytes_per_batch(pull_max_bytes_per_batch().0 as usize) + .max_bytes_per_batch(pull_max_bytes_per_batch.0 as usize) // default heartbeat that is disabled when using manual builder, // required for transient conn errors. .heartbeat(PULL_IDLE_HEARTBEAT) @@ -146,11 +151,23 @@ impl NatsSource { .await .context("failed to subscribe to NATS consumer messages")?; - // A fresh partition per pipeline: durable-mode positions are delivery - // counters local to this instance, and pipelines sharing the consumer - // must not collide on a partition. - let partition_id = - PartitionId::from(append_random_suffix(&format!("nats-{consumer_name}"))); + // One partition per pipeline: positions are delivery counters local to + // this pipeline, and pipelines sharing the consumer must not collide on + // a partition. The pipeline uid is stable across respawns, so a respawn + // resumes the counter from the checkpoint instead of opening a new + // partition. + let partition_id = PartitionId::from(format!( + "nats-{consumer_name}-{}", + source_runtime.pipeline_uid() + )); + let checkpoint = source_runtime + .fetch_checkpoint() + .await + .context("failed to fetch the source checkpoint")?; + let delivery_counter = checkpoint + .position_for_partition(&partition_id) + .and_then(Position::as_u64) + .unwrap_or(0); Ok(NatsSource { source_runtime, @@ -159,7 +176,7 @@ impl NatsSource { message_stream, consumer_name, partition_id, - delivery_counter: 0, + delivery_counter, pending_acks: BTreeMap::new(), state: NatsSourceState::default(), }) @@ -221,14 +238,10 @@ impl NatsSource { if acks.is_empty() { return; } + let nats_client = &self.nats_client; let mut ack_results = futures::stream::iter(acks) .map(async |(delivery_counter, ack_subject)| { - // Mirrors `jetstream::Message::double_ack()`, but through the - // client's muxed inbox rather than a subscription per message. - let ack_result = self - .nats_client - .request(ack_subject, AckKind::Ack.into()) - .await; + let ack_result = send_ack(nats_client, ack_subject).await; (delivery_counter, ack_result) }) .buffer_unordered(MAX_CONCURRENT_ACK_REQUESTS); @@ -238,7 +251,7 @@ impl NatsSource { let mut last_ack_error = None; while let Some((delivery_counter, ack_result)) = ack_results.next().await { match ack_result { - Ok(_ack_reply) => { + Ok(()) => { self.pending_acks.remove(&delivery_counter); num_acks += 1; } @@ -253,6 +266,38 @@ impl NatsSource { } debug!(num_acks, "acked published messages"); } + + /// Negatively acknowledges the messages the client prefetched, so a + /// surviving pipeline picks them up promptly instead of after `ack_wait`. + async fn nak_prefetched_messages(&mut self) { + let mut num_naks = 0usize; + while let Some(Some(message_res)) = self.message_stream.next().now_or_never() { + let message = match message_res { + Ok(message) => message, + Err(error) => { + warn!(%error, "failed to pull a prefetched NATS message to negatively acknowledge it"); + continue; + } + }; + let Some(ack_subject) = message.message.reply else { + warn!( + "prefetched NATS message carries no reply subject to negatively acknowledge \ + it on" + ); + continue; + }; + if let Err(error) = self + .nats_client + .publish(ack_subject, AckKind::Nak(Some(NAK_REDELIVERY_DELAY)).into()) + .await + { + warn!(%error, "failed to negatively acknowledge a prefetched NATS message"); + continue; + } + num_naks += 1; + } + debug!(num_naks, "negatively acknowledged prefetched messages"); + } } #[async_trait] @@ -277,7 +322,7 @@ impl Source for NatsSource { let message = match message_res { Ok(message) => message, Err(error) if is_transient_stream_error(&error) => { - warn!(%error, "transient NATS message stream error"); + warn!(%error, "transient NATS message stream error, retrying"); wait_before_next_batch = TRANSIENT_STREAM_ERROR_BACKOFF; break; } @@ -351,38 +396,21 @@ impl Source for NatsSource { ) -> anyhow::Result<()> { // Any ack still pending at this point was not completed by a drain: // its message is redelivered after the consumer's `ack_wait`. - let mut num_naks = 0usize; - while let Some(Some(message_res)) = self.message_stream.next().now_or_never() { - let message = match message_res { - Ok(message) => message, - Err(error) => { - warn!(%error, "failed to pull a prefetched NATS message to negatively acknowledge it"); - continue; - } - }; - let Some(ack_subject) = message.message.reply else { - warn!( - "prefetched NATS message carries no reply subject to negatively acknowledge \ - it on" - ); - continue; - }; - if let Err(error) = self - .nats_client - .publish(ack_subject, AckKind::Nak(Some(NAK_REDELIVERY_DELAY)).into()) - .await - { - warn!(%error, "failed to negatively acknowledge a prefetched NATS message"); - continue; + // + // Bounded for the same reason as `ACK_REQUEST_TIMEOUT`: a dead server + // must not hold the pipeline teardown, and node shutdown with it. Past + // the deadline, `ack_wait` redelivery takes over. + let teardown = async { + self.nak_prefetched_messages().await; + if let Err(error) = self.nats_client.flush().await { + warn!(%error, "failed to flush NATS negative acknowledgments"); } - num_naks += 1; - } - - if let Err(error) = self.nats_client.flush().await { - warn!(%error, num_naks, "failed to flush NATS negative acknowledgments"); - } - if let Err(error) = self.nats_client.drain().await { - warn!(%error, "failed to drain the NATS connection"); + if let Err(error) = self.nats_client.drain().await { + warn!(%error, "failed to drain the NATS connection"); + } + }; + if time::timeout(FINALIZE_TIMEOUT, teardown).await.is_err() { + warn!("timed out negatively acknowledging prefetched NATS messages"); } Ok(()) } @@ -402,6 +430,25 @@ impl Source for NatsSource { } } +/// Server-confirmed acknowledgment ("double ack"). Mirrors +/// `jetstream::Message::double_ack()`, but through the client's muxed inbox +/// rather than a subscription per message. +async fn send_ack(nats_client: &async_nats::Client, ack_subject: Subject) -> anyhow::Result<()> { + // Nothing drains the client's outbound queue while it reconnects, so a + // request issued now would only sit there. Unconfirmed acks stay pending + // and are retried by the next `SuggestTruncate`. + ensure!( + nats_client.connection_state() == State::Connected, + "NATS connection is down" + ); + let request = nats_client.request(ack_subject, AckKind::Ack.into()); + let ack_reply_res = time::timeout(ACK_REQUEST_TIMEOUT, request) + .await + .context("ack request timed out")?; + ack_reply_res.context("ack request failed")?; + Ok(()) +} + /// Bounds the concurrent server-confirmed ack requests: enough to hide the /// round-trip latency without flooding the connection. /// @@ -426,6 +473,70 @@ fn pull_max_bytes_per_batch() -> ByteSize { ) } +/// Warns about pull settings the server or the consumer would reject, or that +/// lose messages. The consumer's limits belong to whoever provisioned it, so +/// the source reports them rather than refusing to start. +fn warn_on_incompatible_pull_settings( + nats_client: &async_nats::Client, + consumer_config: &jetstream::consumer::Config, + pull_max_bytes_per_batch: ByteSize, + pull_max_messages_per_batch: usize, +) { + let pull_max_bytes = pull_max_bytes_per_batch.as_u64(); + if pull_max_bytes < BATCH_NUM_BYTES_LIMIT { + warn!( + %pull_max_bytes_per_batch, + indexing_batch_num_bytes=BATCH_NUM_BYTES_LIMIT, + "`QW_NATS_PULL_MAX_BYTES_PER_BATCH` is below the indexing batch size: every batch \ + costs several pull round trips" + ); + } + let max_payload = nats_client.server_info().max_payload as u64; + if max_payload > 0 && pull_max_bytes < max_payload { + warn!( + %pull_max_bytes_per_batch, + max_payload, + "`QW_NATS_PULL_MAX_BYTES_PER_BATCH` is below the server's `max_payload`: messages \ + larger than it can never be delivered" + ); + } + let consumer_max_bytes = consumer_config.max_bytes; + if consumer_max_bytes > 0 && pull_max_bytes > consumer_max_bytes as u64 { + warn!( + %pull_max_bytes_per_batch, + consumer_max_bytes, + "`QW_NATS_PULL_MAX_BYTES_PER_BATCH` exceeds the consumer's `max_bytes`: the server \ + rejects every pull request" + ); + } + let consumer_max_batch = consumer_config.max_batch; + if consumer_max_batch > 0 && pull_max_messages_per_batch as i64 > consumer_max_batch { + warn!( + pull_max_messages_per_batch, + consumer_max_batch, + "`QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` exceeds the consumer's `max_batch`: the server \ + rejects every pull request" + ); + } + let consumer_max_expires = consumer_config.max_expires; + if !consumer_max_expires.is_zero() && PULL_EXPIRES > consumer_max_expires { + warn!( + ?consumer_max_expires, + pull_expires=?PULL_EXPIRES, + "the consumer's `max_expires` is below the source's pull expiry: the server \ + rejects every pull request" + ); + } + let consumer_max_deliver = consumer_config.max_deliver; + if consumer_max_deliver > 0 { + warn!( + consumer_max_deliver, + "the consumer's `max_deliver` is finite: a message that exhausts it is never \ + redelivered, so a crash or a lost delivery can leave it unindexed" + ); + } +} + /// Messages per pull batch. The byte cap above is what actually bounds a batch; /// this only binds for small messages, where it caps a pull far below the byte /// budget (200 messages of 1 KiB is under 200 KiB against a 10 MiB budget) and @@ -442,6 +553,15 @@ fn pull_max_messages_per_batch() -> usize { ) } +/// Capacity of the client's per-subscription channel, in pull batches. The +/// connection handler drops a message that finds the channel full, and the +/// server already counts it as delivered, so it only comes back after +/// `ack_wait`. The client keeps up to 1.5 batches outstanding and re-pulls a +/// full batch every 35 s while the source is not polling (e.g. during a +/// drain), so the channel must hold several. Memory is bounded by the byte +/// cap rather than by this: eight batches are at most 80 MiB by default. +const SUBSCRIPTION_CAPACITY_IN_PULL_BATCHES: usize = 8; + // Both kept at the `Consumer::messages` defaults, so switching to the builder // changes the batch bounds and nothing else. const PULL_IDLE_HEARTBEAT: Duration = Duration::from_secs(15); @@ -476,6 +596,12 @@ fn is_transient_stream_error(error: &MessagesError) -> bool { /// messages: 1024 is indistinguishable from 128, and 8192 is 20 % slower. const MAX_CONCURRENT_ACK_REQUESTS: usize = 128; +/// Bounds an ack request end to end, enqueueing included. +const ACK_REQUEST_TIMEOUT: Duration = Duration::from_secs(10); + +/// Bounds the NAK sweep and connection drain at teardown. +const FINALIZE_TIMEOUT: Duration = Duration::from_secs(10); + async fn fetch_durable_consumer( jetstream_stream: &jetstream::stream::Stream, consumer_name: &str, @@ -492,8 +618,8 @@ async fn fetch_durable_consumer( Ok(consumer) } -/// W3C trace context extraction from NATS message headers, following the -/// pattern of `quickwit_common::tracing_utils` for gRPC metadata. +/// Lets `quickwit_common::tracing_utils` read the W3C trace context from +/// NATS message headers. struct NatsHeaderExtractor<'a>(&'a HeaderMap); impl Extractor for NatsHeaderExtractor<'_> { @@ -507,10 +633,6 @@ impl Extractor for NatsHeaderExtractor<'_> { } } -fn extract_remote_context(headers: &HeaderMap) -> OtelContext { - global::get_text_map_propagator(|propagator| propagator.extract(&NatsHeaderExtractor(headers))) -} - /// Builds a span parented on the publisher's trace when the message carries /// a W3C `traceparent` header, stitching the processing of the message into /// the publisher's distributed trace. Messages without a propagated context @@ -521,10 +643,7 @@ fn remote_parented_span( source_runtime: &SourceRuntime, ) -> Option { let headers = message.headers.as_ref()?; - let parent_context = extract_remote_context(headers); - if !parent_context.span().span_context().is_valid() { - return None; - } + let parent_context = tracing_utils::extract_remote_context(&NatsHeaderExtractor(headers))?; let span = tracing::info_span!( "process_nats_message", index_id = %source_runtime.index_id(), @@ -532,12 +651,14 @@ fn remote_parented_span( subject = %message.subject, stream_sequence, ); - let _ = span.set_parent(parent_context); + tracing_utils::set_span_parent(&span, parent_context); Some(span) } async fn connect_nats(params: &NatsSourceParams) -> anyhow::Result { - let mut connect_options = ConnectOptions::new(); + let subscription_capacity = + SUBSCRIPTION_CAPACITY_IN_PULL_BATCHES * pull_max_messages_per_batch(); + let mut connect_options = ConnectOptions::new().subscription_capacity(subscription_capacity); match params.authentication.clone() { None => {} Some(NatsSourceAuth::UserPassword { user, password }) => { @@ -595,6 +716,7 @@ mod tests { #[test] fn extract_trace_context_from_headers() { use opentelemetry::propagation::TextMapPropagator; + use opentelemetry::trace::TraceContextExt; use opentelemetry_sdk::propagation::TraceContextPropagator; // The global propagator is not installed in tests, so the extractor @@ -618,6 +740,19 @@ mod tests { ); assert_eq!(span_context.span_id().to_string(), "00f067aa0ba902b7"); + // Go publishers canonicalize the header name. + let mut capitalized_headers = HeaderMap::new(); + capitalized_headers.insert( + "Traceparent", + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + ); + let span_context = propagator + .extract(&NatsHeaderExtractor(&capitalized_headers)) + .span() + .span_context() + .clone(); + assert!(span_context.is_valid()); + let empty_headers = HeaderMap::new(); let span_context = propagator .extract(&NatsHeaderExtractor(&empty_headers)) @@ -662,11 +797,12 @@ mod nats_broker_tests { use bytes::Bytes; use quickwit_actors::{ActorHandle, Inbox, Universe}; + use quickwit_common::rand::append_random_suffix; use quickwit_config::{SourceConfig, SourceInputFormat, SourceParams}; - use quickwit_metastore::checkpoint::SourceCheckpointDelta; + use quickwit_metastore::checkpoint::{PartitionDelta, SourceCheckpointDelta}; use quickwit_metastore::metastore_for_test; use quickwit_proto::metastore::MetastoreServiceClient; - use quickwit_proto::types::IndexUid; + use quickwit_proto::types::{IndexUid, PipelineUid}; use super::*; use crate::actors::DocProcessor; @@ -1137,6 +1273,64 @@ mod nats_broker_tests { universe.assert_quit().await; } + #[tokio::test] + async fn durable_mode_respawn_resumes_partition_from_checkpoint() { + let universe = Universe::with_accelerated_time(); + let metastore = metastore_for_test(); + let stream = append_random_suffix("test-nats-source--durable-resume--stream"); + let jetstream_ctx = setup_nats_stream(&stream).await; + let consumer_name = "durable-resume-consumer"; + provision_durable_consumer(&jetstream_ctx, &stream, consumer_name).await; + + let index_id = append_random_suffix("test-nats-source--durable-resume--index"); + let source_config = get_durable_source_config(&stream, consumer_name); + // The partition of the pipeline the test runtime builds, as left by a + // previous incarnation that published 7 messages. + let partition_id = PartitionId::from(format!( + "nats-{consumer_name}-{}", + PipelineUid::for_test(0u128) + )); + let index_uid = setup_index( + metastore.clone(), + &index_id, + &source_config, + &[( + partition_id.clone(), + Position::Beginning, + Position::offset(7u64), + )], + ) + .await; + + let (source_handle, doc_processor_inbox) = + create_source_actor(&universe, metastore, index_uid, source_config).await; + let subject = format!("{stream}.logs"); + publish_docs(&jetstream_ctx, &subject, 0..3).await; + wait_for_processed_messages(&source_handle, 3, Duration::from_secs(30)).await; + source_handle.quit().await; + + let batches: Vec = doc_processor_inbox.drain_for_test_typed(); + let checkpoint_delta = merge_doc_batches(batches).checkpoint_delta; + let partition_deltas: Vec<(PartitionId, PartitionDelta)> = + checkpoint_delta.iter().collect(); + // Same partition, counter continued from the checkpoint: the delta + // chains onto the published position instead of opening a new + // partition from the beginning. + assert_eq!( + partition_deltas, + vec![( + partition_id, + PartitionDelta { + from: Position::offset(7u64), + to: Position::offset(10u64), + } + )] + ); + + jetstream_ctx.delete_stream(&stream).await.unwrap(); + universe.assert_quit().await; + } + #[tokio::test] async fn durable_mode_nak_redelivers_prefetched_messages_promptly() { let universe = Universe::with_accelerated_time(); From 76ccf71c3d9889b9f50835363051e48de75c8f65 Mon Sep 17 00:00:00 2001 From: tprevot Date: Wed, 9 Sep 2026 17:01:50 +0200 Subject: [PATCH 18/21] feat(indexers): remove useless hb --- .../src/source/nats_source.rs | 35 +++---------------- 1 file changed, 5 insertions(+), 30 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 58091ba7dea..4f698e71287 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -143,9 +143,6 @@ impl NatsSource { .stream() .max_messages_per_batch(pull_max_messages_per_batch()) .max_bytes_per_batch(pull_max_bytes_per_batch.0 as usize) - // default heartbeat that is disabled when using manual builder, - // required for transient conn errors. - .heartbeat(PULL_IDLE_HEARTBEAT) .expires(PULL_EXPIRES) .messages() .await @@ -562,9 +559,8 @@ fn pull_max_messages_per_batch() -> usize { /// cap rather than by this: eight batches are at most 80 MiB by default. const SUBSCRIPTION_CAPACITY_IN_PULL_BATCHES: usize = 8; -// Both kept at the `Consumer::messages` defaults, so switching to the builder -// changes the batch bounds and nothing else. -const PULL_IDLE_HEARTBEAT: Duration = Duration::from_secs(15); +/// Server-side lifetime of a pull request. The client re-pulls 5 s after it +/// elapses, which is also what recovers a pull request the server lost. const PULL_EXPIRES: Duration = Duration::from_secs(30); /// Pause before pulling again after a transient stream error, to avoid a hot @@ -579,15 +575,11 @@ const TRANSIENT_STREAM_ERROR_BACKOFF: Duration = Duration::from_secs(1); const NAK_REDELIVERY_DELAY: Duration = Duration::from_secs(1); /// Errors the pull stream recovers from on its own: the subscription stays -/// usable and polling simply resumes. `MissingHeartbeat` in particular also -/// fires when the stream was not polled for a while (e.g. under downstream -/// backpressure), not only when the server is unreachable. +/// usable and polling simply resumes. fn is_transient_stream_error(error: &MessagesError) -> bool { matches!( error.kind(), - MessagesErrorKind::MissingHeartbeat - | MessagesErrorKind::Pull - | MessagesErrorKind::NoResponders + MessagesErrorKind::Pull | MessagesErrorKind::NoResponders ) } @@ -740,19 +732,6 @@ mod tests { ); assert_eq!(span_context.span_id().to_string(), "00f067aa0ba902b7"); - // Go publishers canonicalize the header name. - let mut capitalized_headers = HeaderMap::new(); - capitalized_headers.insert( - "Traceparent", - "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", - ); - let span_context = propagator - .extract(&NatsHeaderExtractor(&capitalized_headers)) - .span() - .span_context() - .clone(); - assert!(span_context.is_valid()); - let empty_headers = HeaderMap::new(); let span_context = propagator .extract(&NatsHeaderExtractor(&empty_headers)) @@ -764,11 +743,7 @@ mod tests { #[test] fn transient_stream_error_classification() { - let transient_kinds = [ - MessagesErrorKind::MissingHeartbeat, - MessagesErrorKind::Pull, - MessagesErrorKind::NoResponders, - ]; + let transient_kinds = [MessagesErrorKind::Pull, MessagesErrorKind::NoResponders]; for kind in transient_kinds { assert!( is_transient_stream_error(&MessagesError::new(kind)), From 857a3fe72ce50ee784234ce549f0edc2f6a8d44b Mon Sep 17 00:00:00 2001 From: tprevot Date: Wed, 9 Sep 2026 22:17:11 +0200 Subject: [PATCH 19/21] fix: minor fixes --- .../src/source/nats_source.rs | 37 ++++++------------- 1 file changed, 11 insertions(+), 26 deletions(-) diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 4f698e71287..01407691723 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -143,7 +143,6 @@ impl NatsSource { .stream() .max_messages_per_batch(pull_max_messages_per_batch()) .max_bytes_per_batch(pull_max_bytes_per_batch.0 as usize) - .expires(PULL_EXPIRES) .messages() .await .context("failed to subscribe to NATS consumer messages")?; @@ -266,8 +265,8 @@ impl NatsSource { /// Negatively acknowledges the messages the client prefetched, so a /// surviving pipeline picks them up promptly instead of after `ack_wait`. - async fn nak_prefetched_messages(&mut self) { - let mut num_naks = 0usize; + async fn nak_prefetched_messages(&mut self) -> usize { + let mut num_naks: usize = 0; while let Some(Some(message_res)) = self.message_stream.next().now_or_never() { let message = match message_res { Ok(message) => message, @@ -293,7 +292,8 @@ impl NatsSource { } num_naks += 1; } - debug!(num_naks, "negatively acknowledged prefetched messages"); + + num_naks } } @@ -398,7 +398,13 @@ impl Source for NatsSource { // must not hold the pipeline teardown, and node shutdown with it. Past // the deadline, `ack_wait` redelivery takes over. let teardown = async { - self.nak_prefetched_messages().await; + let naks_count = self.nak_prefetched_messages().await; + info!( + naks_count, + "sending negative acks to nats, they will be redelivered in {}s", + ACK_REQUEST_TIMEOUT.as_secs() + ); + if let Err(error) = self.nats_client.flush().await { warn!(%error, "failed to flush NATS negative acknowledgments"); } @@ -515,23 +521,6 @@ fn warn_on_incompatible_pull_settings( rejects every pull request" ); } - let consumer_max_expires = consumer_config.max_expires; - if !consumer_max_expires.is_zero() && PULL_EXPIRES > consumer_max_expires { - warn!( - ?consumer_max_expires, - pull_expires=?PULL_EXPIRES, - "the consumer's `max_expires` is below the source's pull expiry: the server \ - rejects every pull request" - ); - } - let consumer_max_deliver = consumer_config.max_deliver; - if consumer_max_deliver > 0 { - warn!( - consumer_max_deliver, - "the consumer's `max_deliver` is finite: a message that exhausts it is never \ - redelivered, so a crash or a lost delivery can leave it unindexed" - ); - } } /// Messages per pull batch. The byte cap above is what actually bounds a batch; @@ -559,10 +548,6 @@ fn pull_max_messages_per_batch() -> usize { /// cap rather than by this: eight batches are at most 80 MiB by default. const SUBSCRIPTION_CAPACITY_IN_PULL_BATCHES: usize = 8; -/// Server-side lifetime of a pull request. The client re-pulls 5 s after it -/// elapses, which is also what recovers a pull request the server lost. -const PULL_EXPIRES: Duration = Duration::from_secs(30); - /// Pause before pulling again after a transient stream error, to avoid a hot /// error loop while the connection recovers. `SuggestTruncate` is still /// processed in the meantime. From 3be262cb66ecbd7c6e02cb201edb443327c83548 Mon Sep 17 00:00:00 2001 From: tprevot Date: Wed, 9 Sep 2026 22:29:39 +0200 Subject: [PATCH 20/21] fix: minor fixes --- docker-compose.yml | 1 - docs/configuration/source-config.md | 6 +++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index 9198d13aba6..a2e88600caa 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -128,7 +128,6 @@ services: retries: 100 nats: - # 2.10 is the oldest version supporting multiple consumer filter subjects. image: nats:${NATS_VERSION:-2.10-alpine} container_name: nats command: --jetstream --http_port 8222 diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index de5251d90bf..7046843f3c8 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -219,9 +219,9 @@ achievable rate ≈ max_ack_pending / commit_timeout documents per second Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents per second where the formula predicts 100, and `20000` gave 2,115 where it predicts 2,000. A bound that looks generous for an ordinary queue consumer throttles indexing. If the in-flight window has to be bounded, size it above `throughput × commit_timeout` rather than at a small absolute number, and note the trade: unlimited also makes the crash-replay window the whole delivered span above the ack floor rather than a bounded slice. -**`max_deliver` should stay unlimited (`-1`).** Redelivery is what recovers from a crash or a lost delivery. A finite value turns it into data loss: a message that exhausts it is never delivered again and is silently skipped. The source warns at start when it is finite. +**`max_deliver` should stay unlimited (`-1`).** Redelivery is what recovers from a crash or a lost delivery. A finite value turns it into data loss: a message that exhausts it is never delivered again and is silently skipped. -**`max_batch`, `max_bytes` and `max_expires` must not be below what the source pulls.** Each pull request asks for up to `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` messages and `QW_NATS_PULL_MAX_BYTES_PER_BATCH` bytes with a 30 s expiry. A consumer limit below any of those makes the server reject every pull request and the pipeline idles. The source warns at start. +**`max_batch`, `max_bytes` must not be below what the source pulls.** Each pull request asks for up to `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` messages and `QW_NATS_PULL_MAX_BYTES_PER_BATCH` bytes. A consumer limit below any of those makes the server reject every pull request and the pipeline idles. The source warns at start. **The consumer must outlive the source.** Deleting it, or letting an `inactive_threshold` expire, while a source is bound to it does not stop the pipelines: the server answers their pull requests with "no responders", which is also what a JetStream outage produces, so the source retries indefinitely and warns. Recreate the consumer under the same name to resume, or disable the source. @@ -233,7 +233,7 @@ Measured at four pipelines with a 10 s commit timeout: `1000` gave 97 documents | Setting | Where | Effect | | --- | --- | --- | -| `num_pipelines` | source config | Pipelines bound to the consumer. Scales cleanly on large messages: 52 → 147 MiB/s from one to four pipelines at 512 KiB. On small messages the acknowledgment path caps the aggregate, and it is worth only about 1.15× from one to sixteen — there, parallelism has to come from more consumers. | +| `num_pipelines` | source config | Pipelines bound to the consumer. Scales cleanly on large messages: 52 → 147 MiB/s from one to four pipelines at 512 KiB. On small messages (1 KiB) the acknowledgment path caps the aggregate, and it is worth only about 1.15× from one to sixteen — there, parallelism has to come from more consumers. | | [`commit_timeout_secs`](index-config.md#indexing-settings) | index config | Sets the size of the always-ack-pending window, so it interacts with `max_ack_pending` and with `ack_wait`. | | `QW_NATS_PULL_MAX_BYTES_PER_BATCH` | env, default 10 MiB | Bounds the bytes the server may push per pull request. It must stay well under the server's per-connection `max_pending` (64 MiB by default), or the server declares a slow consumer and closes the connection; the messages are already counted as delivered, so the pipeline then idles for a full `ack_wait`. It must also stay above the server's `max_payload`, or larger messages are never delivered, and at or below the consumer's `max_bytes`; the source warns at start otherwise. 10 MiB and 20 MiB measure identically. | | `QW_NATS_PULL_MAX_MESSAGES_PER_BATCH` | env, default 100,000 | Messages per pull request. Can be used along `QW_NATS_PULL_MAX_BYTES_PER_BATCH` to limit a batch if there is too much contention around acknowledgements. The client buffers eight such batches per subscription; messages beyond that are dropped client-side and come back after `ack_wait`. | From f2d60fb505eda9594b18d810a988b162268c03e6 Mon Sep 17 00:00:00 2001 From: tprevot Date: Wed, 9 Sep 2026 23:11:51 +0200 Subject: [PATCH 21/21] feat: improve tracing --- docs/configuration/source-config.md | 2 + quickwit/quickwit-common/src/tracing_utils.rs | 16 ++++- .../src/source/nats_source.rs | 71 +++++++++++++++---- 3 files changed, 76 insertions(+), 13 deletions(-) diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 7046843f3c8..44fcfc7981f 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -253,6 +253,8 @@ One confirmed acknowledgment per message is the source's dominant broker cost, a Being durable, the consumer is observable through NATS's own monitoring (`nats consumer info`, exporters): `num_pending` is the indexing lag and `num_ack_pending` the in-flight window, both available even while the pipelines are down. The source also reports `num_pending_acks` (messages indexed but whose split is not published yet) in its observable state. +When a message carries a W3C `traceparent` header and the indexers export their traces (see [distributed tracing](../distributed-tracing/plug-quickwit-to-jaeger.md)), the publisher's trace extends to Quickwit: a `process_nats_message` span covers the message's processing until it is acknowledged. + **NATS source parameters** | Property | Description | Default value | diff --git a/quickwit/quickwit-common/src/tracing_utils.rs b/quickwit/quickwit-common/src/tracing_utils.rs index c1711bb6b4e..4fb3774cfc6 100644 --- a/quickwit/quickwit-common/src/tracing_utils.rs +++ b/quickwit/quickwit-common/src/tracing_utils.rs @@ -20,10 +20,11 @@ use std::str::FromStr; +pub use opentelemetry::Context; +use opentelemetry::global; pub use opentelemetry::propagation::Extractor; use opentelemetry::propagation::Injector; use opentelemetry::trace::TraceContextExt; -use opentelemetry::{Context, global}; use tonic::Status; use tonic::metadata::{KeyAndValueRef, MetadataKey, MetadataMap, MetadataValue}; use tonic::service::Interceptor; @@ -102,6 +103,19 @@ pub fn set_span_parent(span: &Span, parent_context: Context) { let _ = span.set_parent(parent_context); } +/// The OpenTelemetry context of `span`, to parent or link spans created +/// later, possibly after `span` closed. +pub fn span_context(span: &Span) -> Context { + span.context() +} + +/// Links `span` to `linked_span`: a causal relation that is not a +/// parent/child one, e.g. a batch and each message it carries. No-op when +/// `linked_span` has no OpenTelemetry context. +pub fn link_span(span: &Span, linked_span: &Span) { + span.add_link(linked_span.context().span().span_context().clone()); +} + /// Extracts a W3C trace context from incoming gRPC request metadata and /// installs it as the parent of the currently-active tracing span. Use this /// at the entry of a gRPC handler that is itself wrapped in a diff --git a/quickwit/quickwit-indexing/src/source/nats_source.rs b/quickwit/quickwit-indexing/src/source/nats_source.rs index 01407691723..eb9051b03d0 100644 --- a/quickwit/quickwit-indexing/src/source/nats_source.rs +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -25,7 +25,9 @@ //! message and is what bounds throughput on small messages. //! //! When a message carries a W3C `traceparent` header, the source stitches the -//! processing of the message into the publisher's distributed trace. +//! processing and the acknowledgment of the message into the publisher's +//! distributed trace, and links the processing span to the batch that +//! carried the message. use std::collections::BTreeMap; use std::fmt; @@ -46,14 +48,14 @@ use async_trait::async_trait; use bytesize::ByteSize; use futures::{FutureExt, StreamExt}; use quickwit_actors::ActorExitStatus; -use quickwit_common::tracing_utils::{self, Extractor}; +use quickwit_common::tracing_utils::{self, Context as TraceContext, Extractor}; use quickwit_config::{NatsSourceAuth, NatsSourceParams}; use quickwit_metastore::checkpoint::{PartitionId, SourceCheckpoint}; use quickwit_proto::metastore::SourceType; use quickwit_proto::types::Position; use serde_json::{Value as JsonValue, json}; use tokio::time; -use tracing::{Span, debug, info, warn}; +use tracing::{Instrument, Span, debug, info, warn}; use crate::source::{ BATCH_NUM_BYTES_LIMIT, BatchBuilder, EMIT_BATCHES_TIMEOUT, Source, SourceContext, @@ -96,11 +98,20 @@ pub struct NatsSource { /// Messages delivered but not published yet, keyed by delivery counter. /// Bounded by the consumer's `max_ack_pending`: the server stops delivering /// when too many messages are unacknowledged. - pending_acks: BTreeMap, + pending_acks: BTreeMap, state: NatsSourceState, } +/// A delivered message whose split is not published yet. +#[derive(Clone)] +struct PendingAck { + ack_subject: Subject, + /// Context of the message's processing span, when the publisher propagated + /// a trace: the ack is reported as its child. + trace_context: Option, +} + impl fmt::Debug for NatsSource { fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { formatter @@ -187,8 +198,11 @@ impl NatsSource { .info() .map_err(|error| anyhow!("failed to parse NATS message metadata: {error}"))? .stream_sequence; - let _span_guard = remote_parented_span(&message, stream_sequence, &self.source_runtime) - .map(Span::entered); + let batch_span = Span::current(); + let message_span = + remote_parented_span(&message, stream_sequence, &self.source_runtime, &batch_span); + let trace_context = message_span.as_ref().map(tracing_utils::span_context); + let _span_guard = message_span.map(Span::entered); let Some(ack_subject) = message.message.reply else { bail!("NATS message carries no reply subject to acknowledge it on"); }; @@ -215,7 +229,11 @@ impl NatsSource { Position::offset(self.delivery_counter), ) .context("failed to record partition delta")?; - self.pending_acks.insert(self.delivery_counter, ack_subject); + let pending_ack = PendingAck { + ack_subject, + trace_context, + }; + self.pending_acks.insert(self.delivery_counter, pending_ack); self.state.num_bytes_processed += num_bytes; self.state.num_messages_processed += 1; @@ -226,7 +244,7 @@ impl NatsSource { /// Sends server-confirmed acknowledgments ("double acks") for the messages up /// to `published_up_to`. async fn ack_up_to(&mut self, published_up_to: u64) { - let acks: Vec<(u64, Subject)> = self + let acks: Vec<(u64, PendingAck)> = self .pending_acks .range(..=published_up_to) .map(|(delivery_counter, pending_ack)| (*delivery_counter, pending_ack.clone())) @@ -235,9 +253,17 @@ impl NatsSource { return; } let nats_client = &self.nats_client; + let source_runtime = &self.source_runtime; let mut ack_results = futures::stream::iter(acks) - .map(async |(delivery_counter, ack_subject)| { - let ack_result = send_ack(nats_client, ack_subject).await; + .map(async |(delivery_counter, pending_ack)| { + let ack_span = ack_span(pending_ack.trace_context, source_runtime); + let ack_result = send_ack(nats_client, pending_ack.ack_subject) + .instrument(ack_span.clone()) + .await; + if let Err(error) = &ack_result { + ack_span.record("otel.status_code", "ERROR"); + ack_span.record("otel.status_description", tracing::field::display(error)); + } (delivery_counter, ack_result) }) .buffer_unordered(MAX_CONCURRENT_ACK_REQUESTS); @@ -612,12 +638,14 @@ impl Extractor for NatsHeaderExtractor<'_> { /// Builds a span parented on the publisher's trace when the message carries /// a W3C `traceparent` header, stitching the processing of the message into -/// the publisher's distributed trace. Messages without a propagated context -/// cost nothing: no span is created. +/// the publisher's distributed trace, and links it to `batch_span`, the +/// batch carrying the message. Messages without a propagated context cost +/// nothing: no span is created. fn remote_parented_span( message: &jetstream::Message, stream_sequence: u64, source_runtime: &SourceRuntime, + batch_span: &Span, ) -> Option { let headers = message.headers.as_ref()?; let parent_context = tracing_utils::extract_remote_context(&NatsHeaderExtractor(headers))?; @@ -629,9 +657,28 @@ fn remote_parented_span( stream_sequence, ); tracing_utils::set_span_parent(&span, parent_context); + tracing_utils::link_span(&span, batch_span); Some(span) } +/// Reports the ack of a message as a child of its processing span, so the +/// publisher's trace extends to durability. Disabled when the publisher +/// propagated no trace. +fn ack_span(trace_context: Option, source_runtime: &SourceRuntime) -> Span { + let Some(trace_context) = trace_context else { + return Span::none(); + }; + let span = tracing::info_span!( + "ack_nats_message", + index_id = %source_runtime.index_id(), + source_id = %source_runtime.source_id(), + otel.status_code = tracing::field::Empty, + otel.status_description = tracing::field::Empty, + ); + tracing_utils::set_span_parent(&span, trace_context); + span +} + async fn connect_nats(params: &NatsSourceParams) -> anyhow::Result { let subscription_capacity = SUBSCRIPTION_CAPACITY_IN_PULL_BATCHES * pull_max_messages_per_batch();