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..a2e88600caa 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -127,6 +127,21 @@ services: timeout: 10s retries: 100 + nats: + 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..44fcfc7981f 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 @@ -183,6 +183,105 @@ 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 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. + +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 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. + +**`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 +``` + +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. + +**`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. 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. + +**`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. + +**`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. + +#### 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 (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`. | + +#### Scaling + +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. +- 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. + +#### 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. + +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 | +| --- | --- | --- | +| `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 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)* + +```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. @@ -216,7 +315,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 new file mode 100644 index 00000000000..42e71e75691 --- /dev/null +++ b/docs/ingest-data/nats.md @@ -0,0 +1,267 @@ +--- +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 +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); +} + +/// 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 @@ -87,7 +123,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 +185,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/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..afcedc3c747 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 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 + /// `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..e00402aafcd 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,12 @@ 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..bf8aaaffc60 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 } @@ -83,6 +84,8 @@ kinesis = [ "quickwit-aws/kinesis", ] kinesis-localstack-tests = [] +nats = ["dep:async-nats"] +nats-broker-tests = ["nats"] pulsar = ["dep:pulsar"] pulsar-broker-tests = [] queue-sources = [] @@ -117,6 +120,8 @@ 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 } prost = { workspace = true } diff --git a/quickwit/quickwit-indexing/src/source/mod.rs b/quickwit/quickwit-indexing/src/source/mod.rs index 27d8958d145..69fd59ee76e 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; @@ -486,6 +490,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); @@ -553,6 +559,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(()), } } @@ -605,6 +622,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, @@ -704,7 +722,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); @@ -862,6 +884,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..eb9051b03d0 --- /dev/null +++ b/quickwit/quickwit-indexing/src/source/nats_source.rs @@ -0,0 +1,1563 @@ +// 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 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. That +//! shape requires `AckPolicy::Explicit`, which costs one acknowledgment per +//! message and is what bounds throughput on small messages. +//! +//! When a message carries a W3C `traceparent` header, the source stitches the +//! 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; +use std::path::PathBuf; +use std::str::FromStr; +use std::time::{Duration, Instant}; + +use anyhow::{Context as _, anyhow, bail, ensure}; +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, HeaderName, Subject, jetstream}; +use async_trait::async_trait; +use bytesize::ByteSize; +use futures::{FutureExt, StreamExt}; +use quickwit_actors::ActorExitStatus; +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::{Instrument, Span, debug, info, warn}; + +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, +} + +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, + /// 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, +} + +/// 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 + .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(); + 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_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) + .messages() + .await + .context("failed to subscribe to NATS consumer messages")?; + + // 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, + source_params, + nats_client, + message_stream, + consumer_name, + partition_id, + delivery_counter, + 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 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"); + }; + 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")?; + 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; + + 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, PendingAck)> = self + .pending_acks + .range(..=published_up_to) + .map(|(delivery_counter, pending_ack)| (*delivery_counter, pending_ack.clone())) + .collect(); + if acks.is_empty() { + 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, 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); + + 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(()) => { + 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"); + } + + /// 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) -> 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, + 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; + } + + num_naks + } +} + +#[async_trait] +impl Source for NatsSource { + #[tracing::instrument(skip(source_sink, ctx))] + 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, retrying"); + 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) + } + + #[tracing::instrument(skip(checkpoint, ctx))] + 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(self.ack_up_to(published_up_to)).await; + 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 + // 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<()> { + // Any ack still pending at this point was not completed by a drain: + // its message is redelivered after the consumer's `ack_wait`. + // + // 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 { + 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"); + } + 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(()) + } + + 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, + }) + } +} + +/// 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. +/// +/// 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`. +const DEFAULT_PULL_MAX_BYTES_PER_BATCH: ByteSize = ByteSize::mib(10); + +fn pull_max_bytes_per_batch() -> ByteSize { + quickwit_common::get_from_env_cached!( + ByteSize, + "QW_NATS_PULL_MAX_BYTES_PER_BATCH", + DEFAULT_PULL_MAX_BYTES_PER_BATCH, + false + ) +} + +/// 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" + ); + } +} + +/// 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 + ) +} + +/// 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; + +/// 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. +fn is_transient_stream_error(error: &MessagesError) -> bool { + matches!( + error.kind(), + MessagesErrorKind::Pull | MessagesErrorKind::NoResponders + ) +} + +/// 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; + +/// 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, +) -> 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) +} + +/// Lets `quickwit_common::tracing_utils` read the W3C trace context from +/// NATS message headers. +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() + } +} + +/// 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, 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))?; + 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, + ); + 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(); + let mut connect_options = ConnectOptions::new().subscription_capacity(subscription_capacity); + 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) +} + +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)) +} + +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 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 + // 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()); + } + + #[test] + fn transient_stream_error_classification() { + let transient_kinds = [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"))] +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_common::rand::append_random_suffix; + use quickwit_config::{SourceConfig, SourceInputFormat, SourceParams}; + use quickwit_metastore::checkpoint::{PartitionDelta, SourceCheckpointDelta}; + use quickwit_metastore::metastore_for_test; + use quickwit_proto::metastore::MetastoreServiceClient; + use quickwit_proto::types::{IndexUid, PipelineUid}; + + 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) + } + + /// 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 + .state + .get("num_messages_processed") + .unwrap() + .as_u64() + .unwrap(); + 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 { + 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_with_ack_wait( + jetstream_ctx: &jetstream::Context, + stream: &str, + consumer_name: &str, + ack_wait: Duration, + ) { + 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, + ack_wait, + ..Default::default() + }, + stream, + ) + .await + .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 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, Duration::from_secs(60)).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(); + 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; + } + + /// 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 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 { + 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; + 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 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 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_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(); + 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; + + 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, diff --git a/quickwit/quickwit-serve/src/index_api/rest_handler.rs b/quickwit/quickwit-serve/src/index_api/rest_handler.rs index bc4c8d9b105..1324508ff68 100644 --- a/quickwit/quickwit-serve/src/index_api/rest_handler.rs +++ b/quickwit/quickwit-serve/src/index_api/rest_handler.rs @@ -1100,8 +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 or Kafka \ - sources" + "Quickwit currently supports multiple pipelines only for GCP PubSub, Kafka, or \ + NATS sources" )); } {