From 6fa3c295fa9518712a8dd3cb04476850d45ef24c Mon Sep 17 00:00:00 2001 From: Remi Dettai Date: Wed, 31 Jul 2024 12:18:09 +0200 Subject: [PATCH] Add File source with SQS notifications (#5148) * Refactor file source * More explicit reponse for read_batch * Fix clippy * Remove batch logic from DocFileReader * Address styling comments * Replace FileSourceParams path with URI * Additional URI related cleanup * Fix unit tests to now use the URI * Add queue source with SQS implementation * Fix the publish_token strategy * Fix never-used warnings * Fix unit tests * Abort visibility task after acknowledging * Address smaller review comments * Shorten visibility extension task * Fix pulsar tests * Adjustments after rebase * Move object backed source to file source * Simpler flow for adhoc file processing * Fix tests and refactor batch creation to BatchReader * Add max_messages param to Queue.receive * Move use_shard_api to the metastore crate * Dedup within batches * Improve visibility task * Re-acquire partitions aggressively * Address simpler review comments * Add test for visibility actor (failing) * Fix visibility actor drop * Fix reader edge case * Add end to end tests * Improve integration test scenario * Chunk receive future to avoid hanging actor * Improve error handling * Fix flaky test * New SourceConfig format with notifications field * Improvements to error handling * Clarification of Queue contract * Address new round of review comments * Remove SqsSource for now * Fix panic * Revert to forbidding any adhoc file source through the API * Add docs * Fix panic on empty file * Documentation improvments * Improve documentation * Improve error handling code and associated docs * Nitpic and TODO cleanup * Add tip about ingests directly from object stores * Ack notifications of undesired type * Add docs about situations where messages require a DLQ * Fix integ test after rebase --- docker-compose.yml | 4 +- docs/assets/sqs-file-source.tf | 134 +++++ docs/configuration/source-config.md | 55 +- docs/ingest-data/ingest-local-file.md | 6 + docs/ingest-data/sqs-files.md | 248 ++++++++ docs/overview/concepts/indexing.md | 11 +- quickwit/Cargo.lock | 32 +- quickwit/Cargo.toml | 1 + quickwit/quickwit-actors/src/actor_context.rs | 7 +- quickwit/quickwit-aws/Cargo.toml | 2 + quickwit/quickwit-aws/src/error.rs | 174 +++--- quickwit/quickwit-cli/Cargo.toml | 3 + quickwit/quickwit-cli/src/source.rs | 9 +- quickwit/quickwit-cli/src/tool.rs | 10 +- quickwit/quickwit-cli/tests/cli.rs | 61 +- quickwit/quickwit-cli/tests/helpers.rs | 32 +- quickwit/quickwit-common/src/fs.rs | 2 +- quickwit/quickwit-config/src/lib.rs | 15 +- .../quickwit-config/src/source_config/mod.rs | 199 +++++-- .../src/source_config/serialize.rs | 38 +- .../src/indexing_scheduler/mod.rs | 30 +- .../src/ingest/ingest_controller.rs | 2 + quickwit/quickwit-indexing/Cargo.toml | 12 +- .../src/actors/indexing_pipeline.rs | 17 +- .../quickwit-indexing/src/actors/publisher.rs | 2 +- .../src/source/doc_file_reader.rs | 561 ++++++++++++++++++ .../src/source/file_source.rs | 529 ++++++++--------- .../src/source/ingest/mod.rs | 9 +- .../src/source/kafka_source.rs | 118 ++-- quickwit/quickwit-indexing/src/source/mod.rs | 116 +++- .../src/source/pulsar_source.rs | 86 +-- .../src/source/queue_sources/coordinator.rs | 521 ++++++++++++++++ .../src/source/queue_sources/helpers.rs | 130 ++++ .../src/source/queue_sources/local_state.rs | 134 +++++ .../src/source/queue_sources/memory_queue.rs | 233 ++++++++ .../src/source/queue_sources/message.rs | 352 +++++++++++ .../src/source/queue_sources/mod.rs | 89 +++ .../src/source/queue_sources/shared_state.rs | 371 ++++++++++++ .../src/source/queue_sources/sqs_queue.rs | 468 +++++++++++++++ .../src/source/queue_sources/visibility.rs | 341 +++++++++++ .../src/source/stdin_source.rs | 133 +++++ .../quickwit-integration-tests/Cargo.toml | 9 + .../src/test_utils/cluster_sandbox.rs | 8 +- .../src/tests/mod.rs | 2 + .../src/tests/sqs_tests.rs | 167 ++++++ .../quickwit-lambda/src/indexer/handler.rs | 2 +- .../src/indexer/ingest/helpers.rs | 6 +- .../quickwit-lambda/src/indexer/ingest/mod.rs | 4 +- quickwit/quickwit-lambda/src/indexer/model.rs | 27 +- quickwit/quickwit-metastore/src/checkpoint.rs | 2 +- .../file_backed/file_backed_index/mod.rs | 11 +- .../file_backed/file_backed_index/shards.rs | 5 +- .../quickwit-metastore/src/metastore/mod.rs | 22 +- .../src/metastore/postgres/metastore.rs | 12 +- .../postgres/queries/shards/open.sql | 4 +- .../quickwit-metastore/src/tests/shard.rs | 31 + .../quickwit-metastore/src/tests/split.rs | 51 +- .../protos/quickwit/metastore.proto | 2 + .../codegen/quickwit/quickwit.metastore.rs | 5 + quickwit/quickwit-proto/src/metastore/mod.rs | 2 + .../src/index_api/rest_handler.rs | 68 ++- quickwit/quickwit-storage/src/storage.rs | 4 +- 62 files changed, 4968 insertions(+), 773 deletions(-) create mode 100644 docs/assets/sqs-file-source.tf create mode 100644 docs/ingest-data/sqs-files.md create mode 100644 quickwit/quickwit-indexing/src/source/doc_file_reader.rs create mode 100644 quickwit/quickwit-indexing/src/source/queue_sources/coordinator.rs create mode 100644 quickwit/quickwit-indexing/src/source/queue_sources/helpers.rs create mode 100644 quickwit/quickwit-indexing/src/source/queue_sources/local_state.rs create mode 100644 quickwit/quickwit-indexing/src/source/queue_sources/memory_queue.rs create mode 100644 quickwit/quickwit-indexing/src/source/queue_sources/message.rs create mode 100644 quickwit/quickwit-indexing/src/source/queue_sources/mod.rs create mode 100644 quickwit/quickwit-indexing/src/source/queue_sources/shared_state.rs create mode 100644 quickwit/quickwit-indexing/src/source/queue_sources/sqs_queue.rs create mode 100644 quickwit/quickwit-indexing/src/source/queue_sources/visibility.rs create mode 100644 quickwit/quickwit-indexing/src/source/stdin_source.rs create mode 100644 quickwit/quickwit-integration-tests/src/tests/sqs_tests.rs diff --git a/docker-compose.yml b/docker-compose.yml index 58b9d8b99f7..0667d9ac434 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -27,7 +27,7 @@ networks: services: localstack: - image: localstack/localstack:${LOCALSTACK_VERSION:-2.3.2} + image: localstack/localstack:${LOCALSTACK_VERSION:-3.5.0} container_name: localstack ports: - "${MAP_HOST_LOCALSTACK:-127.0.0.1}:4566:4566" @@ -37,7 +37,7 @@ services: - all - localstack environment: - SERVICES: kinesis,s3 + SERVICES: kinesis,s3,sqs PERSISTENCE: 1 volumes: - .localstack:/etc/localstack/init/ready.d diff --git a/docs/assets/sqs-file-source.tf b/docs/assets/sqs-file-source.tf new file mode 100644 index 00000000000..ffd348c1193 --- /dev/null +++ b/docs/assets/sqs-file-source.tf @@ -0,0 +1,134 @@ +terraform { + required_version = "1.7.5" + required_providers { + aws = { + source = "hashicorp/aws" + version = "~> 5.39.1" + } + } +} + +provider "aws" { + region = "us-east-1" + default_tags { + tags = { + provisioner = "terraform" + author = "Quickwit" + } + } +} + +locals { + sqs_notification_queue_name = "qw-tuto-s3-event-notifications" + source_bucket_name = "qw-tuto-source-bucket" +} + +resource "aws_s3_bucket" "file_source" { + bucket_prefix = local.source_bucket_name + force_destroy = true +} + +data "aws_iam_policy_document" "sqs_notification" { + statement { + effect = "Allow" + + principals { + type = "*" + identifiers = ["*"] + } + + actions = ["sqs:SendMessage"] + resources = ["arn:aws:sqs:*:*:${local.sqs_notification_queue_name}"] + + condition { + test = "ArnEquals" + variable = "aws:SourceArn" + values = [aws_s3_bucket.file_source.arn] + } + } +} + + +resource "aws_sqs_queue" "s3_events" { + name = local.sqs_notification_queue_name + policy = data.aws_iam_policy_document.sqs_notification.json + + redrive_policy = jsonencode({ + deadLetterTargetArn = aws_sqs_queue.s3_events_deadletter.arn + maxReceiveCount = 5 + }) +} + +resource "aws_sqs_queue" "s3_events_deadletter" { + name = "${locals.sqs_notification_queue_name}-deadletter" +} + +resource "aws_sqs_queue_redrive_allow_policy" "s3_events_deadletter" { + queue_url = aws_sqs_queue.s3_events_deadletter.id + + redrive_allow_policy = jsonencode({ + redrivePermission = "byQueue", + sourceQueueArns = [aws_sqs_queue.s3_events.arn] + }) +} + +resource "aws_s3_bucket_notification" "bucket_notification" { + bucket = aws_s3_bucket.file_source.id + + queue { + queue_arn = aws_sqs_queue.s3_events.arn + events = ["s3:ObjectCreated:*"] + } +} + +data "aws_iam_policy_document" "quickwit_node" { + statement { + effect = "Allow" + actions = [ + "sqs:ReceiveMessage", + "sqs:DeleteMessage", + "sqs:ChangeMessageVisibility", + "sqs:GetQueueAttributes", + ] + resources = [aws_sqs_queue.s3_events.arn] + } + statement { + effect = "Allow" + actions = ["s3:GetObject"] + resources = ["${aws_s3_bucket.file_source.arn}/*"] + } +} + +resource "aws_iam_user" "quickwit_node" { + name = "quickwit-filesource-tutorial" + path = "/system/" +} + +resource "aws_iam_user_policy" "quickwit_node" { + name = "quickwit-filesource-tutorial" + user = aws_iam_user.quickwit_node.name + policy = data.aws_iam_policy_document.quickwit_node.json +} + +resource "aws_iam_access_key" "quickwit_node" { + user = aws_iam_user.quickwit_node.name +} + +output "source_bucket_name" { + value = aws_s3_bucket.file_source.bucket + +} + +output "notification_queue_url" { + value = aws_sqs_queue.s3_events.id +} + +output "quickwit_node_access_key_id" { + value = aws_iam_access_key.quickwit_node.id + sensitive = true +} + +output "quickwit_node_secret_access_key" { + value = aws_iam_access_key.quickwit_node.secret + sensitive = true +} diff --git a/docs/configuration/source-config.md b/docs/configuration/source-config.md index 479c97e2365..83bdace6f96 100644 --- a/docs/configuration/source-config.md +++ b/docs/configuration/source-config.md @@ -29,15 +29,62 @@ The source type designates the kind of source being configured. As of version 0. The source parameters indicate how to connect to a data store and are specific to the source type. -### File source (CLI only) +### File source -A file source reads data from a local file. The file must consist of JSON objects separated by a newline (NDJSON). -As of version 0.5, a file source can only be ingested with the [CLI command](/docs/reference/cli.md#tool-local-ingest). Compressed files (bz2, gzip, ...) and remote files (Amazon S3, HTTP, ...) are not supported. +A file source reads data from files containing JSON objects separated by newlines (NDJSON). Gzip compression is supported provided that the file name ends with the `.gz` suffix. + +#### Ingest a single file (CLI only) + +To ingest a specific file, run the indexing directly in an adhoc CLI process with: + +```bash +./quickwit tool local-ingest --index --input-path +``` + +Both local and object files are supported, provided that the environment is configured with the appropriate permissions. A tutorial is available [here](/docs/ingest-data/ingest-local-file.md). + +#### Notification based file ingestion (beta) + +Quickwit can automatically ingest all new files that are uploaded to an S3 bucket. This requires creating and configuring an [SQS notification queue](https://docs.aws.amazon.com/AmazonS3/latest/userguide/ways-to-add-notification-config-to-bucket.html). A complete example can be found [in this tutorial](/docs/ingest-data/sqs-files.md). + + +The `notifications` parameter takes an array of notification settings. Currently one notifier can be configured per source and only the SQS notification `type` is supported. + +Required fields for the SQS `notifications` parameter items: +- `type`: `sqs` +- `queue_url`: complete URL of the SQS queue (e.g `https://sqs.us-east-1.amazonaws.com/123456789012/queue-name`) +- `message_type`: format of the message payload, either + - `s3_notification`: an [S3 event notification](https://docs.aws.amazon.com/AmazonS3/latest/userguide/EventNotifications.html) + - `raw_uri`: a message containing just the file object URI (e.g. `s3://mybucket/mykey`) + +*Adding a file source with SQS notifications to an index with the [CLI](../reference/cli.md#source)* ```bash -./quickwit tool local-ingest --input-path +cat << EOF > source-config.yaml +version: 0.8 +source_id: my-sqs-file-source +source_type: file +num_pipelines: 2 +params: + notifications: + - type: sqs + queue_url: https://sqs.us-east-1.amazonaws.com/123456789012/queue-name + message_type: s3_notification +EOF +./quickwit source create --index my-index --source-config source-config.yaml ``` +:::note + +- Quickwit does not automatically delete the source files after a successful ingestion. You can use [S3 object expiration](https://docs.aws.amazon.com/AmazonS3/latest/userguide/lifecycle-expire-general-considerations.html) to configure how long they should be retained in the bucket. +- Configure the notification to only forward events of type `s3:ObjectCreated:*`. Other events are acknowledged by the source without further processing and an warning is logged. +- We strongly recommend using a [dead letter queue](https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-dead-letter-queues.html) to receive all messages that couldn't be processed by the file source. A `maxReceiveCount` of 5 is a good default value. Here are some common situations where the notification message ends up in the dead letter queue: + - the notification message could not be parsed (e.g it is not a valid S3 notification) + - the file was not found + - the file is corrupted (e.g unexpected compression) + +::: + ### Ingest API source An ingest API source reads data from the [Ingest API](/docs/reference/rest-api.md#ingest-data-into-an-index). This source is automatically created at the index creation and cannot be deleted nor disabled. diff --git a/docs/ingest-data/ingest-local-file.md b/docs/ingest-data/ingest-local-file.md index 2a5b1bced03..6eb37e7c3eb 100644 --- a/docs/ingest-data/ingest-local-file.md +++ b/docs/ingest-data/ingest-local-file.md @@ -72,6 +72,12 @@ Clearing local cache directory... ✔ Documents successfully indexed. ``` +:::tip + +Object store URIs like `s3://mybucket/mykey.json` are also supported as `--input-path`, provided that your environment is configured with the appropriate permissions. + +::: + ## Tear down resources (optional) That's it! You can now tear down the resources you created. You can do so by running the following command: diff --git a/docs/ingest-data/sqs-files.md b/docs/ingest-data/sqs-files.md new file mode 100644 index 00000000000..ebca49629d7 --- /dev/null +++ b/docs/ingest-data/sqs-files.md @@ -0,0 +1,248 @@ +--- +title: S3 with SQS notifications +description: A short tutorial describing how to set up Quickwit to ingest data from S3 files using an SQS notifier +tags: [s3, sqs, integration] +icon_url: /img/tutorials/file-ndjson.svg +sidebar_position: 5 +--- + +In this tutorial, we describe how to set up Quickwit to ingest data from S3 +with bucket notification events flowing through SQS. We will first create the +AWS resources (S3 bucket, SQS queue, notifications) using terraform. We will +then configure the Quickwit index and file source. Finally we will send some +data to the source bucket and verify that it gets indexed. + +## AWS resources + +The complete terraform script can be downloaded [here](../assets/sqs-file-source.tf). + +First, create the bucket that will receive the source data files (NDJSON format): + +``` +resource "aws_s3_bucket" "file_source" { + bucket_prefix = "qw-tuto-source-bucket" +} +``` + +Then setup the SQS queue that will carry the notifications when files are added +to the bucket. The queue is configured with a policy that allows the source +bucket to write the S3 notification messages to it. Also create a dead letter +queue (DLQ) to receive the messages that couldn't be processed by the file +source (e.g corrupted files). Messages are moved to the DLQ after 5 indexing +attempts. + +``` +locals { + sqs_notification_queue_name = "qw-tuto-s3-event-notifications" +} + +data "aws_iam_policy_document" "sqs_notification" { + statement { + effect = "Allow" + + principals { + type = "*" + identifiers = ["*"] + } + + actions = ["sqs:SendMessage"] + resources = ["arn:aws:sqs:*:*:${local.sqs_notification_queue_name}"] + + condition { + test = "ArnEquals" + variable = "aws:SourceArn" + values = [aws_s3_bucket.file_source.arn] + } + } +} + +resource "aws_sqs_queue" "s3_events_deadletter" { + name = "${locals.sqs_notification_queue_name}-deadletter" +} + +resource "aws_sqs_queue" "s3_events" { + name = local.sqs_notification_queue_name + policy = data.aws_iam_policy_document.sqs_notification.json + + redrive_policy = jsonencode({ + deadLetterTargetArn = aws_sqs_queue.s3_events_deadletter.arn + maxReceiveCount = 5 + }) +} + +resource "aws_sqs_queue_redrive_allow_policy" "s3_events_deadletter" { + queue_url = aws_sqs_queue.s3_events_deadletter.id + + redrive_allow_policy = jsonencode({ + redrivePermission = "byQueue", + sourceQueueArns = [aws_sqs_queue.s3_events.arn] + }) +} +``` + +Configure the bucket notification that writes messages to SQS each time a new +file is created in the source bucket: + +``` +resource "aws_s3_bucket_notification" "bucket_notification" { + bucket = aws_s3_bucket.file_source.id + + queue { + queue_arn = aws_sqs_queue.s3_events.arn + events = ["s3:ObjectCreated:*"] + } +} +``` + +:::note + +Only events of type `s3:ObjectCreated:*` are supported. Other types (e.g. +`ObjectRemoved`) are acknowledged and a warning is logged. + +::: + +The source needs to have access to both the notification queue and the source +bucket. The following policy document contains the minimum permissions required +by the source: + +``` +data "aws_iam_policy_document" "quickwit_node" { + statement { + effect = "Allow" + actions = [ + "sqs:ReceiveMessage", + "sqs:DeleteMessage", + "sqs:ChangeMessageVisibility", + "sqs:GetQueueAttributes", + ] + resources = [aws_sqs_queue.s3_events.arn] + } + statement { + effect = "Allow" + actions = ["s3:GetObject"] + resources = ["${aws_s3_bucket.file_source.arn}/*"] + } +} +``` + +Create the IAM user and credentials that will be used to +associate this policy to your local Quickwit instance: + +``` +resource "aws_iam_user" "quickwit_node" { + name = "quickwit-filesource-tutorial" + path = "/system/" +} + +resource "aws_iam_user_policy" "quickwit_node" { + name = "quickwit-filesource-tutorial" + user = aws_iam_user.quickwit_node.name + policy = data.aws_iam_policy_document.quickwit_node.json +} + +resource "aws_iam_access_key" "quickwit_node" { + user = aws_iam_user.quickwit_node.name +} +``` + + +:::warning + +We don't recommend using IAM user credentials for running Quickwit nodes in +production. This is just a simplified setup for the sake of the tutorial. When +running on EC2/ECS, attach the policy document to an IAM roles instead. + +::: + +Download the [complete terraform script](../assets/sqs-file-source.tf) and +deploy it using `terraform init` and `terraform apply`. After a successful +execution, the outputs required to configure Quickwit will be listed. You can +display the values of the sensitive outputs (key id and secret key) with: + + +```bash +terraform output quickwit_node_access_key_id +terraform output quickwit_node_secret_access_key +``` + +## Run Quickwit + +[Install Quickwit locally](/docs/get-started/installation), then in your install +directory, run Quickwit with the necessary access rights by replacing the +`` and `` with the +matching Terraform output values: + +```bash +AWS_ACCESS_KEY_ID= \ +AWS_SECRET_ACCESS_KEY= \ +AWS_REGION=us-east-1 \ +./quickwit run +``` + +## Configure the index and the source + +In another terminal, in the Quickwit install directory, create an index: + +```bash +cat << EOF > tutorial-sqs-file-index.yaml +version: 0.7 +index_id: tutorial-sqs-file +doc_mapping: + mode: dynamic +indexing_settings: + commit_timeout_secs: 30 +EOF + +./quickwit index create --index-config tutorial-sqs-file-index.yaml +``` + +Replacing `` with the corresponding Terraform output +value, create a file source for that index: + +```bash +cat << EOF > tutorial-sqs-file-source.yaml +version: 0.8 +source_id: sqs-filesource +source_type: file +num_pipelines: 2 +params: + notifications: + - type: sqs + queue_url: + message_type: s3_notification +EOF + +./quickwit source create --index tutorial-sqs-file --source-config tutorial-sqs-file-source.yaml +``` + +:::tip + +The `num_pipeline` configuration controls how many consumers will poll from the queue in parallel. Choose the number according to the indexer compute resources you want to dedicate to this source. As a rule of thumb, configure 1 pipeline for every 2 cores. + +::: + +## Ingest data + +We can now ingest data into Quickwit by uploading files to S3. If you have the +AWS CLI installed, run the following command, replacing `` +with the associated Terraform output: + +```bash +curl https://quickwit-datasets-public.s3.amazonaws.com/hdfs-logs-multitenants-10000.json | \ + aws s3 cp - s3:///hdfs-logs-multitenants-10000.json +``` + +If you prefer not to use the AWS CLI, you can also download the file and upload +it manually to the source bucket using the AWS console. + +Wait approximately 1 minute and the data should appear in the index: + +```bash +./quickwit index describe --index tutorial-sqs-file +``` + +## Tear down the resources + +The AWS resources instantiated in this tutorial don't incur any fixed costs, but +we still recommend deleting them when you are done. In the directory with the +Terraform script, run `terraform destroy`. diff --git a/docs/overview/concepts/indexing.md b/docs/overview/concepts/indexing.md index 4feba085b41..6dab81f392a 100644 --- a/docs/overview/concepts/indexing.md +++ b/docs/overview/concepts/indexing.md @@ -30,16 +30,7 @@ The disk space allocated to the split store is controlled by the config paramete ## Data sources -A data source designates the location and set of parameters that allow to connect to and ingest data from an external data store, which can be a file, a stream, or a database. Often, Quickwit simply refers to data sources as "sources". The indexing engine supports file-based and stream-based sources. Quickwit can insert data into an index from one or multiple sources, defined in the index config. - - -### File sources - -File sources are sources that read data from a file stored on the local file system. - -### Streaming sources - -Streaming sources are sources that read data from a streaming service such as Apache Kafka. As of version 0.2, Quickwit only supports Apache Kafka. Future versions of Quickwit will support additional streaming services such as Amazon Kinesis. +A data source designates the location and set of parameters that allow to connect to and ingest data from an external data store, which can be a file, a stream, or a database. Often, Quickwit simply refers to data sources as "sources". The indexing engine supports local adhoc file ingests using [the CLI](/docs/reference/cli#tool-local-ingest) and streaming sources (e.g. the Kafka source). Quickwit can insert data into an index from one or multiple sources. More details can be found [in the source configuration page](https://quickwit.io/docs/configuration/source-config). ## Checkpoint diff --git a/quickwit/Cargo.lock b/quickwit/Cargo.lock index 5551d880c3d..adf2ac01480 100644 --- a/quickwit/Cargo.lock +++ b/quickwit/Cargo.lock @@ -458,6 +458,28 @@ dependencies = [ "url", ] +[[package]] +name = "aws-sdk-sqs" +version = "1.36.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3587fbaf540d65337c2356ebf3f78fba160025b3d69634175f1ea3a7895738e9" +dependencies = [ + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "http 0.2.12", + "once_cell", + "regex-lite", + "tracing", +] + [[package]] name = "aws-sdk-sso" version = "1.36.0" @@ -5627,6 +5649,7 @@ dependencies = [ "aws-config", "aws-sdk-kinesis", "aws-sdk-s3", + "aws-sdk-sqs", "aws-smithy-async", "aws-smithy-runtime", "aws-types", @@ -5952,6 +5975,7 @@ dependencies = [ "async-compression", "async-trait", "aws-sdk-kinesis", + "aws-sdk-sqs", "bytes", "bytesize", "criterion", @@ -5988,6 +6012,7 @@ dependencies = [ "quickwit-storage", "rand 0.8.5", "rdkafka", + "regex", "reqwest", "serde", "serde_json", @@ -6045,6 +6070,7 @@ name = "quickwit-integration-tests" version = "0.8.0" dependencies = [ "anyhow", + "aws-sdk-sqs", "futures-util", "hyper 0.14.29", "itertools 0.13.0", @@ -6052,6 +6078,7 @@ dependencies = [ "quickwit-cli", "quickwit-common", "quickwit-config", + "quickwit-indexing", "quickwit-metastore", "quickwit-proto", "quickwit-rest-client", @@ -6063,6 +6090,7 @@ dependencies = [ "tokio", "tonic", "tracing", + "tracing-subscriber", ] [[package]] @@ -8471,9 +8499,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.39.1" +version = "1.39.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d040ac2b29ab03b09d4129c2f5bbd012a3ac2f79d38ff506a4bf8dd34b0eac8a" +checksum = "daa4fb1bc778bd6f04cbfc4bb2d06a7396a8f299dc33ea1900cedaa316f467b1" dependencies = [ "backtrace", "bytes", diff --git a/quickwit/Cargo.toml b/quickwit/Cargo.toml index 025f8e193ac..d6888a2df0f 100644 --- a/quickwit/Cargo.toml +++ b/quickwit/Cargo.toml @@ -279,6 +279,7 @@ aws-config = "1.5.4" aws-credential-types = { version = "1.2", features = ["hardcoded-credentials"] } aws-sdk-kinesis = "1.36" aws-sdk-s3 = "1.42" +aws-sdk-sqs = "1.36" aws-smithy-async = "1.2" aws-smithy-runtime = "1.6.2" aws-smithy-types = { version = "1.2", features = ["byte-stream-poll-next"] } diff --git a/quickwit/quickwit-actors/src/actor_context.rs b/quickwit/quickwit-actors/src/actor_context.rs index af7a8a3f7c9..0dbe1194bba 100644 --- a/quickwit/quickwit-actors/src/actor_context.rs +++ b/quickwit/quickwit-actors/src/actor_context.rs @@ -339,8 +339,11 @@ impl ActorContext { self.self_mailbox.try_send_message(msg) } - /// Schedules a message that will be sent to the high-priority - /// queue of the actor Mailbox once `after_duration` has elapsed. + /// Schedules a message that will be sent to the high-priority queue of the + /// actor Mailbox once `after_duration` has elapsed. + /// + /// Note that this holds a reference to the actor mailbox until the message + /// is actually sent. pub fn schedule_self_msg(&self, after_duration: Duration, message: M) where A: DeferableReplyHandler, diff --git a/quickwit/quickwit-aws/Cargo.toml b/quickwit/quickwit-aws/Cargo.toml index f67d36155b4..19bf0ceb6e0 100644 --- a/quickwit/quickwit-aws/Cargo.toml +++ b/quickwit/quickwit-aws/Cargo.toml @@ -14,6 +14,7 @@ license.workspace = true aws-config = { workspace = true } aws-sdk-kinesis = { workspace = true, optional = true } aws-sdk-s3 = { workspace = true } +aws-sdk-sqs = { workspace = true, optional = true } aws-smithy-async = { workspace = true } aws-smithy-runtime = { workspace = true } aws-types = { workspace = true } @@ -27,3 +28,4 @@ quickwit-common = { workspace = true } [features] kinesis = ["aws-sdk-kinesis"] +sqs = ["aws-sdk-sqs"] diff --git a/quickwit/quickwit-aws/src/error.rs b/quickwit/quickwit-aws/src/error.rs index e7c6dfdd077..ba2e620a27e 100644 --- a/quickwit/quickwit-aws/src/error.rs +++ b/quickwit/quickwit-aws/src/error.rs @@ -19,13 +19,6 @@ #![allow(clippy::match_like_matches_macro)] -#[cfg(feature = "kinesis")] -use aws_sdk_kinesis::operation::{ - create_stream::CreateStreamError, delete_stream::DeleteStreamError, - describe_stream::DescribeStreamError, get_records::GetRecordsError, - get_shard_iterator::GetShardIteratorError, list_shards::ListShardsError, - list_streams::ListStreamsError, merge_shards::MergeShardsError, split_shard::SplitShardError, -}; use aws_sdk_s3::error::SdkError; use aws_sdk_s3::operation::abort_multipart_upload::AbortMultipartUploadError; use aws_sdk_s3::operation::complete_multipart_upload::CompleteMultipartUploadError; @@ -109,89 +102,124 @@ impl AwsRetryable for HeadObjectError { } #[cfg(feature = "kinesis")] -impl AwsRetryable for GetRecordsError { - fn is_retryable(&self) -> bool { - match self { - GetRecordsError::KmsThrottlingException(_) => true, - GetRecordsError::ProvisionedThroughputExceededException(_) => true, - _ => false, +mod kinesis { + use aws_sdk_kinesis::operation::create_stream::CreateStreamError; + use aws_sdk_kinesis::operation::delete_stream::DeleteStreamError; + use aws_sdk_kinesis::operation::describe_stream::DescribeStreamError; + use aws_sdk_kinesis::operation::get_records::GetRecordsError; + use aws_sdk_kinesis::operation::get_shard_iterator::GetShardIteratorError; + use aws_sdk_kinesis::operation::list_shards::ListShardsError; + use aws_sdk_kinesis::operation::list_streams::ListStreamsError; + use aws_sdk_kinesis::operation::merge_shards::MergeShardsError; + use aws_sdk_kinesis::operation::split_shard::SplitShardError; + + use super::*; + + impl AwsRetryable for GetRecordsError { + fn is_retryable(&self) -> bool { + match self { + GetRecordsError::KmsThrottlingException(_) => true, + GetRecordsError::ProvisionedThroughputExceededException(_) => true, + _ => false, + } } } -} -#[cfg(feature = "kinesis")] -impl AwsRetryable for GetShardIteratorError { - fn is_retryable(&self) -> bool { - matches!( - self, - GetShardIteratorError::ProvisionedThroughputExceededException(_) - ) + impl AwsRetryable for GetShardIteratorError { + fn is_retryable(&self) -> bool { + matches!( + self, + GetShardIteratorError::ProvisionedThroughputExceededException(_) + ) + } } -} -#[cfg(feature = "kinesis")] -impl AwsRetryable for ListShardsError { - fn is_retryable(&self) -> bool { - matches!( - self, - ListShardsError::ResourceInUseException(_) | ListShardsError::LimitExceededException(_) - ) + impl AwsRetryable for ListShardsError { + fn is_retryable(&self) -> bool { + matches!( + self, + ListShardsError::ResourceInUseException(_) + | ListShardsError::LimitExceededException(_) + ) + } } -} -#[cfg(feature = "kinesis")] -impl AwsRetryable for CreateStreamError { - fn is_retryable(&self) -> bool { - matches!( - self, - CreateStreamError::ResourceInUseException(_) - | CreateStreamError::LimitExceededException(_) - ) + impl AwsRetryable for CreateStreamError { + fn is_retryable(&self) -> bool { + matches!( + self, + CreateStreamError::ResourceInUseException(_) + | CreateStreamError::LimitExceededException(_) + ) + } } -} -#[cfg(feature = "kinesis")] -impl AwsRetryable for DeleteStreamError { - fn is_retryable(&self) -> bool { - matches!( - self, - DeleteStreamError::ResourceInUseException(_) - | DeleteStreamError::LimitExceededException(_) - ) + impl AwsRetryable for DeleteStreamError { + fn is_retryable(&self) -> bool { + matches!( + self, + DeleteStreamError::ResourceInUseException(_) + | DeleteStreamError::LimitExceededException(_) + ) + } } -} -#[cfg(feature = "kinesis")] -impl AwsRetryable for DescribeStreamError { - fn is_retryable(&self) -> bool { - matches!(self, DescribeStreamError::LimitExceededException(_)) + impl AwsRetryable for DescribeStreamError { + fn is_retryable(&self) -> bool { + matches!(self, DescribeStreamError::LimitExceededException(_)) + } } -} -#[cfg(feature = "kinesis")] -impl AwsRetryable for ListStreamsError { - fn is_retryable(&self) -> bool { - matches!(self, ListStreamsError::LimitExceededException(_)) + impl AwsRetryable for ListStreamsError { + fn is_retryable(&self) -> bool { + matches!(self, ListStreamsError::LimitExceededException(_)) + } } -} -#[cfg(feature = "kinesis")] -impl AwsRetryable for MergeShardsError { - fn is_retryable(&self) -> bool { - matches!( - self, - MergeShardsError::ResourceInUseException(_) - | MergeShardsError::LimitExceededException(_) - ) + impl AwsRetryable for MergeShardsError { + fn is_retryable(&self) -> bool { + matches!( + self, + MergeShardsError::ResourceInUseException(_) + | MergeShardsError::LimitExceededException(_) + ) + } + } + + impl AwsRetryable for SplitShardError { + fn is_retryable(&self) -> bool { + matches!( + self, + SplitShardError::ResourceInUseException(_) + | SplitShardError::LimitExceededException(_) + ) + } } } -#[cfg(feature = "kinesis")] -impl AwsRetryable for SplitShardError { - fn is_retryable(&self) -> bool { - matches!( - self, - SplitShardError::ResourceInUseException(_) | SplitShardError::LimitExceededException(_) - ) +#[cfg(feature = "sqs")] +mod sqs { + use aws_sdk_sqs::operation::change_message_visibility::ChangeMessageVisibilityError; + use aws_sdk_sqs::operation::delete_message_batch::DeleteMessageBatchError; + use aws_sdk_sqs::operation::receive_message::ReceiveMessageError; + + use super::*; + + impl AwsRetryable for ReceiveMessageError { + fn is_retryable(&self) -> bool { + false + } + } + + impl AwsRetryable for DeleteMessageBatchError { + fn is_retryable(&self) -> bool { + false + } + } + + impl AwsRetryable for ChangeMessageVisibilityError { + fn is_retryable(&self) -> bool { + false + } } } diff --git a/quickwit/quickwit-cli/Cargo.toml b/quickwit/quickwit-cli/Cargo.toml index 36cfc371aa7..d403eef2922 100644 --- a/quickwit/quickwit-cli/Cargo.toml +++ b/quickwit/quickwit-cli/Cargo.toml @@ -92,6 +92,7 @@ release-feature-set = [ "quickwit-indexing/kafka", "quickwit-indexing/kinesis", "quickwit-indexing/pulsar", + "quickwit-indexing/sqs", "quickwit-indexing/vrl", "quickwit-storage/azure", "quickwit-storage/gcs", @@ -104,6 +105,7 @@ release-feature-vendored-set = [ "pprof", "quickwit-indexing/kinesis", "quickwit-indexing/pulsar", + "quickwit-indexing/sqs", "quickwit-indexing/vrl", "quickwit-indexing/vendored-kafka", "quickwit-storage/azure", @@ -116,6 +118,7 @@ release-macos-feature-vendored-set = [ "openssl-support", "quickwit-indexing/kinesis", "quickwit-indexing/pulsar", + "quickwit-indexing/sqs", "quickwit-indexing/vrl", "quickwit-indexing/vendored-kafka-macos", "quickwit-storage/azure", diff --git a/quickwit/quickwit-cli/src/source.rs b/quickwit/quickwit-cli/src/source.rs index 1a1948fdd99..ee90689ca94 100644 --- a/quickwit/quickwit-cli/src/source.rs +++ b/quickwit/quickwit-cli/src/source.rs @@ -744,7 +744,7 @@ mod tests { source_id: "foo-source".to_string(), num_pipelines: NonZeroUsize::new(1).unwrap(), enabled: true, - source_params: SourceParams::file("path/to/file"), + source_params: SourceParams::file_from_str("path/to/file").unwrap(), transform_config: None, input_format: SourceInputFormat::Json, }]; @@ -753,9 +753,10 @@ mod tests { source_type: "file".to_string(), enabled: "true".to_string(), }]; + let expected_uri = Uri::from_str("path/to/file").unwrap(); let expected_params = vec![ParamsRow { key: "filepath".to_string(), - value: JsonValue::String("path/to/file".to_string()), + value: JsonValue::String(expected_uri.to_string()), }]; let expected_checkpoint = vec![ CheckpointRow { @@ -820,12 +821,12 @@ mod tests { let expected_sources = [ SourceRow { source_id: "bar-source".to_string(), - source_type: "file".to_string(), + source_type: "stdin".to_string(), enabled: "true".to_string(), }, SourceRow { source_id: "foo-source".to_string(), - source_type: "file".to_string(), + source_type: "stdin".to_string(), enabled: "true".to_string(), }, ]; diff --git a/quickwit/quickwit-cli/src/tool.rs b/quickwit/quickwit-cli/src/tool.rs index b936cbc4897..c7ab1911205 100644 --- a/quickwit/quickwit-cli/src/tool.rs +++ b/quickwit/quickwit-cli/src/tool.rs @@ -173,7 +173,7 @@ pub fn build_tool_command() -> Command { pub struct LocalIngestDocsArgs { pub config_uri: Uri, pub index_id: IndexId, - pub input_path_opt: Option, + pub input_path_opt: Option, pub input_format: SourceInputFormat, pub overwrite: bool, pub vrl_script: Option, @@ -251,9 +251,7 @@ impl ToolCliCommand { .remove_one::("index") .expect("`index` should be a required arg."); let input_path_opt = if let Some(input_path) = matches.remove_one::("input-path") { - Uri::from_str(&input_path)? - .filepath() - .map(|path| path.to_path_buf()) + Some(Uri::from_str(&input_path)?) } else { None }; @@ -410,8 +408,8 @@ pub async fn local_ingest_docs_cli(args: LocalIngestDocsArgs) -> anyhow::Result< get_resolvers(&config.storage_configs, &config.metastore_configs); let mut metastore = metastore_resolver.resolve(&config.metastore_uri).await?; - let source_params = if let Some(filepath) = args.input_path_opt.as_ref() { - SourceParams::file(filepath) + let source_params = if let Some(uri) = args.input_path_opt.as_ref() { + SourceParams::file_from_uri(uri.clone()) } else { SourceParams::stdin() }; diff --git a/quickwit/quickwit-cli/tests/cli.rs b/quickwit/quickwit-cli/tests/cli.rs index 205fd778a85..524098537b6 100644 --- a/quickwit/quickwit-cli/tests/cli.rs +++ b/quickwit/quickwit-cli/tests/cli.rs @@ -26,7 +26,7 @@ use std::path::Path; use anyhow::Result; use clap::error::ErrorKind; -use helpers::{TestEnv, TestStorageType}; +use helpers::{uri_from_path, TestEnv, TestStorageType}; use quickwit_cli::checklist::ChecklistError; use quickwit_cli::cli::build_cli; use quickwit_cli::index::{ @@ -38,6 +38,7 @@ use quickwit_cli::tool::{ }; use quickwit_common::fs::get_cache_directory_path; use quickwit_common::rand::append_random_suffix; +use quickwit_common::uri::Uri; use quickwit_config::{RetentionPolicy, SourceInputFormat, CLI_SOURCE_ID}; use quickwit_metastore::{ ListSplitsRequestExt, MetastoreResolver, MetastoreServiceExt, MetastoreServiceStreamSplitsExt, @@ -62,11 +63,11 @@ async fn create_logs_index(test_env: &TestEnv) -> anyhow::Result<()> { create_index_cli(args).await } -async fn local_ingest_docs(input_path: &Path, test_env: &TestEnv) -> anyhow::Result<()> { +async fn local_ingest_docs(uri: Uri, test_env: &TestEnv) -> anyhow::Result<()> { let args = LocalIngestDocsArgs { config_uri: test_env.resource_files.config.clone(), index_id: test_env.index_id.clone(), - input_path_opt: Some(input_path.to_path_buf()), + input_path_opt: Some(uri), input_format: SourceInputFormat::Json, overwrite: false, clear_cache: true, @@ -75,6 +76,10 @@ async fn local_ingest_docs(input_path: &Path, test_env: &TestEnv) -> anyhow::Res local_ingest_docs_cli(args).await } +async fn local_ingest_log_docs(test_env: &TestEnv) -> anyhow::Result<()> { + local_ingest_docs(test_env.resource_files.log_docs.clone(), test_env).await +} + #[test] fn test_cmd_help() { let cmd = build_cli(); @@ -253,14 +258,17 @@ async fn test_ingest_docs_cli() { // Ensure cache directory is empty. let cache_directory_path = get_cache_directory_path(&test_env.data_dir_path); - assert!(cache_directory_path.read_dir().unwrap().next().is_none()); + let does_not_exist_uri = uri_from_path(&test_env.data_dir_path) + .join("file-does-not-exist.json") + .unwrap(); + // Ingest a non-existing file should fail. let args = LocalIngestDocsArgs { config_uri: test_env.resource_files.config, index_id: test_env.index_id, - input_path_opt: Some(test_env.data_dir_path.join("file-does-not-exist.json")), + input_path_opt: Some(does_not_exist_uri), input_format: SourceInputFormat::Json, overwrite: false, clear_cache: true, @@ -333,9 +341,7 @@ async fn test_cmd_search_aggregation() { test_env.start_server().await.unwrap(); create_logs_index(&test_env).await.unwrap(); - local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env) - .await - .unwrap(); + local_ingest_log_docs(&test_env).await.unwrap(); let aggregation: Value = json!( { @@ -433,9 +439,7 @@ async fn test_cmd_search_with_snippets() -> Result<()> { test_env.start_server().await.unwrap(); create_logs_index(&test_env).await.unwrap(); - local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env) - .await - .unwrap(); + local_ingest_log_docs(&test_env).await.unwrap(); // search with snippets let args = SearchIndexArgs { @@ -488,9 +492,7 @@ async fn test_search_index_cli() { sort_by_score: false, }; - local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env) - .await - .unwrap(); + local_ingest_log_docs(&test_env).await.unwrap(); let args = create_search_args("level:info"); @@ -601,9 +603,7 @@ async fn test_delete_index_cli_dry_run() { .unwrap(); assert!(metastore.index_exists(&index_id).await.unwrap()); - local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env) - .await - .unwrap(); + local_ingest_log_docs(&test_env).await.unwrap(); // On non-empty index let args = create_delete_args(true); @@ -627,9 +627,7 @@ async fn test_delete_index_cli() { test_env.start_server().await.unwrap(); create_logs_index(&test_env).await.unwrap(); - local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env) - .await - .unwrap(); + local_ingest_log_docs(&test_env).await.unwrap(); let args = DeleteIndexArgs { client_args: test_env.default_client_args(), @@ -653,9 +651,7 @@ async fn test_garbage_collect_cli_no_grace() { test_env.start_server().await.unwrap(); create_logs_index(&test_env).await.unwrap(); let index_uid = test_env.index_metadata().await.unwrap().index_uid; - local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env) - .await - .unwrap(); + local_ingest_log_docs(&test_env).await.unwrap(); let metastore = MetastoreResolver::unconfigured() .resolve(&test_env.metastore_uri) @@ -763,9 +759,7 @@ async fn test_garbage_collect_index_cli() { test_env.start_server().await.unwrap(); create_logs_index(&test_env).await.unwrap(); let index_uid = test_env.index_metadata().await.unwrap().index_uid; - local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env) - .await - .unwrap(); + local_ingest_log_docs(&test_env).await.unwrap(); let refresh_metastore = |metastore| async { // In this test we rely on the file backed metastore and @@ -915,9 +909,7 @@ async fn test_all_local_index() { .unwrap(); assert!(metadata_file_exists); - local_ingest_docs(test_env.resource_files.log_docs.as_path(), &test_env) - .await - .unwrap(); + local_ingest_log_docs(&test_env).await.unwrap(); let query_response = reqwest::get(format!( "http://127.0.0.1:{}/api/v1/{}/search?query=level:info", @@ -971,16 +963,21 @@ async fn test_all_with_s3_localstack_cli() { test_env.start_server().await.unwrap(); create_logs_index(&test_env).await.unwrap(); - let s3_path = upload_test_file( + let s3_uri = upload_test_file( test_env.storage_resolver.clone(), - test_env.resource_files.log_docs.clone(), + test_env + .resource_files + .log_docs + .filepath() + .unwrap() + .to_path_buf(), "quickwit-integration-tests", "sources/", &append_random_suffix("test-all--cli-s3-localstack"), ) .await; - local_ingest_docs(&s3_path, &test_env).await.unwrap(); + local_ingest_docs(s3_uri, &test_env).await.unwrap(); // Cli search let args = SearchIndexArgs { diff --git a/quickwit/quickwit-cli/tests/helpers.rs b/quickwit/quickwit-cli/tests/helpers.rs index 0a52f6b9792..67839f7368b 100644 --- a/quickwit/quickwit-cli/tests/helpers.rs +++ b/quickwit/quickwit-cli/tests/helpers.rs @@ -17,9 +17,8 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::borrow::Borrow; use std::fs; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::str::FromStr; use std::sync::Arc; @@ -114,8 +113,8 @@ pub struct TestResourceFiles { pub index_config: Uri, pub index_config_without_uri: Uri, pub index_config_with_retention: Uri, - pub log_docs: PathBuf, - pub wikipedia_docs: PathBuf, + pub log_docs: Uri, + pub wikipedia_docs: Uri, } /// A struct to hold few info about the test environment. @@ -192,8 +191,8 @@ pub enum TestStorageType { LocalFileSystem, } -fn uri_from_path(path: PathBuf) -> Uri { - Uri::from_str(&format!("file://{}", path.display())).unwrap() +pub fn uri_from_path(path: &Path) -> Uri { + Uri::from_str(path.to_str().unwrap()).unwrap() } /// Creates all necessary artifacts in a test environment. @@ -265,12 +264,12 @@ pub async fn create_test_env( .context("failed to parse cluster endpoint")?; let resource_files = TestResourceFiles { - config: uri_from_path(node_config_path), - index_config: uri_from_path(index_config_path), - index_config_without_uri: uri_from_path(index_config_without_uri_path), - index_config_with_retention: uri_from_path(index_config_with_retention_path), - log_docs: log_docs_path, - wikipedia_docs: wikipedia_docs_path, + config: uri_from_path(&node_config_path), + index_config: uri_from_path(&index_config_path), + index_config_without_uri: uri_from_path(&index_config_without_uri_path), + index_config_with_retention: uri_from_path(&index_config_with_retention_path), + log_docs: uri_from_path(&log_docs_path), + wikipedia_docs: uri_from_path(&wikipedia_docs_path), }; Ok(TestEnv { @@ -297,15 +296,14 @@ pub async fn upload_test_file( bucket: &str, prefix: &str, filename: &str, -) -> PathBuf { +) -> Uri { let test_data = tokio::fs::read(local_src_path).await.unwrap(); - let mut src_location: PathBuf = [r"s3://", bucket, prefix].iter().collect(); - let storage_uri = Uri::from_str(src_location.to_string_lossy().borrow()).unwrap(); + let src_location = format!("s3://{}/{}", bucket, prefix); + let storage_uri = Uri::from_str(&src_location).unwrap(); let storage = storage_resolver.resolve(&storage_uri).await.unwrap(); storage .put(&PathBuf::from(filename), Box::new(test_data)) .await .unwrap(); - src_location.push(filename); - src_location + storage_uri.join(filename).unwrap() } diff --git a/quickwit/quickwit-common/src/fs.rs b/quickwit/quickwit-common/src/fs.rs index adcb432e1b1..1aaa43d8286 100644 --- a/quickwit/quickwit-common/src/fs.rs +++ b/quickwit/quickwit-common/src/fs.rs @@ -34,7 +34,7 @@ pub async fn empty_dir>(path: P) -> anyhow::Result<()> { Ok(()) } -/// Helper function to get the cache path. +/// Helper function to get the indexer split cache path. pub fn get_cache_directory_path(data_dir_path: &Path) -> PathBuf { data_dir_path.join("indexer-split-cache").join("splits") } diff --git a/quickwit/quickwit-config/src/lib.rs b/quickwit/quickwit-config/src/lib.rs index e41c46dce5b..5e256793fcd 100644 --- a/quickwit/quickwit-config/src/lib.rs +++ b/quickwit/quickwit-config/src/lib.rs @@ -55,11 +55,13 @@ pub use quickwit_doc_mapper::DocMapping; use serde::de::DeserializeOwned; use serde::Serialize; use serde_json::Value as JsonValue; +use source_config::FileSourceParamsForSerde; pub use source_config::{ - load_source_config_from_user_config, FileSourceParams, KafkaSourceParams, KinesisSourceParams, - PubSubSourceParams, PulsarSourceAuth, PulsarSourceParams, RegionOrEndpoint, SourceConfig, - SourceInputFormat, SourceParams, TransformConfig, VecSourceParams, VoidSourceParams, - CLI_SOURCE_ID, INGEST_API_SOURCE_ID, INGEST_V2_SOURCE_ID, + load_source_config_from_user_config, FileSourceMessageType, FileSourceNotification, + FileSourceParams, FileSourceSqs, KafkaSourceParams, KinesisSourceParams, PubSubSourceParams, + PulsarSourceAuth, PulsarSourceParams, RegionOrEndpoint, SourceConfig, SourceInputFormat, + SourceParams, TransformConfig, VecSourceParams, VoidSourceParams, CLI_SOURCE_ID, + INGEST_API_SOURCE_ID, INGEST_V2_SOURCE_ID, }; use tracing::warn; @@ -112,7 +114,10 @@ pub fn disable_ingest_v1() -> bool { IndexTemplateV0_8, SourceInputFormat, SourceParams, - FileSourceParams, + FileSourceMessageType, + FileSourceNotification, + FileSourceParamsForSerde, + FileSourceSqs, PubSubSourceParams, KafkaSourceParams, KinesisSourceParams, diff --git a/quickwit/quickwit-config/src/source_config/mod.rs b/quickwit/quickwit-config/src/source_config/mod.rs index bc1c0cf3168..b9fcaa15018 100644 --- a/quickwit/quickwit-config/src/source_config/mod.rs +++ b/quickwit/quickwit-config/src/source_config/mod.rs @@ -19,8 +19,8 @@ pub(crate) mod serialize; +use std::borrow::Cow; use std::num::NonZeroUsize; -use std::path::{Path, PathBuf}; use std::str::FromStr; use bytes::Bytes; @@ -82,6 +82,7 @@ impl SourceConfig { SourceParams::Kinesis(_) => SourceType::Kinesis, SourceParams::PubSub(_) => SourceType::PubSub, SourceParams::Pulsar(_) => SourceType::Pulsar, + SourceParams::Stdin => SourceType::Stdin, SourceParams::Vec(_) => SourceType::Vec, SourceParams::Void(_) => SourceType::Void, } @@ -98,6 +99,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::Stdin => serde_json::to_value(()), SourceParams::Vec(params) => serde_json::to_value(params), SourceParams::Void(params) => serde_json::to_value(params), } @@ -214,6 +216,7 @@ impl FromStr for SourceInputFormat { #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, utoipa::ToSchema)] #[serde(tag = "source_type", content = "params", rename_all = "snake_case")] pub enum SourceParams { + #[schema(value_type = FileSourceParamsForSerde)] File(FileSourceParams), Ingest, #[serde(rename = "ingest-api")] @@ -225,17 +228,22 @@ pub enum SourceParams { #[serde(rename = "pubsub")] PubSub(PubSubSourceParams), Pulsar(PulsarSourceParams), + Stdin, Vec(VecSourceParams), Void(VoidSourceParams), } impl SourceParams { - pub fn file>(filepath: P) -> Self { - Self::File(FileSourceParams::file(filepath)) + pub fn file_from_uri(uri: Uri) -> Self { + Self::File(FileSourceParams::Filepath(uri)) + } + + pub fn file_from_str>(filepath: P) -> anyhow::Result { + Uri::from_str(filepath.as_ref()).map(Self::file_from_uri) } pub fn stdin() -> Self { - Self::File(FileSourceParams::stdin()) + Self::Stdin } pub fn void() -> Self { @@ -243,41 +251,92 @@ impl SourceParams { } } +#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize, utoipa::ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum FileSourceMessageType { + /// See + S3Notification, + /// A string with the URI of the file (e.g `s3://bucket/key`) + RawUri, +} + #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, utoipa::ToSchema)] +pub struct FileSourceSqs { + pub queue_url: String, + pub message_type: FileSourceMessageType, +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, utoipa::ToSchema)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum FileSourceNotification { + Sqs(FileSourceSqs), +} + +#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize, utoipa::ToSchema)] #[serde(deny_unknown_fields)] -pub struct FileSourceParams { - /// Path of the file to read. Assume stdin if None. - #[schema(value_type = String)] - #[serde(skip_serializing_if = "Option::is_none")] - #[serde(default)] - #[serde(deserialize_with = "absolute_filepath_from_str")] - pub filepath: Option, //< If None read from stdin. +pub(super) struct FileSourceParamsForSerde { + #[serde(default, skip_serializing_if = "Vec::is_empty")] + notifications: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + filepath: Option, } -/// Deserializing as an URI first to validate the input. -/// -/// TODO: we might want to replace `PathBuf` with `Uri` directly in -/// `FileSourceParams` -fn absolute_filepath_from_str<'de, D>(deserializer: D) -> Result, D::Error> -where D: Deserializer<'de> { - let filepath_opt: Option = Deserialize::deserialize(deserializer)?; - if let Some(filepath) = filepath_opt { - let uri = Uri::from_str(&filepath).map_err(D::Error::custom)?; - Ok(Some(PathBuf::from(uri.as_str()))) - } else { - Ok(None) +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +#[serde( + try_from = "FileSourceParamsForSerde", + into = "FileSourceParamsForSerde" +)] +pub enum FileSourceParams { + Notifications(FileSourceNotification), + Filepath(Uri), +} + +impl TryFrom for FileSourceParams { + type Error = Cow<'static, str>; + + fn try_from(mut value: FileSourceParamsForSerde) -> Result { + if value.filepath.is_some() && !value.notifications.is_empty() { + return Err( + "File source parameters `notifications` and `filepath` are mutually exclusive" + .into(), + ); + } + if let Some(filepath) = value.filepath { + let uri = Uri::from_str(&filepath).map_err(|err| err.to_string())?; + Ok(FileSourceParams::Filepath(uri)) + } else if value.notifications.len() == 1 { + Ok(FileSourceParams::Notifications( + value.notifications.remove(0), + )) + } else if value.notifications.len() > 1 { + return Err("Only one notification can be specified for now".into()); + } else { + return Err( + "Either `notifications` or `filepath` must be specified as file source parameters" + .into(), + ); + } } } -impl FileSourceParams { - pub fn file>(filepath: P) -> Self { - FileSourceParams { - filepath: Some(filepath.as_ref().to_path_buf()), +impl From for FileSourceParamsForSerde { + fn from(value: FileSourceParams) -> Self { + match value { + FileSourceParams::Filepath(uri) => Self { + filepath: Some(uri.to_string()), + notifications: vec![], + }, + FileSourceParams::Notifications(notification) => Self { + filepath: None, + notifications: vec![notification], + }, } } +} - pub fn stdin() -> Self { - FileSourceParams { filepath: None } +impl FileSourceParams { + pub fn from_filepath>(filepath: P) -> anyhow::Result { + Uri::from_str(filepath.as_ref()).map(Self::Filepath) } } @@ -802,18 +861,88 @@ mod tests { } #[test] - fn test_file_source_params_serialization() { + fn test_file_source_params_serde() { { let yaml = r#" filepath: source-path.json "#; - let file_params = serde_yaml::from_str::(yaml).unwrap(); + let file_params_deserialized = serde_yaml::from_str::(yaml).unwrap(); let uri = Uri::from_str("source-path.json").unwrap(); + assert_eq!(file_params_deserialized, FileSourceParams::Filepath(uri)); + let file_params_reserialized = serde_json::to_value(file_params_deserialized).unwrap(); + file_params_reserialized + .get("filepath") + .unwrap() + .as_str() + .unwrap() + .contains("source-path.json"); + } + { + let yaml = r#" + notifications: + - type: sqs + queue_url: https://sqs.us-east-1.amazonaws.com/123456789012/queue-name + message_type: s3_notification + "#; + let file_params_deserialized = serde_yaml::from_str::(yaml).unwrap(); + assert_eq!( + file_params_deserialized, + FileSourceParams::Notifications(FileSourceNotification::Sqs(FileSourceSqs { + queue_url: "https://sqs.us-east-1.amazonaws.com/123456789012/queue-name" + .to_string(), + message_type: FileSourceMessageType::S3Notification, + })), + ); + let file_params_reserialized = serde_json::to_value(&file_params_deserialized).unwrap(); assert_eq!( - file_params.filepath.unwrap().as_path(), - Path::new(uri.as_str()) + file_params_reserialized, + json!({"notifications": [{"type": "sqs", "queue_url": "https://sqs.us-east-1.amazonaws.com/123456789012/queue-name", "message_type": "s3_notification"}]}) ); } + { + let yaml = r#" + filepath: source-path.json + notifications: + - type: sqs + queue_url: https://sqs.us-east-1.amazonaws.com/123456789012/queue-name + message_type: s3_notification + "#; + let error = serde_yaml::from_str::(yaml).unwrap_err(); + assert_eq!( + error.to_string(), + "File source parameters `notifications` and `filepath` are mutually exclusive" + ); + } + { + let yaml = r#" + notifications: + - type: sqs + queue_url: https://sqs.us-east-1.amazonaws.com/123456789012/queue1 + message_type: s3_notification + - type: sqs + queue_url: https://sqs.us-east-1.amazonaws.com/123456789012/queue2 + message_type: s3_notification + "#; + let error = serde_yaml::from_str::(yaml).unwrap_err(); + assert_eq!( + error.to_string(), + "Only one notification can be specified for now" + ); + } + { + let json = r#" + { + "notifications": [ + { + "queue_url": "https://sqs.us-east-1.amazonaws.com/123456789012/queue", + "message_type": "s3_notification" + } + ] + } + "#; + let error = serde_json::from_str::(json).unwrap_err(); + assert!(error.to_string().contains("missing field `type`")); + } } #[test] @@ -1199,7 +1328,9 @@ mod tests { "desired_num_pipelines": 1, "max_num_pipelines_per_indexer": 1, "source_type": "file", - "params": {"filepath": "/test_non_json_corpus.txt"}, + "params": { + "filepath": "s3://mybucket/test_non_json_corpus.txt" + }, "input_format": "plain_text" }"#; let source_config = diff --git a/quickwit/quickwit-config/src/source_config/serialize.rs b/quickwit/quickwit-config/src/source_config/serialize.rs index 00b76923775..a99f9371eed 100644 --- a/quickwit/quickwit-config/src/source_config/serialize.rs +++ b/quickwit/quickwit-config/src/source_config/serialize.rs @@ -24,7 +24,10 @@ use quickwit_proto::types::SourceId; use serde::{Deserialize, Serialize}; use super::{TransformConfig, RESERVED_SOURCE_IDS}; -use crate::{validate_identifier, ConfigFormat, SourceConfig, SourceInputFormat, SourceParams}; +use crate::{ + validate_identifier, ConfigFormat, FileSourceParams, SourceConfig, SourceInputFormat, + SourceParams, +}; type SourceConfigForSerialization = SourceConfigV0_8; @@ -65,12 +68,11 @@ impl SourceConfigForSerialization { /// Checks the validity of the `SourceConfig` as a "deserializable source". /// /// Two remarks: - /// - This does not check connectivity. (See `check_connectivity(..)`) - /// This just validate configuration, without performing any IO. - /// - This is only here to validate user input. - /// When ingesting from stdin, we programmatically create an invalid `SourceConfig`. - /// - /// TODO refactor #1065 + /// - This does not check connectivity, it just validate configuration, + /// without performing any IO. See `check_connectivity(..)`. + /// - This is used each time the `SourceConfig` is deserialized (at creation but also during + /// communications with the metastore). When ingesting from stdin, we programmatically create + /// an invalid `SourceConfig` and only use it locally. fn validate_and_build(self) -> anyhow::Result { if !RESERVED_SOURCE_IDS.contains(&self.source_id.as_str()) { validate_identifier("source", &self.source_id)?; @@ -78,16 +80,16 @@ impl SourceConfigForSerialization { let num_pipelines = NonZeroUsize::new(self.num_pipelines) .ok_or_else(|| anyhow::anyhow!("`desired_num_pipelines` must be strictly positive"))?; match &self.source_params { - // We want to forbid source_config with no filepath - SourceParams::File(file_params) => { - if file_params.filepath.is_none() { - bail!( - "source `{}` of type `file` must contain a filepath", - self.source_id - ) - } + SourceParams::Stdin => { + bail!( + "stdin can only be used as source through the CLI command `quickwit tool \ + local-ingest`" + ); } - SourceParams::Kafka(_) | SourceParams::Kinesis(_) | SourceParams::Pulsar(_) => { + SourceParams::File(_) + | SourceParams::Kafka(_) + | SourceParams::Kinesis(_) + | SourceParams::Pulsar(_) => { // TODO consider any validation opportunity } SourceParams::PubSub(_) @@ -98,7 +100,9 @@ impl SourceConfigForSerialization { | SourceParams::Void(_) => {} } match &self.source_params { - SourceParams::PubSub(_) | SourceParams::Kafka(_) => {} + SourceParams::PubSub(_) + | SourceParams::Kafka(_) + | SourceParams::File(FileSourceParams::Notifications(_)) => {} _ => { 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"); diff --git a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs index ff1bd5658ad..c396e1ac23a 100644 --- a/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs +++ b/quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs @@ -30,11 +30,11 @@ use fnv::{FnvHashMap, FnvHashSet}; use itertools::Itertools; use once_cell::sync::OnceCell; use quickwit_common::pretty::PrettySample; +use quickwit_config::{FileSourceParams, SourceParams}; use quickwit_proto::indexing::{ ApplyIndexingPlanRequest, CpuCapacity, IndexingService, IndexingTask, PIPELINE_FULL_CAPACITY, PIPELINE_THROUGHPUT, }; -use quickwit_proto::metastore::SourceType; use quickwit_proto::types::NodeId; use scheduling::{SourceToSchedule, SourceToScheduleType}; use serde::Serialize; @@ -168,22 +168,22 @@ fn get_sources_to_schedule(model: &ControlPlaneModel) -> Vec { if !source_config.enabled { continue; } - match source_config.source_type() { - SourceType::Cli - | SourceType::File - | SourceType::Vec - | SourceType::Void - | SourceType::Unspecified => { - // We don't need to schedule those. + match source_config.source_params { + SourceParams::File(FileSourceParams::Filepath(_)) + | SourceParams::IngestCli + | SourceParams::Stdin + | SourceParams::Void(_) + | SourceParams::Vec(_) => { // We don't need to schedule those. } - SourceType::IngestV1 => { + + SourceParams::IngestApi => { // TODO ingest v1 is scheduled differently sources.push(SourceToSchedule { source_uid, source_type: SourceToScheduleType::IngestV1, }); } - SourceType::IngestV2 => { + SourceParams::Ingest => { // Expect: the source should exist since we just read it from `get_source_configs`. // Note that we keep all shards, including Closed shards: // A closed shards still needs to be indexed. @@ -208,11 +208,11 @@ fn get_sources_to_schedule(model: &ControlPlaneModel) -> Vec { }, }); } - SourceType::Kafka - | SourceType::Kinesis - | SourceType::PubSub - | SourceType::Nats - | SourceType::Pulsar => { + SourceParams::Kafka(_) + | SourceParams::Kinesis(_) + | SourceParams::PubSub(_) + | SourceParams::Pulsar(_) + | SourceParams::File(FileSourceParams::Notifications(_)) => { sources.push(SourceToSchedule { source_uid, source_type: SourceToScheduleType::NonSharded { diff --git a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs index 29111f28b23..8636cca7379 100644 --- a/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs +++ b/quickwit/quickwit-control-plane/src/ingest/ingest_controller.rs @@ -820,6 +820,8 @@ impl IngestController { leader_id: shard.leader_id.clone(), follower_id: shard.follower_id.clone(), doc_mapping_uid: shard.doc_mapping_uid, + // Shards are acquired by the ingest sources + publish_token: None, } }) .collect(); diff --git a/quickwit/quickwit-indexing/Cargo.toml b/quickwit/quickwit-indexing/Cargo.toml index d7008a4579c..7b005301ea9 100644 --- a/quickwit/quickwit-indexing/Cargo.toml +++ b/quickwit/quickwit-indexing/Cargo.toml @@ -11,12 +11,12 @@ authors.workspace = true license.workspace = true [dependencies] -aws-sdk-kinesis = { workspace = true, optional = true } - anyhow = { workspace = true } arc-swap = { workspace = true } async-compression = { workspace = true } async-trait = { workspace = true } +aws-sdk-kinesis = { workspace = true, optional = true } +aws-sdk-sqs = { workspace = true, optional = true } bytes = { workspace = true } bytesize = { workspace = true } fail = { workspace = true } @@ -34,6 +34,7 @@ oneshot = { workspace = true } openssl = { workspace = true, optional = true } pulsar = { workspace = true, optional = true } quickwit-query = { workspace = true } +regex = { workspace = true } rdkafka = { workspace = true, optional = true } serde = { workspace = true } serde_json = { workspace = true } @@ -77,6 +78,13 @@ kinesis = [ kinesis-localstack-tests = [] pulsar = ["dep:pulsar"] pulsar-broker-tests = [] +queue-sources = [] +sqs = [ + "aws-sdk-sqs", + "queue-sources", + "quickwit-aws/sqs", +] +sqs-localstack-tests = [] vendored-kafka = [ "kafka", "libz-sys/static", diff --git a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs index d0a34300473..4087f2ed230 100644 --- a/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs +++ b/quickwit/quickwit-indexing/src/actors/indexing_pipeline.rs @@ -436,6 +436,7 @@ impl IndexingPipeline { queues_dir_path: self.params.queues_dir_path.clone(), storage_resolver: self.params.source_storage_resolver.clone(), event_broker: self.params.event_broker.clone(), + indexing_setting: self.params.indexing_settings.clone(), }; let source = ctx .protect_future(quickwit_supported_sources().load_source(source_runtime)) @@ -598,7 +599,7 @@ mod tests { use quickwit_actors::{Command, Universe}; use quickwit_common::ServiceStream; - use quickwit_config::{IndexingSettings, SourceInputFormat, SourceParams, VoidSourceParams}; + use quickwit_config::{IndexingSettings, SourceInputFormat, SourceParams}; use quickwit_doc_mapper::{default_doc_mapper_for_test, DefaultDocMapper}; use quickwit_metastore::checkpoint::IndexCheckpointDelta; use quickwit_metastore::{IndexMetadata, IndexMetadataResponseExt, PublishSplitsRequestExt}; @@ -639,7 +640,7 @@ mod tests { source_id: "test-source".to_string(), num_pipelines: NonZeroUsize::MIN, enabled: true, - source_params: SourceParams::file(PathBuf::from(test_file)), + source_params: SourceParams::file_from_str(test_file).unwrap(), transform_config: None, input_format: SourceInputFormat::Json, }; @@ -689,7 +690,7 @@ mod tests { && publish_splits_request.staged_split_ids.len() == 1 && publish_splits_request.replaced_split_ids.is_empty() && format!("{:?}", checkpoint_delta.source_delta) - .ends_with(":(00000000000000000000..00000000000000001030])") + .ends_with(":(00000000000000000000..~00000000000000001030])") }) .returning(|_| Ok(EmptyResponse {})); @@ -758,7 +759,7 @@ mod tests { source_id: "test-source".to_string(), num_pipelines: NonZeroUsize::MIN, enabled: true, - source_params: SourceParams::file(PathBuf::from(test_file)), + source_params: SourceParams::file_from_str(test_file).unwrap(), transform_config: None, input_format: SourceInputFormat::Json, }; @@ -801,7 +802,7 @@ mod tests { && publish_splits_request.replaced_split_ids.is_empty() && checkpoint_delta.source_id == "test-source" && format!("{:?}", checkpoint_delta.source_delta) - .ends_with(":(00000000000000000000..00000000000000001030])") + .ends_with(":(00000000000000000000..~00000000000000001030])") }) .returning(|_| Ok(EmptyResponse {})); @@ -862,7 +863,7 @@ mod tests { source_id: "test-source".to_string(), num_pipelines: NonZeroUsize::MIN, enabled: true, - source_params: SourceParams::Void(VoidSourceParams), + source_params: SourceParams::void(), transform_config: None, input_format: SourceInputFormat::Json, }; @@ -965,7 +966,7 @@ mod tests { source_id: "test-source".to_string(), num_pipelines: NonZeroUsize::MIN, enabled: true, - source_params: SourceParams::file(PathBuf::from(test_file)), + source_params: SourceParams::file_from_str(test_file).unwrap(), transform_config: None, input_format: SourceInputFormat::Json, }; @@ -1008,7 +1009,7 @@ mod tests { && publish_splits_request.replaced_split_ids.is_empty() && checkpoint_delta.source_id == "test-source" && format!("{:?}", checkpoint_delta.source_delta) - .ends_with(":(00000000000000000000..00000000000000001030])") + .ends_with(":(00000000000000000000..~00000000000000001030])") }) .returning(|_| Ok(EmptyResponse {})); let universe = Universe::new(); diff --git a/quickwit/quickwit-indexing/src/actors/publisher.rs b/quickwit/quickwit-indexing/src/actors/publisher.rs index 1f24ef51477..4e999e7f7a2 100644 --- a/quickwit/quickwit-indexing/src/actors/publisher.rs +++ b/quickwit/quickwit-indexing/src/actors/publisher.rs @@ -147,7 +147,7 @@ impl Handler for Publisher { ); return Ok(()); } - info!(new_splits=?split_ids, checkpoint_delta=?checkpoint_delta_opt, "publish-new-splits"); + info!("publish-new-splits"); if let Some(source_mailbox) = self.source_mailbox_opt.as_ref() { if let Some(checkpoint) = checkpoint_delta_opt { // We voluntarily do not log anything here. diff --git a/quickwit/quickwit-indexing/src/source/doc_file_reader.rs b/quickwit/quickwit-indexing/src/source/doc_file_reader.rs new file mode 100644 index 00000000000..f8fc3c0b8d3 --- /dev/null +++ b/quickwit/quickwit-indexing/src/source/doc_file_reader.rs @@ -0,0 +1,561 @@ +// Copyright (C) 2024 Quickwit, Inc. +// +// Quickwit is offered under the AGPL v3.0 and as commercial software. +// For commercial licensing, contact us at hello@quickwit.io. +// +// AGPL: +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as +// published by the Free Software Foundation, either version 3 of the +// License, or (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::io; +use std::path::Path; + +use anyhow::Context; +use async_compression::tokio::bufread::GzipDecoder; +use bytes::Bytes; +use quickwit_common::uri::Uri; +use quickwit_common::Progress; +use quickwit_metastore::checkpoint::PartitionId; +use quickwit_proto::metastore::SourceType; +use quickwit_proto::types::Position; +use quickwit_storage::StorageResolver; +use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, BufReader}; + +use super::{BatchBuilder, BATCH_NUM_BYTES_LIMIT}; + +pub struct FileRecord { + pub next_offset: u64, + pub doc: Bytes, + pub is_last: bool, +} + +/// A helper wrapper that lets you skip bytes in compressed files where you +/// cannot seek (e.g. gzip files). +struct SkipReader { + reader: BufReader>, + num_bytes_to_skip: usize, +} + +impl SkipReader { + fn new(reader: Box, num_bytes_to_skip: usize) -> Self { + Self { + reader: BufReader::new(reader), + num_bytes_to_skip, + } + } + + async fn skip(&mut self) -> io::Result<()> { + // allocate on the heap to avoid stack overflows + let mut buf = vec![0u8; 64_000]; + while self.num_bytes_to_skip > 0 { + let num_bytes_to_read = self.num_bytes_to_skip.min(buf.len()); + let num_bytes_read = self + .reader + .read_exact(&mut buf[..num_bytes_to_read]) + .await?; + self.num_bytes_to_skip -= num_bytes_read; + } + Ok(()) + } + + /// Reads a line and peeks into the readers buffer. Returns the number of + /// bytes read and true the end of the file is reached. + async fn read_line_and_peek<'a>(&mut self, buf: &'a mut String) -> io::Result<(usize, bool)> { + if self.num_bytes_to_skip > 0 { + self.skip().await?; + } + let line_size = self.reader.read_line(buf).await?; + if line_size == 0 { + return Ok((0, true)); + } + let next_bytes = self.reader.fill_buf().await?; + Ok((line_size, next_bytes.is_empty())) + } +} + +pub struct DocFileReader { + reader: SkipReader, + next_offset: u64, +} + +impl DocFileReader { + pub fn empty() -> Self { + DocFileReader { + reader: SkipReader::new(Box::new(tokio::io::empty()), 0), + next_offset: 0, + } + } + + pub async fn from_uri( + storage_resolver: &StorageResolver, + uri: &Uri, + offset: usize, + ) -> anyhow::Result { + let (dir_uri, file_name) = dir_and_filename(uri)?; + let storage = storage_resolver.resolve(&dir_uri).await?; + let file_size = storage.file_num_bytes(file_name).await?.try_into().unwrap(); + if file_size == 0 { + return Ok(DocFileReader::empty()); + } + // If it's a gzip file, we can't seek to a specific offset. `SkipReader` + // starts from the beginning of the file, decompresses and skips the + // first `offset` bytes. + let reader = if uri.extension() == Some("gz") { + let stream = storage.get_slice_stream(file_name, 0..file_size).await?; + let decompressed_stream = Box::new(GzipDecoder::new(BufReader::new(stream))); + DocFileReader { + reader: SkipReader::new(decompressed_stream, offset), + next_offset: offset as u64, + } + } else { + let stream = storage + .get_slice_stream(file_name, offset..file_size) + .await?; + DocFileReader { + reader: SkipReader::new(stream, 0), + next_offset: offset as u64, + } + }; + Ok(reader) + } + + /// Reads the next record from the underlying file. Returns `None` when EOF + /// is reached. + pub async fn next_record(&mut self) -> anyhow::Result> { + let mut buf = String::new(); + // TODO retry if stream is broken (#5243) + let (bytes_read, is_last) = self.reader.read_line_and_peek(&mut buf).await?; + if bytes_read == 0 { + Ok(None) + } else { + self.next_offset += bytes_read as u64; + Ok(Some(FileRecord { + next_offset: self.next_offset, + doc: Bytes::from(buf), + is_last, + })) + } + } +} + +pub struct ObjectUriBatchReader { + partition_id: PartitionId, + reader: DocFileReader, + current_offset: usize, + is_eof: bool, +} + +impl ObjectUriBatchReader { + pub async fn try_new( + storage_resolver: &StorageResolver, + partition_id: PartitionId, + uri: &Uri, + position: Position, + ) -> anyhow::Result { + let current_offset = match position { + Position::Beginning => 0, + Position::Offset(offset) => offset + .as_usize() + .context("file offset should be stored as usize")?, + Position::Eof(_) => { + return Ok(ObjectUriBatchReader { + partition_id, + reader: DocFileReader::empty(), + current_offset: 0, + is_eof: true, + }) + } + }; + let reader = DocFileReader::from_uri(storage_resolver, uri, current_offset).await?; + Ok(ObjectUriBatchReader { + partition_id, + reader, + current_offset, + is_eof: false, + }) + } + + pub async fn read_batch( + &mut self, + source_progress: &Progress, + source_type: SourceType, + ) -> anyhow::Result { + let limit_num_bytes = self.current_offset + BATCH_NUM_BYTES_LIMIT as usize; + let mut new_offset = self.current_offset; + let mut batch_builder = BatchBuilder::new(source_type); + while new_offset < limit_num_bytes { + if let Some(record) = source_progress + .protect_future(self.reader.next_record()) + .await? + { + new_offset = record.next_offset as usize; + batch_builder.add_doc(record.doc); + if record.is_last { + self.is_eof = true; + break; + } + } else { + self.is_eof = true; + break; + } + } + let to_position = if self.is_eof { + Position::eof(new_offset) + } else { + Position::offset(new_offset) + }; + batch_builder.checkpoint_delta.record_partition_delta( + self.partition_id.clone(), + Position::offset(self.current_offset), + to_position, + )?; + self.current_offset = new_offset; + Ok(batch_builder) + } + + pub fn is_eof(&self) -> bool { + self.is_eof + } +} + +pub(crate) fn dir_and_filename(filepath: &Uri) -> anyhow::Result<(Uri, &Path)> { + let dir_uri: Uri = filepath + .parent() + .context("Parent directory could not be resolved")?; + let file_name = filepath + .file_name() + .context("Path does not appear to be a file")?; + Ok((dir_uri, file_name)) +} + +#[cfg(test)] +pub mod file_test_helpers { + use std::io::Write; + + use async_compression::tokio::write::GzipEncoder; + use tempfile::NamedTempFile; + + pub const DUMMY_DOC: &[u8] = r#"{"body": "hello happy tax payer!"}"#.as_bytes(); + + async fn gzip_bytes(bytes: &[u8]) -> Vec { + let mut gzip_documents = Vec::new(); + let mut encoder = GzipEncoder::new(&mut gzip_documents); + tokio::io::AsyncWriteExt::write_all(&mut encoder, bytes) + .await + .unwrap(); + // flush is not sufficient here and reading the file will raise a unexpected end of file + // error. + tokio::io::AsyncWriteExt::shutdown(&mut encoder) + .await + .unwrap(); + gzip_documents + } + + async fn write_to_tmp(data: Vec, gzip: bool) -> NamedTempFile { + let mut temp_file: tempfile::NamedTempFile = if gzip { + tempfile::Builder::new().suffix(".gz").tempfile().unwrap() + } else { + tempfile::NamedTempFile::new().unwrap() + }; + if gzip { + let gzip_documents = gzip_bytes(&data).await; + temp_file.write_all(&gzip_documents).unwrap(); + } else { + temp_file.write_all(&data).unwrap(); + } + temp_file.flush().unwrap(); + temp_file + } + + pub async fn generate_dummy_doc_file(gzip: bool, lines: usize) -> (NamedTempFile, usize) { + let mut documents_bytes = Vec::with_capacity(DUMMY_DOC.len() * lines); + for _ in 0..lines { + documents_bytes.write_all(DUMMY_DOC).unwrap(); + documents_bytes.write_all("\n".as_bytes()).unwrap(); + } + let size = documents_bytes.len(); + let file = write_to_tmp(documents_bytes, gzip).await; + (file, size) + } + + /// Generates a file with increasing padded numbers. Each line is 8 bytes + /// including the newline char. + /// + /// 0000000\n0000001\n0000002\n... + pub async fn generate_index_doc_file(gzip: bool, lines: usize) -> NamedTempFile { + assert!(lines < 9999999, "each line is 7 digits + newline"); + let mut documents_bytes = Vec::new(); + for i in 0..lines { + documents_bytes + .write_all(format!("{:0>7}\n", i).as_bytes()) + .unwrap(); + } + write_to_tmp(documents_bytes, gzip).await + } +} + +#[cfg(test)] +mod tests { + use std::io::Cursor; + use std::str::FromStr; + + use file_test_helpers::generate_index_doc_file; + use quickwit_metastore::checkpoint::SourceCheckpointDelta; + + use super::*; + + #[tokio::test] + async fn test_skip_reader() { + { + // Skip 0 bytes. + let mut reader = SkipReader::new(Box::new("hello".as_bytes()), 0); + let mut buf = String::new(); + let (bytes_read, eof) = reader.read_line_and_peek(&mut buf).await.unwrap(); + assert_eq!(buf, "hello"); + assert!(eof); + assert_eq!(bytes_read, 5) + } + { + // Skip 2 bytes. + let mut reader = SkipReader::new(Box::new("hello".as_bytes()), 2); + let mut buf = String::new(); + let (bytes_read, eof) = reader.read_line_and_peek(&mut buf).await.unwrap(); + assert_eq!(buf, "llo"); + assert!(eof); + assert_eq!(bytes_read, 3) + } + { + let input = "hello"; + let cursor = Cursor::new(input); + let mut reader = SkipReader::new(Box::new(cursor), 5); + let mut buf = String::new(); + let (bytes_read, eof) = reader.read_line_and_peek(&mut buf).await.unwrap(); + assert!(eof); + assert_eq!(bytes_read, 0) + } + { + let input = "hello"; + let cursor = Cursor::new(input); + let mut reader = SkipReader::new(Box::new(cursor), 10); + let mut buf = String::new(); + assert!(reader.read_line_and_peek(&mut buf).await.is_err()); + } + { + let input = "hello world".repeat(10000); + let cursor = Cursor::new(input.clone()); + let mut reader = SkipReader::new(Box::new(cursor), 64000); + let mut buf = String::new(); + reader.read_line_and_peek(&mut buf).await.unwrap(); + assert_eq!(buf, input[64000..]); + } + { + let input = "hello world".repeat(10000); + let cursor = Cursor::new(input.clone()); + let mut reader = SkipReader::new(Box::new(cursor), 64001); + let mut buf = String::new(); + reader.read_line_and_peek(&mut buf).await.unwrap(); + assert_eq!(buf, input[64001..]); + } + } + + async fn aux_test_full_read_record(file: impl AsRef, expected_lines: usize) { + let storage_resolver = StorageResolver::for_test(); + let uri = Uri::from_str(file.as_ref()).unwrap(); + let mut doc_reader = DocFileReader::from_uri(&storage_resolver, &uri, 0) + .await + .unwrap(); + let mut parsed_lines = 0; + while doc_reader.next_record().await.unwrap().is_some() { + parsed_lines += 1; + } + assert_eq!(parsed_lines, expected_lines); + } + + #[tokio::test] + async fn test_full_read_record() { + aux_test_full_read_record("data/test_corpus.json", 4).await; + } + + #[tokio::test] + async fn test_full_read_record_gz() { + aux_test_full_read_record("data/test_corpus.json.gz", 4).await; + } + + #[tokio::test] + async fn test_empty_file() { + let empty_file = tempfile::NamedTempFile::new().unwrap(); + let empty_file_uri = empty_file.path().to_str().unwrap(); + aux_test_full_read_record(empty_file_uri, 0).await; + } + + async fn aux_test_resumed_read_record( + file: impl AsRef, + expected_lines: usize, + stop_at_line: usize, + ) { + let storage_resolver = StorageResolver::for_test(); + let uri = Uri::from_str(file.as_ref()).unwrap(); + // read the first part of the file + let mut first_part_reader = DocFileReader::from_uri(&storage_resolver, &uri, 0) + .await + .unwrap(); + let mut resume_offset = 0; + let mut parsed_lines = 0; + for _ in 0..stop_at_line { + let rec = first_part_reader + .next_record() + .await + .unwrap() + .expect("EOF happened before stop_at_line"); + resume_offset = rec.next_offset as usize; + assert_eq!(Bytes::from(format!("{:0>7}\n", parsed_lines)), rec.doc); + parsed_lines += 1; + } + // read the second part of the file + let mut second_part_reader = + DocFileReader::from_uri(&storage_resolver, &uri, resume_offset) + .await + .unwrap(); + while let Some(rec) = second_part_reader.next_record().await.unwrap() { + assert_eq!(Bytes::from(format!("{:0>7}\n", parsed_lines)), rec.doc); + parsed_lines += 1; + } + assert_eq!(parsed_lines, expected_lines); + } + + #[tokio::test] + async fn test_resumed_read_record() { + let dummy_doc_file = generate_index_doc_file(false, 1000).await; + let dummy_doc_file_uri = dummy_doc_file.path().to_str().unwrap(); + aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 1).await; + aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 40).await; + aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 999).await; + aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 1000).await; + } + + #[tokio::test] + async fn test_resumed_read_record_gz() { + let dummy_doc_file = generate_index_doc_file(true, 1000).await; + let dummy_doc_file_uri = dummy_doc_file.path().to_str().unwrap(); + aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 1).await; + aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 40).await; + aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 999).await; + aux_test_resumed_read_record(dummy_doc_file_uri, 1000, 1000).await; + } + + async fn aux_test_full_read_batch( + file: impl AsRef, + expected_lines: usize, + expected_batches: usize, + file_size: usize, + from: Position, + ) { + let progress = Progress::default(); + let storage_resolver = StorageResolver::for_test(); + let uri = Uri::from_str(file.as_ref()).unwrap(); + let partition = PartitionId::from("test"); + let mut batch_reader = + ObjectUriBatchReader::try_new(&storage_resolver, partition.clone(), &uri, from) + .await + .unwrap(); + + let mut parsed_lines = 0; + let mut parsed_batches = 0; + let mut checkpoint_delta = SourceCheckpointDelta::default(); + while !batch_reader.is_eof() { + let batch = batch_reader + .read_batch(&progress, SourceType::Unspecified) + .await + .unwrap(); + parsed_lines += batch.docs.len(); + parsed_batches += 1; + checkpoint_delta.extend(batch.checkpoint_delta).unwrap(); + } + assert_eq!(parsed_lines, expected_lines); + assert_eq!(parsed_batches, expected_batches); + let position = checkpoint_delta + .get_source_checkpoint() + .position_for_partition(&partition) + .unwrap() + .clone(); + assert_eq!(position, Position::eof(file_size)) + } + + #[tokio::test] + async fn test_read_batch_empty_file() { + let empty_file = tempfile::NamedTempFile::new().unwrap(); + let empty_file_uri = empty_file.path().to_str().unwrap(); + aux_test_full_read_batch(empty_file_uri, 0, 1, 0, Position::Beginning).await; + } + + #[tokio::test] + async fn test_full_read_single_batch() { + let num_lines = 10; + let dummy_doc_file = generate_index_doc_file(false, num_lines).await; + let dummy_doc_file_uri = dummy_doc_file.path().to_str().unwrap(); + aux_test_full_read_batch( + dummy_doc_file_uri, + num_lines, + 1, + num_lines * 8, + Position::Beginning, + ) + .await; + } + + #[tokio::test] + async fn test_full_read_single_batch_max_size() { + let num_lines = BATCH_NUM_BYTES_LIMIT as usize / 8; + let dummy_doc_file = generate_index_doc_file(false, num_lines).await; + let dummy_doc_file_uri = dummy_doc_file.path().to_str().unwrap(); + aux_test_full_read_batch( + dummy_doc_file_uri, + num_lines, + 1, + num_lines * 8, + Position::Beginning, + ) + .await; + } + + #[tokio::test] + async fn test_full_read_two_batches() { + let num_lines = BATCH_NUM_BYTES_LIMIT as usize / 8 + 10; + let dummy_doc_file = generate_index_doc_file(false, num_lines).await; + let dummy_doc_file_uri = dummy_doc_file.path().to_str().unwrap(); + aux_test_full_read_batch( + dummy_doc_file_uri, + num_lines, + 2, + num_lines * 8, + Position::Beginning, + ) + .await; + } + + #[tokio::test] + async fn test_resume_read_batches() { + let total_num_lines = BATCH_NUM_BYTES_LIMIT as usize / 8 * 3; + let resume_after_lines = total_num_lines / 2; + let dummy_doc_file = generate_index_doc_file(false, total_num_lines).await; + let dummy_doc_file_uri = dummy_doc_file.path().to_str().unwrap(); + aux_test_full_read_batch( + dummy_doc_file_uri, + total_num_lines - resume_after_lines, + 2, + total_num_lines * 8, + Position::offset(resume_after_lines * 8), + ) + .await; + } +} diff --git a/quickwit/quickwit-indexing/src/source/file_source.rs b/quickwit/quickwit-indexing/src/source/file_source.rs index e4674f7a5ce..e3be553ae1b 100644 --- a/quickwit/quickwit-indexing/src/source/file_source.rs +++ b/quickwit/quickwit-indexing/src/source/file_source.rs @@ -17,45 +17,36 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::borrow::Borrow; -use std::ffi::OsStr; -use std::path::Path; +use std::fmt; use std::time::Duration; -use std::{fmt, io}; -use anyhow::Context; -use async_compression::tokio::bufread::GzipDecoder; use async_trait::async_trait; -use bytes::Bytes; use quickwit_actors::{ActorExitStatus, Mailbox}; -use quickwit_common::uri::Uri; use quickwit_config::FileSourceParams; -use quickwit_metastore::checkpoint::PartitionId; +use quickwit_metastore::checkpoint::{PartitionId, SourceCheckpoint}; use quickwit_proto::metastore::SourceType; -use quickwit_proto::types::{Position, SourceId}; -use serde::Serialize; -use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, BufReader}; -use tracing::info; +use quickwit_proto::types::SourceId; -use super::BatchBuilder; +use super::doc_file_reader::ObjectUriBatchReader; +#[cfg(feature = "queue-sources")] +use super::queue_sources::coordinator::QueueCoordinator; use crate::actors::DocProcessor; use crate::source::{Source, SourceContext, SourceRuntime, TypedSourceFactory}; -/// Number of bytes after which a new batch is cut. -pub(crate) const BATCH_NUM_BYTES_LIMIT: u64 = 500_000u64; - -#[derive(Default, Clone, Debug, Eq, PartialEq, Serialize)] -pub struct FileSourceCounters { - pub previous_offset: u64, - pub current_offset: u64, - pub num_lines_processed: u64, +enum FileSourceState { + #[cfg(feature = "queue-sources")] + Notification(QueueCoordinator), + Filepath { + batch_reader: ObjectUriBatchReader, + num_bytes_processed: u64, + num_lines_processed: u64, + }, } pub struct FileSource { source_id: SourceId, - params: FileSourceParams, - counters: FileSourceCounters, - reader: FileSourceReader, + state: FileSourceState, + source_type: SourceType, } impl fmt::Debug for FileSource { @@ -66,66 +57,90 @@ impl fmt::Debug for FileSource { #[async_trait] impl Source for FileSource { + #[allow(unused_variables)] + async fn initialize( + &mut self, + doc_processor_mailbox: &Mailbox, + ctx: &SourceContext, + ) -> Result<(), ActorExitStatus> { + match &mut self.state { + #[cfg(feature = "queue-sources")] + FileSourceState::Notification(coordinator) => { + coordinator.initialize(doc_processor_mailbox, ctx).await + } + FileSourceState::Filepath { .. } => Ok(()), + } + } + + #[allow(unused_variables)] async fn emit_batches( &mut self, doc_processor_mailbox: &Mailbox, ctx: &SourceContext, ) -> Result { - // We collect batches of documents before sending them to the indexer. - let limit_num_bytes = self.counters.previous_offset + BATCH_NUM_BYTES_LIMIT; - let mut reached_eof = false; - let mut batch_builder = BatchBuilder::new(SourceType::File); - - while self.counters.current_offset < limit_num_bytes { - let mut doc_line = String::new(); - // guard the zone in case of slow read, such as reading from someone - // typing to stdin - let num_bytes = ctx - .protect_future(self.reader.read_line(&mut doc_line)) - .await - .map_err(anyhow::Error::from)?; - if num_bytes == 0 { - reached_eof = true; - break; + match &mut self.state { + #[cfg(feature = "queue-sources")] + FileSourceState::Notification(coordinator) => { + coordinator.emit_batches(doc_processor_mailbox, ctx).await?; } - batch_builder.add_doc(Bytes::from(doc_line)); - self.counters.current_offset += num_bytes as u64; - self.counters.num_lines_processed += 1; - } - if !batch_builder.docs.is_empty() { - if let Some(filepath) = &self.params.filepath { - let filepath_str = filepath - .to_str() - .context("path is invalid utf-8")? - .to_string(); - let partition_id = PartitionId::from(filepath_str); - batch_builder - .checkpoint_delta - .record_partition_delta( - partition_id, - Position::offset(self.counters.previous_offset), - Position::offset(self.counters.current_offset), - ) - .unwrap(); + FileSourceState::Filepath { + batch_reader, + num_bytes_processed, + num_lines_processed, + } => { + let batch_builder = batch_reader + .read_batch(ctx.progress(), self.source_type) + .await?; + *num_bytes_processed += batch_builder.num_bytes; + *num_lines_processed += batch_builder.docs.len() as u64; + doc_processor_mailbox + .send_message(batch_builder.build()) + .await?; + if batch_reader.is_eof() { + ctx.send_exit_with_success(doc_processor_mailbox).await?; + return Err(ActorExitStatus::Success); + } } - self.counters.previous_offset = self.counters.current_offset; - ctx.send_message(doc_processor_mailbox, batch_builder.build()) - .await?; } - if reached_eof { - info!("reached end of file"); - ctx.send_exit_with_success(doc_processor_mailbox).await?; - return Err(ActorExitStatus::Success); - } - Ok(Duration::default()) + Ok(Duration::ZERO) } fn name(&self) -> String { format!("{:?}", self) } + #[allow(unused_variables)] + async fn suggest_truncate( + &mut self, + checkpoint: SourceCheckpoint, + ctx: &SourceContext, + ) -> anyhow::Result<()> { + match &mut self.state { + #[cfg(feature = "queue-sources")] + FileSourceState::Notification(coordinator) => { + coordinator.suggest_truncate(checkpoint, ctx).await + } + FileSourceState::Filepath { .. } => Ok(()), + } + } + fn observable_state(&self) -> serde_json::Value { - serde_json::to_value(&self.counters).unwrap() + match &self.state { + #[cfg(feature = "queue-sources")] + FileSourceState::Notification(coordinator) => { + serde_json::to_value(coordinator.observable_state()).unwrap() + } + FileSourceState::Filepath { + num_bytes_processed, + num_lines_processed, + .. + } => { + serde_json::json!({ + "num_bytes_processed": num_bytes_processed, + "num_lines_processed": num_lines_processed, + }) + } + } } } @@ -140,116 +155,71 @@ impl TypedSourceFactory for FileSourceFactory { source_runtime: SourceRuntime, params: FileSourceParams, ) -> anyhow::Result { - let checkpoint = source_runtime.fetch_checkpoint().await?; - let mut offset = 0; - - let reader: FileSourceReader = if let Some(filepath) = ¶ms.filepath { - let partition_id = PartitionId::from(filepath.to_string_lossy().borrow()); - offset = checkpoint - .position_for_partition(&partition_id) - .map(|position| { - position - .as_usize() - .expect("file offset should be stored as usize") - }) - .unwrap_or(0); - let (dir_uri, file_name) = dir_and_filename(filepath)?; - let storage = source_runtime.storage_resolver.resolve(&dir_uri).await?; - let file_size = storage.file_num_bytes(file_name).await?.try_into().unwrap(); - // If it's a gzip file, we can't seek to a specific offset, we need to start from the - // beginning of the file, decompress and skip the first `offset` bytes. - if filepath.extension() == Some(OsStr::new("gz")) { - let stream = storage.get_slice_stream(file_name, 0..file_size).await?; - FileSourceReader::new(Box::new(GzipDecoder::new(BufReader::new(stream))), offset) - } else { - let stream = storage - .get_slice_stream(file_name, offset..file_size) - .await?; - FileSourceReader::new(stream, 0) + let source_id = source_runtime.source_config.source_id.clone(); + let source_type = source_runtime.source_config.source_type(); + let state = match params { + FileSourceParams::Filepath(file_uri) => { + let partition_id = PartitionId::from(file_uri.as_str()); + let position = source_runtime + .fetch_checkpoint() + .await? + .position_for_partition(&partition_id) + .cloned() + .unwrap_or_default(); + let batch_reader = ObjectUriBatchReader::try_new( + &source_runtime.storage_resolver, + partition_id, + &file_uri, + position, + ) + .await?; + FileSourceState::Filepath { + batch_reader, + num_bytes_processed: 0, + num_lines_processed: 0, + } + } + #[cfg(feature = "sqs")] + FileSourceParams::Notifications(quickwit_config::FileSourceNotification::Sqs( + sqs_config, + )) => { + let coordinator = + QueueCoordinator::try_from_sqs_config(sqs_config, source_runtime).await?; + FileSourceState::Notification(coordinator) + } + #[cfg(not(feature = "sqs"))] + FileSourceParams::Notifications(quickwit_config::FileSourceNotification::Sqs(_)) => { + anyhow::bail!("Quickwit was compiled without the `sqs` feature") } - } else { - // We cannot use the checkpoint. - FileSourceReader::new(Box::new(tokio::io::stdin()), 0) - }; - let file_source = FileSource { - source_id: source_runtime.source_id().to_string(), - counters: FileSourceCounters { - previous_offset: offset as u64, - current_offset: offset as u64, - num_lines_processed: 0, - }, - reader, - params, }; - Ok(file_source) - } -} - -struct FileSourceReader { - reader: BufReader>, - num_bytes_to_skip: usize, -} -impl FileSourceReader { - fn new(reader: Box, num_bytes_to_skip: usize) -> Self { - Self { - reader: BufReader::new(reader), - num_bytes_to_skip, - } + Ok(FileSource { + state, + source_id, + source_type, + }) } - - // This function is only called for GZIP file. - // Because they cannot be sought into, we have to scan them to the right initial position. - async fn skip(&mut self) -> io::Result<()> { - // Allocate once a 64kb buffer. - let mut buf = [0u8; 64000]; - while self.num_bytes_to_skip > 0 { - let num_bytes_to_read = self.num_bytes_to_skip.min(buf.len()); - let num_bytes_read = self - .reader - .read_exact(&mut buf[..num_bytes_to_read]) - .await?; - self.num_bytes_to_skip -= num_bytes_read; - } - Ok(()) - } - - async fn read_line<'a>(&mut self, buf: &'a mut String) -> io::Result { - if self.num_bytes_to_skip > 0 { - self.skip().await?; - } - self.reader.read_line(buf).await - } -} - -pub(crate) fn dir_and_filename(filepath: &Path) -> anyhow::Result<(Uri, &Path)> { - let dir_uri: Uri = filepath - .parent() - .context("Parent directory could not be resolved")? - .to_str() - .context("Path cannot be turned to string")? - .parse()?; - let file_name = filepath - .file_name() - .context("Path does not appear to be a file")?; - Ok((dir_uri, file_name.as_ref())) } #[cfg(test)] mod tests { - use std::io::{Cursor, Write}; use std::num::NonZeroUsize; + use std::str::FromStr; - use async_compression::tokio::write::GzipEncoder; + use bytes::Bytes; use quickwit_actors::{Command, Universe}; + use quickwit_common::uri::Uri; use quickwit_config::{SourceConfig, SourceInputFormat, SourceParams}; - use quickwit_metastore::checkpoint::SourceCheckpointDelta; - use quickwit_proto::types::IndexUid; + use quickwit_metastore::checkpoint::{PartitionId, SourceCheckpointDelta}; + use quickwit_proto::types::{IndexUid, Position}; use super::*; use crate::models::RawDocBatch; + use crate::source::doc_file_reader::file_test_helpers::{ + generate_dummy_doc_file, generate_index_doc_file, DUMMY_DOC, + }; use crate::source::tests::SourceRuntimeBuilder; - use crate::source::SourceActor; + use crate::source::{SourceActor, BATCH_NUM_BYTES_LIMIT}; #[tokio::test] async fn test_file_source() { @@ -261,9 +231,9 @@ mod tests { let universe = Universe::with_accelerated_time(); let (doc_processor_mailbox, indexer_inbox) = universe.create_test_mailbox(); let params = if gzip { - FileSourceParams::file("data/test_corpus.json.gz") + FileSourceParams::from_filepath("data/test_corpus.json.gz").unwrap() } else { - FileSourceParams::file("data/test_corpus.json") + FileSourceParams::from_filepath("data/test_corpus.json").unwrap() }; let source_config = SourceConfig { source_id: "test-file-source".to_string(), @@ -289,13 +259,13 @@ mod tests { assert_eq!( counters, serde_json::json!({ - "previous_offset": 1030u64, - "current_offset": 1030u64, + "num_bytes_processed": 1030u64, "num_lines_processed": 4u32 }) ); let batch = indexer_inbox.drain_for_test(); assert_eq!(batch.len(), 2); + batch[0].downcast_ref::().unwrap(); assert!(matches!( batch[1].downcast_ref::().unwrap(), Command::ExitWithSuccess @@ -312,33 +282,11 @@ mod tests { quickwit_common::setup_logging_for_tests(); let universe = Universe::with_accelerated_time(); let (doc_processor_mailbox, doc_processor_inbox) = universe.create_test_mailbox(); - let mut documents_bytes = Vec::new(); - for _ in 0..20_000 { - documents_bytes - .write_all(r#"{"body": "hello happy tax payer!"}"#.as_bytes()) - .unwrap(); - documents_bytes.write_all("\n".as_bytes()).unwrap(); - } - let mut temp_file: tempfile::NamedTempFile = if gzip { - tempfile::Builder::new().suffix(".gz").tempfile().unwrap() - } else { - tempfile::NamedTempFile::new().unwrap() - }; - if gzip { - let gzip_documents = gzip_bytes(&documents_bytes).await; - temp_file.write_all(&gzip_documents).unwrap(); - } else { - temp_file.write_all(&documents_bytes).unwrap(); - } - temp_file.flush().unwrap(); - let params = FileSourceParams::file(temp_file.path()); - let filepath = params - .filepath - .as_ref() - .unwrap() - .to_string_lossy() - .to_string(); - + let lines = BATCH_NUM_BYTES_LIMIT as usize / DUMMY_DOC.len() + 1; + let (temp_file, temp_file_size) = generate_dummy_doc_file(gzip, lines).await; + let filepath = temp_file.path().to_str().unwrap(); + let uri = Uri::from_str(filepath).unwrap(); + let params = FileSourceParams::Filepath(uri.clone()); let source_config = SourceConfig { source_id: "test-file-source".to_string(), num_pipelines: NonZeroUsize::new(1).unwrap(), @@ -363,9 +311,8 @@ mod tests { assert_eq!( counters, serde_json::json!({ - "previous_offset": 700_000u64, - "current_offset": 700_000u64, - "num_lines_processed": 20_000u64 + "num_lines_processed": lines, + "num_bytes_processed": temp_file_size, }) ); let indexer_msgs = doc_processor_inbox.drain_for_test(); @@ -377,27 +324,19 @@ mod tests { format!("{:?}", &batch1.checkpoint_delta), format!( "∆({}:{})", - filepath, "(00000000000000000000..00000000000000500010]" + uri, "(00000000000000000000..00000000000005242895]" ) ); assert_eq!( - &extract_position_delta(&batch1.checkpoint_delta).unwrap(), - "00000000000000000000..00000000000000500010" - ); - assert_eq!( - &extract_position_delta(&batch2.checkpoint_delta).unwrap(), - "00000000000000500010..00000000000000700000" + format!("{:?}", &batch2.checkpoint_delta), + format!( + "∆({}:{})", + uri, "(00000000000005242895..~00000000000005397105]" + ) ); assert!(matches!(command, &Command::ExitWithSuccess)); } - fn extract_position_delta(checkpoint_delta: &SourceCheckpointDelta) -> Option { - let checkpoint_delta_str = format!("{checkpoint_delta:?}"); - let (_left, right) = - &checkpoint_delta_str[..checkpoint_delta_str.len() - 2].rsplit_once('(')?; - Some(right.to_string()) - } - #[tokio::test] async fn test_file_source_resume_from_checkpoint() { aux_test_file_source_resume_from_checkpoint(false).await; @@ -408,27 +347,10 @@ mod tests { quickwit_common::setup_logging_for_tests(); let universe = Universe::with_accelerated_time(); let (doc_processor_mailbox, doc_processor_inbox) = universe.create_test_mailbox(); - let mut documents_bytes = Vec::new(); - for i in 0..100 { - documents_bytes - .write_all(format!("{i}\n").as_bytes()) - .unwrap(); - } - let mut temp_file: tempfile::NamedTempFile = if gzip { - tempfile::Builder::new().suffix(".gz").tempfile().unwrap() - } else { - tempfile::NamedTempFile::new().unwrap() - }; - let temp_file_path = temp_file.path().canonicalize().unwrap(); - if gzip { - let gzipped_documents = gzip_bytes(&documents_bytes).await; - temp_file.write_all(&gzipped_documents).unwrap(); - } else { - temp_file.write_all(&documents_bytes).unwrap(); - } - temp_file.flush().unwrap(); - - let params = FileSourceParams::file(&temp_file_path); + let temp_file = generate_index_doc_file(gzip, 100).await; + let temp_file_path = temp_file.path().to_str().unwrap(); + let uri = Uri::from_str(temp_file_path).unwrap(); + let params = FileSourceParams::Filepath(uri.clone()); let source_config = SourceConfig { source_id: "test-file-source".to_string(), num_pipelines: NonZeroUsize::new(1).unwrap(), @@ -437,11 +359,11 @@ mod tests { transform_config: None, input_format: SourceInputFormat::Json, }; - let partition_id = PartitionId::from(temp_file_path.to_string_lossy().borrow()); + let partition_id = PartitionId::from(uri.as_str()); let source_checkpoint_delta = SourceCheckpointDelta::from_partition_delta( partition_id, Position::Beginning, - Position::offset(4u64), + Position::offset(16usize), ) .unwrap(); @@ -465,74 +387,93 @@ mod tests { assert_eq!( counters, serde_json::json!({ - "previous_offset": 290u64, - "current_offset": 290u64, - "num_lines_processed": 98u64 + "num_bytes_processed": (800-16) as u64, + "num_lines_processed": (100-2) as u64, }) ); let indexer_messages: Vec = doc_processor_inbox.drain_for_test_typed(); - assert!(&indexer_messages[0].docs[0].starts_with(b"2\n")); + assert_eq!( + indexer_messages[0].docs[0], + Bytes::from_static(b"0000002\n") + ); } +} - async fn gzip_bytes(bytes: &[u8]) -> Vec { - let mut gzip_documents = Vec::new(); - let mut encoder = GzipEncoder::new(&mut gzip_documents); - tokio::io::AsyncWriteExt::write_all(&mut encoder, bytes) - .await - .unwrap(); - // flush is not sufficient here and reading the file will raise a unexpected end of file - // error. - tokio::io::AsyncWriteExt::shutdown(&mut encoder) +#[cfg(all(test, feature = "sqs-localstack-tests"))] +mod localstack_tests { + use std::str::FromStr; + + use quickwit_actors::Universe; + use quickwit_common::rand::append_random_suffix; + use quickwit_common::uri::Uri; + use quickwit_config::{ + FileSourceMessageType, FileSourceNotification, FileSourceSqs, SourceConfig, SourceParams, + }; + use quickwit_metastore::metastore_for_test; + + use super::*; + use crate::models::RawDocBatch; + use crate::source::doc_file_reader::file_test_helpers::generate_dummy_doc_file; + use crate::source::queue_sources::sqs_queue::test_helpers::{ + create_queue, get_localstack_sqs_client, send_message, + }; + use crate::source::test_setup_helper::setup_index; + use crate::source::tests::SourceRuntimeBuilder; + use crate::source::SourceActor; + + #[tokio::test] + async fn test_file_source_sqs_notifications() { + // queue setup + let sqs_client = get_localstack_sqs_client().await.unwrap(); + let queue_url = create_queue(&sqs_client, "file-source-sqs-notifications").await; + let (dummy_doc_file, _) = generate_dummy_doc_file(false, 10).await; + let test_uri = Uri::from_str(dummy_doc_file.path().to_str().unwrap()).unwrap(); + send_message(&sqs_client, &queue_url, test_uri.as_str()).await; + + // source setup + let source_params = + FileSourceParams::Notifications(FileSourceNotification::Sqs(FileSourceSqs { + queue_url, + message_type: FileSourceMessageType::RawUri, + })); + let source_config = SourceConfig::for_test( + "test-file-source-sqs-notifications", + SourceParams::File(source_params.clone()), + ); + let metastore = metastore_for_test(); + let index_id = append_random_suffix("test-sqs-index"); + 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 sqs_source = FileSourceFactory::typed_create_source(source_runtime, source_params) .await .unwrap(); - gzip_documents - } - #[tokio::test] - async fn test_skip_reader() { - { - // Skip 0 bytes. - let mut reader = FileSourceReader::new(Box::new("hello".as_bytes()), 0); - let mut buf = String::new(); - reader.read_line(&mut buf).await.unwrap(); - assert_eq!(buf, "hello"); - } - { - // Skip 2 bytes. - let mut reader = FileSourceReader::new(Box::new("hello".as_bytes()), 2); - let mut buf = String::new(); - reader.read_line(&mut buf).await.unwrap(); - assert_eq!(buf, "llo"); - } - { - let input = "hello"; - let cursor = Cursor::new(input); - let mut reader = FileSourceReader::new(Box::new(cursor), 5); - let mut buf = String::new(); - assert!(reader.read_line(&mut buf).await.is_ok()); - } - { - let input = "hello"; - let cursor = Cursor::new(input); - let mut reader = FileSourceReader::new(Box::new(cursor), 10); - let mut buf = String::new(); - assert!(reader.read_line(&mut buf).await.is_err()); - } - { - let input = "hello world".repeat(10000); - let cursor = Cursor::new(input.clone()); - let mut reader = FileSourceReader::new(Box::new(cursor), 64000); - let mut buf = String::new(); - reader.read_line(&mut buf).await.unwrap(); - assert_eq!(buf, input[64000..]); - } + // actor setup + let universe = Universe::with_accelerated_time(); + let (doc_processor_mailbox, doc_processor_inbox) = universe.create_test_mailbox(); { - let input = "hello world".repeat(10000); - let cursor = Cursor::new(input.clone()); - let mut reader = FileSourceReader::new(Box::new(cursor), 64001); - let mut buf = String::new(); - reader.read_line(&mut buf).await.unwrap(); - assert_eq!(buf, input[64001..]); + let actor = SourceActor { + source: Box::new(sqs_source), + doc_processor_mailbox: doc_processor_mailbox.clone(), + }; + let (_mailbox, handle) = universe.spawn_builder().spawn(actor); + + // run the source actor for a while + tokio::time::timeout(Duration::from_millis(500), handle.join()) + .await + .unwrap_err(); + + let next_message = doc_processor_inbox + .drain_for_test() + .into_iter() + .flat_map(|box_any| box_any.downcast::().ok()) + .map(|box_raw_doc_batch| *box_raw_doc_batch) + .next() + .unwrap(); + assert_eq!(next_message.docs.len(), 10); } + universe.assert_quit().await; } } diff --git a/quickwit/quickwit-indexing/src/source/ingest/mod.rs b/quickwit/quickwit-indexing/src/source/ingest/mod.rs index 1d1855bbaed..76aafab3344 100644 --- a/quickwit/quickwit-indexing/src/source/ingest/mod.rs +++ b/quickwit/quickwit-indexing/src/source/ingest/mod.rs @@ -677,7 +677,7 @@ mod tests { use quickwit_common::metrics::MEMORY_METRICS; use quickwit_common::stream_utils::InFlightValue; use quickwit_common::ServiceStream; - use quickwit_config::{SourceConfig, SourceParams}; + use quickwit_config::{IndexingSettings, SourceConfig, SourceParams}; use quickwit_proto::indexing::IndexingPipelineId; use quickwit_proto::ingest::ingester::{ FetchMessage, IngesterServiceClient, MockIngesterService, TruncateShardsResponse, @@ -944,6 +944,7 @@ mod tests { queues_dir_path: PathBuf::from("./queues"), storage_resolver: StorageResolver::for_test(), event_broker, + indexing_setting: IndexingSettings::default(), }; let retry_params = RetryParams::no_retries(); let mut source = IngestSource::try_new(source_runtime, retry_params) @@ -1145,6 +1146,7 @@ mod tests { queues_dir_path: PathBuf::from("./queues"), storage_resolver: StorageResolver::for_test(), event_broker, + indexing_setting: IndexingSettings::default(), }; let retry_params = RetryParams::for_test(); let mut source = IngestSource::try_new(source_runtime, retry_params) @@ -1307,6 +1309,7 @@ mod tests { queues_dir_path: PathBuf::from("./queues"), storage_resolver: StorageResolver::for_test(), event_broker, + indexing_setting: IndexingSettings::default(), }; let retry_params = RetryParams::for_test(); let mut source = IngestSource::try_new(source_runtime, retry_params) @@ -1372,6 +1375,7 @@ mod tests { queues_dir_path: PathBuf::from("./queues"), storage_resolver: StorageResolver::for_test(), event_broker, + indexing_setting: IndexingSettings::default(), }; let retry_params = RetryParams::for_test(); let mut source = IngestSource::try_new(source_runtime, retry_params) @@ -1604,6 +1608,7 @@ mod tests { queues_dir_path: PathBuf::from("./queues"), storage_resolver: StorageResolver::for_test(), event_broker, + indexing_setting: IndexingSettings::default(), }; let retry_params = RetryParams::for_test(); let mut source = IngestSource::try_new(source_runtime, retry_params) @@ -1758,6 +1763,7 @@ mod tests { queues_dir_path: PathBuf::from("./queues"), storage_resolver: StorageResolver::for_test(), event_broker, + indexing_setting: IndexingSettings::default(), }; let retry_params = RetryParams::for_test(); let mut source = IngestSource::try_new(source_runtime, retry_params) @@ -1889,6 +1895,7 @@ mod tests { queues_dir_path: PathBuf::from("./queues"), storage_resolver: StorageResolver::for_test(), event_broker: event_broker.clone(), + indexing_setting: IndexingSettings::default(), }; let retry_params = RetryParams::for_test(); let mut source = IngestSource::try_new(source_runtime, retry_params) diff --git a/quickwit/quickwit-indexing/src/source/kafka_source.rs b/quickwit/quickwit-indexing/src/source/kafka_source.rs index ea4e26e77c6..6aae5b7c5bb 100644 --- a/quickwit/quickwit-indexing/src/source/kafka_source.rs +++ b/quickwit/quickwit-indexing/src/source/kafka_source.rs @@ -766,15 +766,9 @@ mod kafka_broker_tests { use quickwit_actors::{ActorContext, Universe}; use quickwit_common::rand::append_random_suffix; - use quickwit_config::{IndexConfig, SourceConfig, SourceInputFormat, SourceParams}; - use quickwit_metastore::checkpoint::{IndexCheckpointDelta, SourceCheckpointDelta}; - use quickwit_metastore::{ - metastore_for_test, CreateIndexRequestExt, SplitMetadata, StageSplitsRequestExt, - }; - use quickwit_proto::metastore::{ - CreateIndexRequest, MetastoreService, MetastoreServiceClient, PublishSplitsRequest, - StageSplitsRequest, - }; + use quickwit_config::{SourceConfig, SourceInputFormat, SourceParams}; + use quickwit_metastore::checkpoint::SourceCheckpointDelta; + use quickwit_metastore::metastore_for_test; use quickwit_proto::types::IndexUid; use rdkafka::admin::{AdminClient, AdminOptions, NewTopic, TopicReplication}; use rdkafka::client::DefaultClientContext; @@ -783,7 +777,7 @@ mod kafka_broker_tests { use tokio::sync::watch; use super::*; - use crate::new_split_id; + use crate::source::test_setup_helper::setup_index; use crate::source::tests::SourceRuntimeBuilder; use crate::source::{quickwit_supported_sources, RawDocBatch, SourceActor}; @@ -915,71 +909,6 @@ mod kafka_broker_tests { Ok(merged_batch) } - async fn setup_index( - metastore: MetastoreServiceClient, - index_id: &str, - source_config: &SourceConfig, - partition_deltas: &[(u64, i64, i64)], - ) -> IndexUid { - let index_uri = format!("ram:///indexes/{index_id}"); - let index_config = IndexConfig::for_test(index_id, &index_uri); - let create_index_request = CreateIndexRequest::try_from_index_and_source_configs( - &index_config, - &[source_config.clone()], - ) - .unwrap(); - let index_uid: IndexUid = metastore - .create_index(create_index_request) - .await - .unwrap() - .index_uid() - .clone(); - - if partition_deltas.is_empty() { - return index_uid; - } - let split_id = new_split_id(); - let split_metadata = SplitMetadata::for_test(split_id.clone()); - let stage_splits_request = - StageSplitsRequest::try_from_split_metadata(index_uid.clone(), &split_metadata) - .unwrap(); - metastore.stage_splits(stage_splits_request).await.unwrap(); - - let mut source_delta = SourceCheckpointDelta::default(); - for (partition_id, from_position, to_position) in partition_deltas.iter().copied() { - source_delta - .record_partition_delta( - partition_id.into(), - { - if from_position < 0 { - Position::Beginning - } else { - Position::offset(from_position as u64) - } - }, - Position::offset(to_position as u64), - ) - .unwrap(); - } - let checkpoint_delta = IndexCheckpointDelta { - source_id: source_config.source_id.to_string(), - source_delta, - }; - let checkpoint_delta_json = serde_json::to_string(&checkpoint_delta).unwrap(); - let publish_splits_request = PublishSplitsRequest { - index_uid: Some(index_uid.clone()), - index_checkpoint_delta_json_opt: Some(checkpoint_delta_json), - staged_split_ids: vec![split_id.clone()], - replaced_split_ids: Vec::new(), - publish_token_opt: None, - }; - metastore - .publish_splits(publish_splits_request) - .await - .unwrap(); - index_uid - } - #[tokio::test] async fn test_kafka_source_process_message() { let admin_client = create_admin_client(); @@ -1110,8 +1039,17 @@ mod kafka_broker_tests { let index_id = append_random_suffix("test-kafka-source--process-assign-partitions--index"); let (_source_id, source_config) = get_source_config(&topic, "earliest"); - let index_uid = - setup_index(metastore.clone(), &index_id, &source_config, &[(2, -1, 42)]).await; + let index_uid = setup_index( + metastore.clone(), + &index_id, + &source_config, + &[( + PartitionId::from(2u64), + Position::Beginning, + Position::offset(42u64), + )], + ) + .await; let SourceParams::Kafka(params) = source_config.clone().source_params else { panic!( @@ -1243,8 +1181,17 @@ mod kafka_broker_tests { let metastore = metastore_for_test(); let index_id = append_random_suffix("test-kafka-source--suggest-truncate--index"); let (_source_id, source_config) = get_source_config(&topic, "earliest"); - let index_uid = - setup_index(metastore.clone(), &index_id, &source_config, &[(2, -1, 42)]).await; + let index_uid = setup_index( + metastore.clone(), + &index_id, + &source_config, + &[( + PartitionId::from(2u64), + Position::Beginning, + Position::offset(42u64), + )], + ) + .await; let SourceParams::Kafka(params) = source_config.clone().source_params else { panic!( @@ -1436,7 +1383,18 @@ mod kafka_broker_tests { metastore.clone(), &index_id, &source_config, - &[(0, -1, 0), (1, -1, 2)], + &[ + ( + PartitionId::from(0u64), + Position::Beginning, + Position::offset(0u64), + ), + ( + PartitionId::from(1u64), + Position::Beginning, + Position::offset(2u64), + ), + ], ) .await; let source_runtime = SourceRuntimeBuilder::new(index_uid, source_config) diff --git a/quickwit/quickwit-indexing/src/source/mod.rs b/quickwit/quickwit-indexing/src/source/mod.rs index db2583d0e95..20d1bad6fc4 100644 --- a/quickwit/quickwit-indexing/src/source/mod.rs +++ b/quickwit/quickwit-indexing/src/source/mod.rs @@ -57,6 +57,7 @@ //! that file. //! - the kafka source: the partition id is a kafka topic partition id, and the position is a kafka //! offset. +mod doc_file_reader; mod file_source; #[cfg(feature = "gcp-pubsub")] mod gcp_pubsub_source; @@ -68,7 +69,10 @@ mod kafka_source; mod kinesis; #[cfg(feature = "pulsar")] mod pulsar_source; +#[cfg(feature = "queue-sources")] +mod queue_sources; mod source_factory; +mod stdin_source; mod vec_source; mod void_source; @@ -89,11 +93,15 @@ pub use kinesis::kinesis_source::{KinesisSource, KinesisSourceFactory}; use once_cell::sync::OnceCell; #[cfg(feature = "pulsar")] pub use pulsar_source::{PulsarSource, PulsarSourceFactory}; +#[cfg(feature = "sqs")] +pub use queue_sources::sqs_queue; use quickwit_actors::{Actor, ActorContext, ActorExitStatus, Handler, Mailbox}; use quickwit_common::metrics::{GaugeGuard, MEMORY_METRICS}; use quickwit_common::pubsub::EventBroker; use quickwit_common::runtimes::RuntimeType; -use quickwit_config::{SourceConfig, SourceParams}; +use quickwit_config::{ + FileSourceNotification, FileSourceParams, IndexingSettings, SourceConfig, SourceParams, +}; use quickwit_ingest::IngesterPool; use quickwit_metastore::checkpoint::{SourceCheckpoint, SourceCheckpointDelta}; use quickwit_metastore::IndexMetadataResponseExt; @@ -111,7 +119,7 @@ use tracing::error; pub use vec_source::{VecSource, VecSourceFactory}; pub use void_source::{VoidSource, VoidSourceFactory}; -use self::file_source::dir_and_filename; +use self::doc_file_reader::dir_and_filename; use crate::actors::DocProcessor; use crate::models::RawDocBatch; use crate::source::ingest::IngestSourceFactory; @@ -143,6 +151,7 @@ pub struct SourceRuntime { pub queues_dir_path: PathBuf, pub storage_resolver: StorageResolver, pub event_broker: EventBroker, + pub indexing_setting: IndexingSettings, } impl SourceRuntime { @@ -232,7 +241,7 @@ pub trait Source: Send + 'static { /// In that case, `batch_sink` will block. /// /// It returns an optional duration specifying how long the batch requester - /// should wait before pooling gain. + /// should wait before polling again. async fn emit_batches( &mut self, doc_processor_mailbox: &Mailbox, @@ -410,15 +419,26 @@ pub async fn check_source_connectivity( source_config: &SourceConfig, ) -> anyhow::Result<()> { match &source_config.source_params { - SourceParams::File(params) => { - if let Some(filepath) = ¶ms.filepath { - let (dir_uri, file_name) = dir_and_filename(filepath)?; - let storage = storage_resolver.resolve(&dir_uri).await?; - storage.file_num_bytes(file_name).await?; - } + SourceParams::File(FileSourceParams::Filepath(file_uri)) => { + let (dir_uri, file_name) = dir_and_filename(file_uri)?; + let storage = storage_resolver.resolve(&dir_uri).await?; + storage.file_num_bytes(file_name).await?; Ok(()) } #[allow(unused_variables)] + SourceParams::File(FileSourceParams::Notifications(FileSourceNotification::Sqs( + sqs_config, + ))) => { + #[cfg(not(feature = "sqs"))] + anyhow::bail!("Quickwit was compiled without the `sqs` feature"); + + #[cfg(feature = "sqs")] + { + queue_sources::sqs_queue::check_connectivity(&sqs_config.queue_url).await?; + Ok(()) + } + } + #[allow(unused_variables)] SourceParams::Kafka(params) => { #[cfg(not(feature = "kafka"))] anyhow::bail!("Quickwit was compiled without the `kafka` feature"); @@ -591,10 +611,11 @@ mod tests { source_config: self.source_config, storage_resolver: StorageResolver::for_test(), event_broker: EventBroker::default(), + indexing_setting: IndexingSettings::default(), } } - #[cfg(feature = "kafka")] + #[cfg(any(feature = "kafka", feature = "sqs"))] pub fn with_metastore(mut self, metastore: MetastoreServiceClient) -> Self { self.metastore_opt = Some(metastore); self @@ -676,7 +697,7 @@ mod tests { source_id: "file".to_string(), num_pipelines: NonZeroUsize::new(1).unwrap(), enabled: true, - source_params: SourceParams::file("file-does-not-exist.json"), + source_params: SourceParams::file_from_str("file-does-not-exist.json").unwrap(), transform_config: None, input_format: SourceInputFormat::Json, }; @@ -691,7 +712,7 @@ mod tests { source_id: "file".to_string(), num_pipelines: NonZeroUsize::new(1).unwrap(), enabled: true, - source_params: SourceParams::file("data/test_corpus.json"), + source_params: SourceParams::file_from_str("data/test_corpus.json").unwrap(), transform_config: None, input_format: SourceInputFormat::Json, }; @@ -704,3 +725,74 @@ mod tests { Ok(()) } } + +#[cfg(all( + test, + any(feature = "sqs-localstack-tests", feature = "kafka-broker-tests") +))] +mod test_setup_helper { + + use quickwit_config::IndexConfig; + use quickwit_metastore::checkpoint::{IndexCheckpointDelta, PartitionId}; + use quickwit_metastore::{CreateIndexRequestExt, SplitMetadata, StageSplitsRequestExt}; + use quickwit_proto::metastore::{CreateIndexRequest, PublishSplitsRequest, StageSplitsRequest}; + use quickwit_proto::types::Position; + + use super::*; + use crate::new_split_id; + + pub async fn setup_index( + metastore: MetastoreServiceClient, + index_id: &str, + source_config: &SourceConfig, + partition_deltas: &[(PartitionId, Position, Position)], + ) -> IndexUid { + let index_uri = format!("ram:///indexes/{index_id}"); + let index_config = IndexConfig::for_test(index_id, &index_uri); + let create_index_request = CreateIndexRequest::try_from_index_and_source_configs( + &index_config, + &[source_config.clone()], + ) + .unwrap(); + let index_uid: IndexUid = metastore + .create_index(create_index_request) + .await + .unwrap() + .index_uid() + .clone(); + + if partition_deltas.is_empty() { + return index_uid; + } + let split_id = new_split_id(); + let split_metadata = SplitMetadata::for_test(split_id.clone()); + let stage_splits_request = + StageSplitsRequest::try_from_split_metadata(index_uid.clone(), &split_metadata) + .unwrap(); + metastore.stage_splits(stage_splits_request).await.unwrap(); + + let mut source_delta = SourceCheckpointDelta::default(); + for (partition_id, from_position, to_position) in partition_deltas.iter().cloned() { + source_delta + .record_partition_delta(partition_id, from_position, to_position) + .unwrap(); + } + let checkpoint_delta = IndexCheckpointDelta { + source_id: source_config.source_id.to_string(), + source_delta, + }; + let checkpoint_delta_json = serde_json::to_string(&checkpoint_delta).unwrap(); + let publish_splits_request = PublishSplitsRequest { + index_uid: Some(index_uid.clone()), + index_checkpoint_delta_json_opt: Some(checkpoint_delta_json), + staged_split_ids: vec![split_id.clone()], + replaced_split_ids: Vec::new(), + publish_token_opt: None, + }; + metastore + .publish_splits(publish_splits_request) + .await + .unwrap(); + index_uid + } +} diff --git a/quickwit/quickwit-indexing/src/source/pulsar_source.rs b/quickwit/quickwit-indexing/src/source/pulsar_source.rs index 6dec087a454..6dcc91abc71 100644 --- a/quickwit/quickwit-indexing/src/source/pulsar_source.rs +++ b/quickwit/quickwit-indexing/src/source/pulsar_source.rs @@ -445,22 +445,15 @@ mod pulsar_broker_tests { use futures::future::join_all; use quickwit_actors::{ActorHandle, Inbox, Universe, HEARTBEAT}; use quickwit_common::rand::append_random_suffix; - use quickwit_config::{IndexConfig, SourceConfig, SourceInputFormat, SourceParams}; - use quickwit_metastore::checkpoint::{ - IndexCheckpointDelta, PartitionId, SourceCheckpointDelta, - }; - use quickwit_metastore::{ - metastore_for_test, CreateIndexRequestExt, SplitMetadata, StageSplitsRequestExt, - }; - use quickwit_proto::metastore::{ - CreateIndexRequest, MetastoreService, MetastoreServiceClient, PublishSplitsRequest, - StageSplitsRequest, - }; + use quickwit_config::{SourceConfig, SourceInputFormat, SourceParams}; + use quickwit_metastore::checkpoint::{PartitionId, SourceCheckpointDelta}; + use quickwit_metastore::metastore_for_test; + use quickwit_proto::metastore::MetastoreServiceClient; use reqwest::StatusCode; use super::*; - use crate::new_split_id; use crate::source::pulsar_source::{msg_id_from_position, msg_id_to_position}; + use crate::source::test_setup_helper::setup_index; use crate::source::tests::SourceRuntimeBuilder; use crate::source::{quickwit_supported_sources, RawDocBatch, SuggestTruncate}; @@ -492,63 +485,6 @@ mod pulsar_broker_tests { }}; } - async fn setup_index( - metastore: MetastoreServiceClient, - index_id: &str, - source_id: &str, - partition_deltas: &[(&str, Position, Position)], - ) -> IndexUid { - let index_uri = format!("ram:///indexes/{index_id}"); - let index_config = IndexConfig::for_test(index_id, &index_uri); - let create_index_request = - CreateIndexRequest::try_from_index_config(&index_config).unwrap(); - let index_uid: IndexUid = metastore - .create_index(create_index_request) - .await - .unwrap() - .index_uid() - .clone(); - - if partition_deltas.is_empty() { - return index_uid; - } - let split_id = new_split_id(); - let split_metadata = SplitMetadata::for_test(split_id.clone()); - let stage_splits_request = - StageSplitsRequest::try_from_split_metadata(index_uid.clone(), &split_metadata) - .unwrap(); - metastore.stage_splits(stage_splits_request).await.unwrap(); - - let mut source_delta = SourceCheckpointDelta::default(); - for (partition_id, from_position, to_position) in partition_deltas { - source_delta - .record_partition_delta( - PartitionId::from(&**partition_id), - from_position.clone(), - to_position.clone(), - ) - .unwrap(); - } - let checkpoint_delta = IndexCheckpointDelta { - source_id: source_id.to_string(), - source_delta, - }; - let publish_splits_request = PublishSplitsRequest { - index_uid: Some(index_uid.clone()), - staged_split_ids: vec![split_id.clone()], - replaced_split_ids: Vec::new(), - index_checkpoint_delta_json_opt: Some( - serde_json::to_string(&checkpoint_delta).unwrap(), - ), - publish_token_opt: None, - }; - metastore - .publish_splits(publish_splits_request) - .await - .unwrap(); - index_uid - } - fn get_source_config>( topics: impl IntoIterator, ) -> (String, SourceConfig) { @@ -895,7 +831,7 @@ mod pulsar_broker_tests { let index_id = append_random_suffix("test-pulsar-source--topic-ingestion--index"); let (source_id, source_config) = get_source_config([&topic]); - let index_uid = setup_index(metastore.clone(), &index_id, &source_id, &[]).await; + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; let (source_handle, doc_processor_inbox) = create_source( &universe, @@ -952,7 +888,7 @@ mod pulsar_broker_tests { let index_id = append_random_suffix("test-pulsar-source--topic-ingestion--index"); let (source_id, source_config) = get_source_config([&topic1, &topic2]); - let index_uid = setup_index(metastore.clone(), &index_id, &source_id, &[]).await; + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; let (source_handle, doc_processor_inbox) = create_source( &universe, @@ -1020,7 +956,7 @@ mod pulsar_broker_tests { let (source_id, source_config) = get_source_config([&topic]); create_partitioned_topic(&topic, 2).await; - let index_uid = setup_index(metastore.clone(), &index_id, &source_id, &[]).await; + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; let (source_handle, doc_processor_inbox) = create_source( &universe, @@ -1074,7 +1010,7 @@ mod pulsar_broker_tests { let (source_id, source_config) = get_source_config([&topic]); create_partitioned_topic(&topic, 2).await; - let index_uid = setup_index(metastore.clone(), &index_id, &source_id, &[]).await; + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; let topic_partition_1 = format!("{topic}-partition-0"); let topic_partition_2 = format!("{topic}-partition-1"); @@ -1158,10 +1094,10 @@ mod pulsar_broker_tests { let index_id = append_random_suffix("test-pulsar-source--partitioned-multi-consumer-failure--index"); - let (source_id, source_config) = get_source_config([&topic]); + let (_, source_config) = get_source_config([&topic]); create_partitioned_topic(&topic, 2).await; - let index_uid = setup_index(metastore.clone(), &index_id, &source_id, &[]).await; + let index_uid = setup_index(metastore.clone(), &index_id, &source_config, &[]).await; let topic_partition_1 = format!("{topic}-partition-0"); let topic_partition_2 = format!("{topic}-partition-1"); diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/coordinator.rs b/quickwit/quickwit-indexing/src/source/queue_sources/coordinator.rs new file mode 100644 index 00000000000..bd00840f657 --- /dev/null +++ b/quickwit/quickwit-indexing/src/source/queue_sources/coordinator.rs @@ -0,0 +1,521 @@ +// Copyright (C) 2024 Quickwit, Inc. +// +// Quickwit is offered under the AGPL v3.0 and as commercial software. +// For commercial licensing, contact us at hello@quickwit.io. +// +// AGPL: +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as +// published by the Free Software Foundation, either version 3 of the +// License, or (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::fmt; +use std::sync::Arc; +use std::time::Duration; + +use itertools::Itertools; +use quickwit_actors::{ActorExitStatus, Mailbox}; +use quickwit_common::rate_limited_error; +use quickwit_config::{FileSourceMessageType, FileSourceSqs}; +use quickwit_metastore::checkpoint::SourceCheckpoint; +use quickwit_proto::indexing::IndexingPipelineId; +use quickwit_proto::metastore::SourceType; +use quickwit_storage::StorageResolver; +use serde::Serialize; +use ulid::Ulid; + +use super::helpers::QueueReceiver; +use super::local_state::QueueLocalState; +use super::message::{MessageType, PreProcessingError, ReadyMessage}; +use super::shared_state::{checkpoint_messages, QueueSharedState}; +use super::visibility::{spawn_visibility_task, VisibilitySettings}; +use super::Queue; +use crate::actors::DocProcessor; +use crate::models::{NewPublishLock, NewPublishToken, PublishLock}; +use crate::source::{SourceContext, SourceRuntime}; + +/// Maximum duration that the `emit_batches()` callback can wait for +/// `queue.receive()` calls. If too small, the actor loop will spin +/// un-necessarily. If too large, the actor loop will be slow to react to new +/// messages (or shutdown). +pub const RECEIVE_POLL_TIMEOUT: Duration = Duration::from_millis(500); + +#[derive(Default, Serialize)] +pub struct QueueCoordinatorObservableState { + /// Number of bytes processed by the source. + pub num_bytes_processed: u64, + /// Number of lines processed by the source. + pub num_lines_processed: u64, + /// Number of messages processed by the source. + pub num_messages_processed: u64, + /// Number of messages that could not be pre-processed. + pub num_messages_failed_preprocessing: u64, + /// Number of messages that could not be moved to in-progress. + pub num_messages_failed_opening: u64, +} + +/// The `QueueCoordinator` fetches messages from a queue, converts them into +/// record batches, and tracks the messages' state until their entire content is +/// published. Its API closely resembles the [`crate::source::Source`] trait, +/// making the implementation of queue sources straightforward. +pub struct QueueCoordinator { + storage_resolver: StorageResolver, + pipeline_id: IndexingPipelineId, + source_type: SourceType, + queue: Arc, + queue_receiver: QueueReceiver, + observable_state: QueueCoordinatorObservableState, + message_type: MessageType, + publish_lock: PublishLock, + shared_state: QueueSharedState, + local_state: QueueLocalState, + publish_token: String, + visible_settings: VisibilitySettings, +} + +impl fmt::Debug for QueueCoordinator { + fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter + .debug_struct("QueueTracker") + .field("index_id", &self.pipeline_id.index_uid.index_id) + .field("queue", &self.queue) + .finish() + } +} + +impl QueueCoordinator { + pub fn new( + source_runtime: SourceRuntime, + queue: Arc, + message_type: MessageType, + ) -> Self { + Self { + shared_state: QueueSharedState { + metastore: source_runtime.metastore, + source_id: source_runtime.pipeline_id.source_id.clone(), + index_uid: source_runtime.pipeline_id.index_uid.clone(), + }, + local_state: QueueLocalState::default(), + pipeline_id: source_runtime.pipeline_id, + source_type: source_runtime.source_config.source_type(), + storage_resolver: source_runtime.storage_resolver, + queue_receiver: QueueReceiver::new(queue.clone(), RECEIVE_POLL_TIMEOUT), + queue, + observable_state: QueueCoordinatorObservableState::default(), + message_type, + publish_lock: PublishLock::default(), + publish_token: Ulid::new().to_string(), + visible_settings: VisibilitySettings::from_commit_timeout( + source_runtime.indexing_setting.commit_timeout_secs, + ), + } + } + + #[cfg(feature = "sqs")] + pub async fn try_from_sqs_config( + config: FileSourceSqs, + source_runtime: SourceRuntime, + ) -> anyhow::Result { + use super::sqs_queue::SqsQueue; + let queue = SqsQueue::try_new(config.queue_url).await?; + let message_type = match config.message_type { + FileSourceMessageType::S3Notification => MessageType::S3Notification, + FileSourceMessageType::RawUri => MessageType::RawUri, + }; + Ok(QueueCoordinator::new( + source_runtime, + Arc::new(queue), + message_type, + )) + } + + pub async fn initialize( + &mut self, + doc_processor_mailbox: &Mailbox, + ctx: &SourceContext, + ) -> Result<(), ActorExitStatus> { + let publish_lock = self.publish_lock.clone(); + ctx.send_message(doc_processor_mailbox, NewPublishLock(publish_lock)) + .await?; + ctx.send_message( + doc_processor_mailbox, + NewPublishToken(self.publish_token.clone()), + ) + .await?; + Ok(()) + } + + /// Polls messages from the queue and prepares them for processing + async fn poll_messages(&mut self, ctx: &SourceContext) -> Result<(), ActorExitStatus> { + let raw_messages = self + .queue_receiver + .receive(1, self.visible_settings.deadline_for_receive) + .await?; + + let mut format_errors = Vec::new(); + let mut discardable_ack_ids = Vec::new(); + let mut preprocessed_messages = Vec::new(); + for message in raw_messages { + match message.pre_process(self.message_type) { + Ok(preprocessed_message) => preprocessed_messages.push(preprocessed_message), + Err(PreProcessingError::UnexpectedFormat(err)) => format_errors.push(err), + Err(PreProcessingError::Discardable { ack_id }) => discardable_ack_ids.push(ack_id), + } + } + if !format_errors.is_empty() { + self.observable_state.num_messages_failed_preprocessing += format_errors.len() as u64; + rate_limited_error!( + limit_per_min = 10, + count = format_errors.len(), + last_err = ?format_errors.last().unwrap(), + "invalid messages not processed, use a dead letter queue to limit retries" + ); + } + if preprocessed_messages.is_empty() { + self.queue.acknowledge(&discardable_ack_ids).await?; + return Ok(()); + } + + // in rare situations, there might be duplicates within a batch + let deduplicated_messages = preprocessed_messages + .into_iter() + .unique_by(|x| x.partition_id()); + + let mut untracked_locally = Vec::new(); + let mut already_completed = Vec::new(); + for message in deduplicated_messages { + let partition_id = message.partition_id(); + if self.local_state.is_completed(&partition_id) { + already_completed.push(message); + } else if !self.local_state.is_tracked(&partition_id) { + untracked_locally.push(message); + } + } + + let checkpointed_messages = + checkpoint_messages(&self.shared_state, &self.publish_token, untracked_locally).await?; + + let mut ready_messages = Vec::new(); + for (message, position) in checkpointed_messages { + if position.is_eof() { + self.local_state.mark_completed(message.partition_id()); + already_completed.push(message); + } else { + ready_messages.push(ReadyMessage { + visibility_handle: spawn_visibility_task( + ctx, + self.queue.clone(), + message.metadata.ack_id.clone(), + message.metadata.initial_deadline, + self.visible_settings.clone(), + ), + content: message, + position, + }) + } + } + + self.local_state.set_ready_for_read(ready_messages); + + // Acknowledge messages that already have been processed + let mut ack_ids = already_completed + .iter() + .map(|msg| msg.metadata.ack_id.clone()) + .collect::>(); + ack_ids.append(&mut discardable_ack_ids); + self.queue.acknowledge(&ack_ids).await?; + + Ok(()) + } + + pub async fn emit_batches( + &mut self, + doc_processor_mailbox: &Mailbox, + ctx: &SourceContext, + ) -> Result { + if let Some(in_progress_ref) = self.local_state.read_in_progress_mut() { + // TODO: should we kill the publish lock if the message visibility extension failed? + let batch_builder = in_progress_ref + .batch_reader + .read_batch(ctx.progress(), self.source_type) + .await?; + self.observable_state.num_lines_processed += batch_builder.docs.len() as u64; + self.observable_state.num_bytes_processed += batch_builder.num_bytes; + doc_processor_mailbox + .send_message(batch_builder.build()) + .await?; + if in_progress_ref.batch_reader.is_eof() { + self.local_state + .drop_currently_read(self.visible_settings.deadline_for_last_extension) + .await?; + self.observable_state.num_messages_processed += 1; + } + } else if let Some(ready_message) = self.local_state.get_ready_for_read() { + match ready_message.start_processing(&self.storage_resolver).await { + Ok(new_in_progress) => { + self.local_state.set_currently_read(new_in_progress)?; + } + Err(err) => { + self.observable_state.num_messages_failed_opening += 1; + rate_limited_error!( + limit_per_min = 5, + err = ?err, + "failed to start message processing" + ); + } + } + } else { + self.poll_messages(ctx).await?; + } + + Ok(Duration::ZERO) + } + + pub async fn suggest_truncate( + &mut self, + checkpoint: SourceCheckpoint, + _ctx: &SourceContext, + ) -> anyhow::Result<()> { + let committed_partition_ids = checkpoint + .iter() + .filter(|(_, pos)| pos.is_eof()) + .map(|(pid, _)| pid) + .collect::>(); + let mut completed = Vec::new(); + for partition_id in committed_partition_ids { + let ack_id_opt = self.local_state.mark_completed(partition_id); + if let Some(ack_id) = ack_id_opt { + completed.push(ack_id); + } + } + self.queue.acknowledge(&completed).await + } + + pub fn observable_state(&self) -> &QueueCoordinatorObservableState { + &self.observable_state + } +} + +#[cfg(test)] +mod tests { + use std::str::FromStr; + + use quickwit_actors::{ActorContext, Universe}; + use quickwit_common::uri::Uri; + use quickwit_proto::types::{NodeId, PipelineUid, Position}; + use tokio::sync::watch; + use ulid::Ulid; + + use super::*; + use crate::models::RawDocBatch; + use crate::source::doc_file_reader::file_test_helpers::{generate_dummy_doc_file, DUMMY_DOC}; + use crate::source::queue_sources::memory_queue::MemoryQueueForTests; + use crate::source::queue_sources::message::PreProcessedPayload; + use crate::source::queue_sources::shared_state::shared_state_for_tests::shared_state_for_tests; + use crate::source::{SourceActor, BATCH_NUM_BYTES_LIMIT}; + + fn setup_coordinator( + queue: Arc, + shared_state: QueueSharedState, + ) -> QueueCoordinator { + let pipeline_id = IndexingPipelineId { + node_id: NodeId::from_str("test-node").unwrap(), + index_uid: shared_state.index_uid.clone(), + source_id: shared_state.source_id.clone(), + pipeline_uid: PipelineUid::random(), + }; + + QueueCoordinator { + local_state: QueueLocalState::default(), + shared_state, + pipeline_id, + observable_state: QueueCoordinatorObservableState::default(), + publish_lock: PublishLock::default(), + // set a very high chunking timeout to make it possible to count the + // number of iterations required to process messages + queue_receiver: QueueReceiver::new(queue.clone(), Duration::from_secs(10)), + queue, + message_type: MessageType::RawUri, + source_type: SourceType::Unspecified, + storage_resolver: StorageResolver::for_test(), + publish_token: Ulid::new().to_string(), + visible_settings: VisibilitySettings::from_commit_timeout(5), + } + } + + async fn process_messages( + coordinator: &mut QueueCoordinator, + queue: Arc, + messages: &[(&Uri, &str)], + ) -> Vec { + let universe = Universe::with_accelerated_time(); + let (source_mailbox, _source_inbox) = universe.create_test_mailbox::(); + let (doc_processor_mailbox, doc_processor_inbox) = + universe.create_test_mailbox::(); + let (observable_state_tx, _observable_state_rx) = watch::channel(serde_json::Value::Null); + let ctx: SourceContext = + ActorContext::for_test(&universe, source_mailbox, observable_state_tx); + + coordinator + .initialize(&doc_processor_mailbox, &ctx) + .await + .unwrap(); + + coordinator + .emit_batches(&doc_processor_mailbox, &ctx) + .await + .unwrap(); + + for (uri, ack_id) in messages { + queue.send_message(uri.to_string(), ack_id); + } + + // Need 3 iterations for each msg to emit the first batch (receive, + // start, emit), assuming the `QueueReceiver` doesn't chunk the receive + // future. + for _ in 0..(messages.len() * 4) { + coordinator + .emit_batches(&doc_processor_mailbox, &ctx) + .await + .unwrap(); + } + + let batches = doc_processor_inbox + .drain_for_test() + .into_iter() + .flat_map(|box_any| box_any.downcast::().ok()) + .map(|box_raw_doc_batch| *box_raw_doc_batch) + .collect::>(); + universe.assert_quit().await; + batches + } + + #[tokio::test] + async fn test_process_empty_queue() { + let queue = Arc::new(MemoryQueueForTests::new()); + let shared_state = shared_state_for_tests("test-index", Default::default()); + let mut coordinator = setup_coordinator(queue.clone(), shared_state); + let batches = process_messages(&mut coordinator, queue, &[]).await; + assert_eq!(batches.len(), 0); + } + + #[tokio::test] + async fn test_process_one_small_message() { + let queue = Arc::new(MemoryQueueForTests::new()); + let shared_state = shared_state_for_tests("test-index", Default::default()); + let mut coordinator = setup_coordinator(queue.clone(), shared_state.clone()); + let (dummy_doc_file, _) = generate_dummy_doc_file(false, 10).await; + let test_uri = Uri::from_str(dummy_doc_file.path().to_str().unwrap()).unwrap(); + let partition_id = PreProcessedPayload::ObjectUri(test_uri.clone()).partition_id(); + let batches = process_messages(&mut coordinator, queue, &[(&test_uri, "ack-id")]).await; + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].docs.len(), 10); + assert!(coordinator.local_state.is_awaiting_commit(&partition_id)); + } + + #[tokio::test] + async fn test_process_one_big_message() { + let queue = Arc::new(MemoryQueueForTests::new()); + let shared_state = shared_state_for_tests("test-index", Default::default()); + let mut coordinator = setup_coordinator(queue.clone(), shared_state); + let lines = BATCH_NUM_BYTES_LIMIT as usize / DUMMY_DOC.len() + 1; + let (dummy_doc_file, _) = generate_dummy_doc_file(true, lines).await; + let test_uri = Uri::from_str(dummy_doc_file.path().to_str().unwrap()).unwrap(); + let batches = process_messages(&mut coordinator, queue, &[(&test_uri, "ack-id")]).await; + assert_eq!(batches.len(), 2); + assert_eq!(batches.iter().map(|b| b.docs.len()).sum::(), lines); + } + + #[tokio::test] + async fn test_process_two_messages_different_compression() { + let queue = Arc::new(MemoryQueueForTests::new()); + let shared_state = shared_state_for_tests("test-index", Default::default()); + let mut coordinator = setup_coordinator(queue.clone(), shared_state); + let (dummy_doc_file_1, _) = generate_dummy_doc_file(false, 10).await; + let test_uri_1 = Uri::from_str(dummy_doc_file_1.path().to_str().unwrap()).unwrap(); + let (dummy_doc_file_2, _) = generate_dummy_doc_file(true, 10).await; + let test_uri_2 = Uri::from_str(dummy_doc_file_2.path().to_str().unwrap()).unwrap(); + let batches = process_messages( + &mut coordinator, + queue, + &[(&test_uri_1, "ack-id-1"), (&test_uri_2, "ack-id-2")], + ) + .await; + // could be generated in 1 or 2 batches, it doesn't matter + assert_eq!(batches.iter().map(|b| b.docs.len()).sum::(), 20); + } + + #[tokio::test] + async fn test_process_local_duplicate_message() { + let queue = Arc::new(MemoryQueueForTests::new()); + let shared_state = shared_state_for_tests("test-index", Default::default()); + let mut coordinator = setup_coordinator(queue.clone(), shared_state); + let (dummy_doc_file, _) = generate_dummy_doc_file(false, 10).await; + let test_uri = Uri::from_str(dummy_doc_file.path().to_str().unwrap()).unwrap(); + let batches = process_messages( + &mut coordinator, + queue, + &[(&test_uri, "ack-id-1"), (&test_uri, "ack-id-2")], + ) + .await; + assert_eq!(batches.len(), 1); + assert_eq!(batches.iter().map(|b| b.docs.len()).sum::(), 10); + } + + #[tokio::test] + async fn test_process_shared_complete_message() { + let (dummy_doc_file, file_size) = generate_dummy_doc_file(false, 10).await; + let test_uri = Uri::from_str(dummy_doc_file.path().to_str().unwrap()).unwrap(); + let partition_id = PreProcessedPayload::ObjectUri(test_uri.clone()).partition_id(); + + let queue = Arc::new(MemoryQueueForTests::new()); + let shared_state = shared_state_for_tests( + "test-index", + &[( + partition_id.clone(), + ("existing_token".to_string(), Position::eof(file_size)), + )], + ); + let mut coordinator = setup_coordinator(queue.clone(), shared_state.clone()); + + assert!(!coordinator.local_state.is_tracked(&partition_id)); + let batches = process_messages(&mut coordinator, queue, &[(&test_uri, "ack-id-1")]).await; + assert_eq!(batches.len(), 0); + assert!(coordinator.local_state.is_completed(&partition_id)); + } + + #[tokio::test] + async fn test_process_multiple_coordinator() { + let queue = Arc::new(MemoryQueueForTests::new()); + let shared_state = shared_state_for_tests("test-index", Default::default()); + let mut proc_1 = setup_coordinator(queue.clone(), shared_state.clone()); + let mut proc_2 = setup_coordinator(queue.clone(), shared_state.clone()); + let (dummy_doc_file, _) = generate_dummy_doc_file(false, 10).await; + let test_uri = Uri::from_str(dummy_doc_file.path().to_str().unwrap()).unwrap(); + let partition_id = PreProcessedPayload::ObjectUri(test_uri.clone()).partition_id(); + + let batches_1 = process_messages(&mut proc_1, queue.clone(), &[(&test_uri, "ack1")]).await; + let batches_2 = process_messages(&mut proc_2, queue, &[(&test_uri, "ack2")]).await; + + assert_eq!(batches_1.len(), 1); + assert_eq!(batches_1[0].docs.len(), 10); + assert!(proc_1.local_state.is_awaiting_commit(&partition_id)); + // proc_2 doesn't know for sure what is happening with the message + // (proc_1 might have crashed), so it just acquires it and takes over + // processing + // + // TODO: this test should fail once we implement the grace + // period before a partition can be re-acquired + assert_eq!(batches_2.len(), 1); + assert_eq!(batches_2[0].docs.len(), 10); + assert!(proc_2.local_state.is_awaiting_commit(&partition_id)); + } +} diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/helpers.rs b/quickwit/quickwit-indexing/src/source/queue_sources/helpers.rs new file mode 100644 index 00000000000..8a215f66710 --- /dev/null +++ b/quickwit/quickwit-indexing/src/source/queue_sources/helpers.rs @@ -0,0 +1,130 @@ +// Copyright (C) 2024 Quickwit, Inc. +// +// Quickwit is offered under the AGPL v3.0 and as commercial software. +// For commercial licensing, contact us at hello@quickwit.io. +// +// AGPL: +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as +// published by the Free Software Foundation, either version 3 of the +// License, or (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::sync::Arc; +use std::time::Duration; + +use futures::future::BoxFuture; + +use super::message::RawMessage; +use super::Queue; + +type ReceiveResult = anyhow::Result>; + +/// A statefull wrapper around a `Queue` that chunks the slow `receive()` call +/// into shorter iterations. This enables yielding back to the actor system +/// without compromising on queue poll durations. Without this, an actor that +/// tries to receive messages from a `Queue` will be blocked for multiple seconds +/// before being able to process new mailbox messages (or shutting down). +pub struct QueueReceiver { + queue: Arc, + receive: Option>, + iteration: Duration, +} + +impl QueueReceiver { + pub fn new(queue: Arc, iteration: Duration) -> Self { + Self { + queue, + receive: None, + iteration, + } + } + + pub async fn receive( + &mut self, + max_messages: usize, + suggested_deadline: Duration, + ) -> anyhow::Result> { + if self.receive.is_none() { + self.receive = Some(self.queue.clone().receive(max_messages, suggested_deadline)); + } + tokio::select! { + res = self.receive.as_mut().unwrap() => { + self.receive = None; + res + } + _ = tokio::time::sleep(self.iteration) => { + Ok(Vec::new()) + } + + } + } +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use anyhow::bail; + use async_trait::async_trait; + + use super::*; + + #[derive(Clone, Debug)] + struct SleepyQueue { + receive_sleep: Duration, + } + + #[async_trait] + impl Queue for SleepyQueue { + async fn receive( + self: Arc, + _max_messages: usize, + _suggested_deadline: Duration, + ) -> anyhow::Result> { + tokio::time::sleep(self.receive_sleep).await; + bail!("Waking up from my nap") + } + + async fn acknowledge(&self, _ack_ids: &[String]) -> anyhow::Result<()> { + unimplemented!() + } + + async fn modify_deadlines( + &self, + _ack_id: &str, + _suggested_deadline: Duration, + ) -> anyhow::Result { + unimplemented!() + } + } + + #[tokio::test] + async fn test_queue_receiver_slow_receive() { + let queue = Arc::new(SleepyQueue { + receive_sleep: Duration::from_millis(100), + }); + let mut receiver = QueueReceiver::new(queue, Duration::from_millis(20)); + let mut iterations = 0; + while receiver.receive(1, Duration::from_secs(1)).await.is_ok() { + iterations += 1; + } + assert!(iterations >= 4); + } + + #[tokio::test] + async fn test_queue_receiver_fast_receive() { + let queue = Arc::new(SleepyQueue { + receive_sleep: Duration::from_millis(10), + }); + let mut receiver = QueueReceiver::new(queue, Duration::from_millis(50)); + assert!(receiver.receive(1, Duration::from_secs(1)).await.is_err()); + } +} diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/local_state.rs b/quickwit/quickwit-indexing/src/source/queue_sources/local_state.rs new file mode 100644 index 00000000000..b0361b69adf --- /dev/null +++ b/quickwit/quickwit-indexing/src/source/queue_sources/local_state.rs @@ -0,0 +1,134 @@ +// Copyright (C) 2024 Quickwit, Inc. +// +// Quickwit is offered under the AGPL v3.0 and as commercial software. +// For commercial licensing, contact us at hello@quickwit.io. +// +// AGPL: +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as +// published by the Free Software Foundation, either version 3 of the +// License, or (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::collections::{BTreeMap, BTreeSet, VecDeque}; +use std::time::Duration; + +use anyhow::bail; +use quickwit_metastore::checkpoint::PartitionId; + +use super::message::{InProgressMessage, ReadyMessage}; + +/// Tracks the state of the queue messages that are known to the owning indexing +/// pipeline. +/// +/// Messages first land in the `ready_for_read` queue. They are then moved to +/// `read_in_progress` to track the reader's progress. Once the reader reaches +/// EOF, the message is transitioned as `awaiting_commit`. Once the message is +/// known to be fully indexed and committed (e.g after receiving the +/// `suggest_truncate` call), it is moved to `completed`. +#[derive(Default)] +pub struct QueueLocalState { + /// Messages that were received from the queue and are ready to be read + ready_for_read: VecDeque, + /// Message that is currently being read and sent to the `DocProcessor` + read_in_progress: Option, + /// Partitions that were read and are still being indexed, with their + /// associated ack_id + awaiting_commit: BTreeMap, + /// Partitions that were fully indexed and committed + completed: BTreeSet, +} + +impl QueueLocalState { + pub fn is_ready_for_read(&self, partition_id: &PartitionId) -> bool { + self.ready_for_read + .iter() + .any(|msg| &msg.partition_id() == partition_id) + } + + pub fn is_read_in_progress(&self, partition_id: &PartitionId) -> bool { + self.read_in_progress + .as_ref() + .map_or(false, |msg| &msg.partition_id == partition_id) + } + + pub fn is_awaiting_commit(&self, partition_id: &PartitionId) -> bool { + self.awaiting_commit.contains_key(partition_id) + } + + pub fn is_completed(&self, partition_id: &PartitionId) -> bool { + self.completed.contains(partition_id) + } + + pub fn is_tracked(&self, partition_id: &PartitionId) -> bool { + self.is_ready_for_read(partition_id) + || self.is_read_in_progress(partition_id) + || self.is_awaiting_commit(partition_id) + || self.is_completed(partition_id) + } + + pub fn set_ready_for_read(&mut self, ready_messages: Vec) { + for message in ready_messages { + self.ready_for_read.push_back(message) + } + } + + pub fn get_ready_for_read(&mut self) -> Option { + while let Some(msg) = self.ready_for_read.pop_front() { + // don't return messages for which we didn't manage to extend the + // visibility, they will pop up in the queue again anyway + if !msg.visibility_handle.extension_failed() { + return Some(msg); + } + } + None + } + + pub fn read_in_progress_mut(&mut self) -> Option<&mut InProgressMessage> { + self.read_in_progress.as_mut() + } + + pub async fn drop_currently_read( + &mut self, + deadline_for_last_extension: Duration, + ) -> anyhow::Result<()> { + if let Some(in_progress) = self.read_in_progress.take() { + self.awaiting_commit.insert( + in_progress.partition_id.clone(), + in_progress.visibility_handle.ack_id().to_string(), + ); + in_progress + .visibility_handle + .request_last_extension(deadline_for_last_extension) + .await?; + } + Ok(()) + } + + /// Tries to set the message that is currently being read. Returns an error + /// if there is already a message being read. + pub fn set_currently_read( + &mut self, + in_progress: Option, + ) -> anyhow::Result<()> { + if self.read_in_progress.is_some() { + bail!("trying to replace in progress message"); + } + self.read_in_progress = in_progress; + Ok(()) + } + + /// Returns the ack_id if that message was awaiting_commit + pub fn mark_completed(&mut self, partition_id: PartitionId) -> Option { + let ack_id_opt = self.awaiting_commit.remove(&partition_id); + self.completed.insert(partition_id); + ack_id_opt + } +} diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/memory_queue.rs b/quickwit/quickwit-indexing/src/source/queue_sources/memory_queue.rs new file mode 100644 index 00000000000..627b7587de4 --- /dev/null +++ b/quickwit/quickwit-indexing/src/source/queue_sources/memory_queue.rs @@ -0,0 +1,233 @@ +// Copyright (C) 2024 Quickwit, Inc. +// +// Quickwit is offered under the AGPL v3.0 and as commercial software. +// For commercial licensing, contact us at hello@quickwit.io. +// +// AGPL: +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as +// published by the Free Software Foundation, either version 3 of the +// License, or (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::collections::{BTreeMap, VecDeque}; +use std::fmt; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use anyhow::bail; +use async_trait::async_trait; +use quickwit_storage::OwnedBytes; +use ulid::Ulid; + +use super::message::{MessageMetadata, RawMessage}; +use super::Queue; + +#[derive(Default)] +struct InnerState { + in_queue: VecDeque, + in_flight: BTreeMap, + acked: Vec, +} + +impl fmt::Debug for InnerState { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Queue") + .field("in_queue_count", &self.in_queue.len()) + .field("in_flight_count", &self.in_flight.len()) + .field("acked_count", &self.acked.len()) + .finish() + } +} + +/// A simple in-memory queue +#[derive(Clone, Debug)] +pub struct MemoryQueueForTests { + inner_state: Arc>, + receive_sleep: Duration, +} + +impl MemoryQueueForTests { + pub fn new() -> Self { + let inner_state = Arc::new(Mutex::new(InnerState::default())); + let inner_weak = Arc::downgrade(&inner_state); + tokio::spawn(async move { + loop { + if let Some(inner_state) = inner_weak.upgrade() { + let mut inner_state = inner_state.lock().unwrap(); + let mut expired = Vec::new(); + for (ack_id, msg) in inner_state.in_flight.iter() { + if msg.metadata.initial_deadline < Instant::now() { + expired.push(ack_id.clone()); + } + } + for ack_id in expired { + let msg = inner_state.in_flight.remove(&ack_id).unwrap(); + inner_state.in_queue.push_back(msg); + } + } else { + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + }); + MemoryQueueForTests { + inner_state: Arc::new(Mutex::new(InnerState::default())), + receive_sleep: Duration::from_millis(50), + } + } + + pub fn send_message(&self, payload: String, ack_id: &str) { + let message = RawMessage { + payload: OwnedBytes::new(payload.into_bytes()), + metadata: MessageMetadata { + ack_id: ack_id.to_string(), + delivery_attempts: 0, + initial_deadline: Instant::now(), + message_id: Ulid::new().to_string(), + }, + }; + self.inner_state.lock().unwrap().in_queue.push_back(message); + } + + /// Returns the next visibility deadline for the message if it is in flight + pub fn next_visibility_deadline(&self, ack_id: &str) -> Option { + let inner_state = self.inner_state.lock().unwrap(); + inner_state + .in_flight + .get(ack_id) + .map(|msg| msg.metadata.initial_deadline) + } +} + +#[async_trait] +impl Queue for MemoryQueueForTests { + async fn receive( + self: Arc, + max_messages: usize, + suggested_deadline: Duration, + ) -> anyhow::Result> { + { + let mut inner_state = self.inner_state.lock().unwrap(); + let mut response = Vec::new(); + while let Some(mut msg) = inner_state.in_queue.pop_front() { + msg.metadata.delivery_attempts += 1; + msg.metadata.initial_deadline = Instant::now() + suggested_deadline; + let msg_cloned = RawMessage { + payload: msg.payload.clone(), + metadata: msg.metadata.clone(), + }; + inner_state + .in_flight + .insert(msg.metadata.ack_id.clone(), msg_cloned); + response.push(msg); + if response.len() >= max_messages { + break; + } + } + if !response.is_empty() { + return Ok(response); + } + } + // `sleep` to avoid using all the CPU when called in a loop + tokio::time::sleep(self.receive_sleep).await; + + Ok(vec![]) + } + + async fn acknowledge(&self, ack_ids: &[String]) -> anyhow::Result<()> { + let mut inner_state = self.inner_state.lock().unwrap(); + for ack_id in ack_ids { + if let Some(msg) = inner_state.in_flight.remove(ack_id) { + inner_state.acked.push(msg); + } + } + Ok(()) + } + + async fn modify_deadlines( + &self, + ack_id: &str, + suggested_deadline: Duration, + ) -> anyhow::Result { + let mut inner_state = self.inner_state.lock().unwrap(); + let in_flight = inner_state.in_flight.get_mut(ack_id); + if let Some(msg) = in_flight { + msg.metadata.initial_deadline = Instant::now() + suggested_deadline; + } else { + bail!("ack_id {} not found in in-flight", ack_id); + } + return Ok(Instant::now() + suggested_deadline); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn prefilled_queue(nb_message: usize) -> Arc { + let memory_queue = MemoryQueueForTests::new(); + for i in 0..nb_message { + let payload = format!("Test message {}", i); + let ack_id = i.to_string(); + memory_queue.send_message(payload.clone(), &ack_id); + } + Arc::new(memory_queue) + } + + #[tokio::test] + async fn test_receive_1_by_1() { + let memory_queue = prefilled_queue(2); + for i in 0..2 { + let messages = memory_queue + .clone() + .receive(1, Duration::from_secs(5)) + .await + .unwrap(); + assert_eq!(messages.len(), 1); + let message = &messages[0]; + let exp_payload = format!("Test message {}", i); + let exp_ack_id = i.to_string(); + assert_eq!(message.payload.as_ref(), exp_payload.as_bytes()); + assert_eq!(message.metadata.ack_id, exp_ack_id); + } + } + + #[tokio::test] + async fn test_receive_2_by_2() { + let memory_queue = prefilled_queue(2); + let messages = memory_queue + .receive(2, Duration::from_secs(5)) + .await + .unwrap(); + assert_eq!(messages.len(), 2); + for (i, message) in messages.iter().enumerate() { + let exp_payload = format!("Test message {}", i); + let exp_ack_id = i.to_string(); + assert_eq!(message.payload.as_ref(), exp_payload.as_bytes()); + assert_eq!(message.metadata.ack_id, exp_ack_id); + } + } + + #[tokio::test] + async fn test_receive_early_if_only_1() { + let memory_queue = prefilled_queue(1); + let messages = memory_queue + .receive(2, Duration::from_secs(5)) + .await + .unwrap(); + assert_eq!(messages.len(), 1); + let message = &messages[0]; + let exp_payload = "Test message 0".to_string(); + let exp_ack_id = "0"; + assert_eq!(message.payload.as_ref(), exp_payload.as_bytes()); + assert_eq!(message.metadata.ack_id, exp_ack_id); + } +} diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/message.rs b/quickwit/quickwit-indexing/src/source/queue_sources/message.rs new file mode 100644 index 00000000000..c0ad5e2b235 --- /dev/null +++ b/quickwit/quickwit-indexing/src/source/queue_sources/message.rs @@ -0,0 +1,352 @@ +// Copyright (C) 2024 Quickwit, Inc. +// +// Quickwit is offered under the AGPL v3.0 and as commercial software. +// For commercial licensing, contact us at hello@quickwit.io. +// +// AGPL: +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as +// published by the Free Software Foundation, either version 3 of the +// License, or (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use core::fmt; +use std::io::read_to_string; +use std::str::FromStr; +use std::time::Instant; + +use anyhow::Context; +use quickwit_common::rate_limited_warn; +use quickwit_common::uri::Uri; +use quickwit_metastore::checkpoint::PartitionId; +use quickwit_proto::types::Position; +use quickwit_storage::{OwnedBytes, StorageResolver}; +use serde_json::Value; +use thiserror::Error; +use tracing::info; + +use super::visibility::VisibilityTaskHandle; +use crate::source::doc_file_reader::ObjectUriBatchReader; + +#[derive(Debug, Clone, Copy)] +pub enum MessageType { + S3Notification, + // GcsNotification, + RawUri, + // RawData, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MessageMetadata { + /// The handle that should be used to acknowledge the message or change its visibility deadline + pub ack_id: String, + + /// The unique message id assigned by the queue + pub message_id: String, + + /// The approximate number of times the message was delivered. 1 means it is + /// the first time this message is being delivered. + pub delivery_attempts: usize, + + /// The first deadline when the message is received. It can be extended later using the ack_id. + pub initial_deadline: Instant, +} + +/// The raw messages as received from the queue abstraction +pub struct RawMessage { + pub metadata: MessageMetadata, + pub payload: OwnedBytes, +} + +impl fmt::Debug for RawMessage { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("RawMessage") + .field("metadata", &self.metadata) + .field("payload", &"") + .finish() + } +} + +#[derive(Error, Debug)] +pub enum PreProcessingError { + #[error("message can be acknowledged without processing")] + Discardable { ack_id: String }, + #[error("unexpected message format: {0}")] + UnexpectedFormat(#[from] anyhow::Error), +} + +impl RawMessage { + pub fn pre_process( + self, + message_type: MessageType, + ) -> Result { + let payload = match message_type { + MessageType::S3Notification => PreProcessedPayload::ObjectUri( + uri_from_s3_notification(&self.payload, &self.metadata.ack_id)?, + ), + MessageType::RawUri => { + let payload_str = read_to_string(self.payload).context("failed to read payload")?; + PreProcessedPayload::ObjectUri(Uri::from_str(&payload_str)?) + } + }; + Ok(PreProcessedMessage { + metadata: self.metadata, + payload, + }) + } +} + +#[derive(Debug, PartialEq, Eq)] +pub enum PreProcessedPayload { + /// The message contains an object URI + ObjectUri(Uri), + // /// The message contains the raw JSON data + // RawData(OwnedBytes), +} + +impl PreProcessedPayload { + pub fn partition_id(&self) -> PartitionId { + match &self { + Self::ObjectUri(uri) => PartitionId::from(uri.as_str()), + } + } +} + +/// A message that went through the minimal transformation to discover its +/// partition id. Indeed, the message might be discarded if the partition was +/// already processed, so it's better to avoid doing unnecessary work at this +/// stage. +#[derive(Debug, PartialEq, Eq)] +pub struct PreProcessedMessage { + pub metadata: MessageMetadata, + pub payload: PreProcessedPayload, +} + +impl PreProcessedMessage { + pub fn partition_id(&self) -> PartitionId { + self.payload.partition_id() + } +} + +fn uri_from_s3_notification(message: &[u8], ack_id: &str) -> Result { + let value: Value = serde_json::from_slice(message).context("invalid JSON message")?; + if matches!(value["Event"].as_str(), Some("s3:TestEvent")) { + info!("discarding S3 test event"); + return Err(PreProcessingError::Discardable { + ack_id: ack_id.to_string(), + }); + } + let event_name = value["Records"][0]["eventName"] + .as_str() + .context("invalid S3 notification: Records[0].eventName not found")?; + if !event_name.starts_with("ObjectCreated:") { + rate_limited_warn!( + limit_per_min = 5, + event = event_name, + "only s3:ObjectCreated:* events are supported" + ); + return Err(PreProcessingError::Discardable { + ack_id: ack_id.to_string(), + }); + } + let key = value["Records"][0]["s3"]["object"]["key"] + .as_str() + .context("invalid S3 notification: Records[0].s3.object.key not found")?; + let bucket = value["Records"][0]["s3"]["bucket"]["name"] + .as_str() + .context("invalid S3 notification: Records[0].s3.bucket.name not found")?; + Uri::from_str(&format!("s3://{}/{}", bucket, key)).map_err(|e| e.into()) +} + +/// A message for which we know as much of the global processing status as +/// possible and that is now ready to be processed. +pub struct ReadyMessage { + pub position: Position, + pub content: PreProcessedMessage, + pub visibility_handle: VisibilityTaskHandle, +} + +impl ReadyMessage { + pub async fn start_processing( + self, + storage_resolver: &StorageResolver, + ) -> anyhow::Result> { + let partition_id = self.partition_id(); + match self.content.payload { + PreProcessedPayload::ObjectUri(uri) => { + let batch_reader = ObjectUriBatchReader::try_new( + storage_resolver, + partition_id.clone(), + &uri, + self.position, + ) + .await?; + if batch_reader.is_eof() { + Ok(None) + } else { + Ok(Some(InProgressMessage { + batch_reader, + partition_id, + visibility_handle: self.visibility_handle, + })) + } + } + } + } + + pub fn partition_id(&self) -> PartitionId { + self.content.partition_id() + } +} + +/// A message that is actively being read +pub struct InProgressMessage { + pub partition_id: PartitionId, + pub visibility_handle: VisibilityTaskHandle, + pub batch_reader: ObjectUriBatchReader, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_uri_from_s3_notification_valid() { + let test_message = r#" + { + "Records": [ + { + "eventVersion": "2.1", + "eventSource": "aws:s3", + "awsRegion": "us-west-2", + "eventTime": "2021-05-22T09:22:41.789Z", + "eventName": "ObjectCreated:Put", + "userIdentity": { + "principalId": "AWS:AIDAJDPLRKLG7UEXAMPLE" + }, + "requestParameters": { + "sourceIPAddress": "127.0.0.1" + }, + "responseElements": { + "x-amz-request-id": "C3D13FE58DE4C810", + "x-amz-id-2": "FMyUVURIx7Zv2cPi/IZb9Fk1/U4QfTaVK5fahHPj/" + }, + "s3": { + "s3SchemaVersion": "1.0", + "configurationId": "testConfigRule", + "bucket": { + "name": "mybucket", + "ownerIdentity": { + "principalId": "A3NL1KOZZKExample" + }, + "arn": "arn:aws:s3:::mybucket" + }, + "object": { + "key": "logs.json", + "size": 1024, + "eTag": "d41d8cd98f00b204e9800998ecf8427e", + "versionId": "096fKKXTRTtl3on89fVO.nfljtsv6qko", + "sequencer": "0055AED6DCD90281E5" + } + } + } + ] + }"#; + let actual_uri = uri_from_s3_notification(test_message.as_bytes(), "myackid").unwrap(); + let expected_uri = Uri::from_str("s3://mybucket/logs.json").unwrap(); + assert_eq!(actual_uri, expected_uri); + } + + #[test] + fn test_uri_from_s3_notification_invalid() { + let invalid_message = r#"{ + "Records": [ + { + "s3": { + "object": { + "key": "test_key" + } + } + } + ] + }"#; + let result = + uri_from_s3_notification(&OwnedBytes::new(invalid_message.as_bytes()), "myackid"); + assert!(matches!( + result, + Err(PreProcessingError::UnexpectedFormat(_)) + )); + } + + #[test] + fn test_uri_from_s3_bad_event_type() { + let invalid_message = r#"{ + "Records": [ + { + "eventVersion": "2.1", + "eventSource": "aws:s3", + "awsRegion": "us-east-1", + "eventTime": "2024-07-29T12:47:14.577Z", + "eventName": "ObjectRemoved:Delete", + "userIdentity": { + "principalId": "AWS:ARGHGOHSDGOKGHOGHMCC4:user" + }, + "requestParameters": { + "sourceIPAddress": "1.1.1.1" + }, + "responseElements": { + "x-amz-request-id": "GHGSH", + "x-amz-id-2": "gndflghndflhmnrflsh+gLLKU6X0PvD6ANdVY1+/hspflhjladgfkelagfkndl" + }, + "s3": { + "s3SchemaVersion": "1.0", + "configurationId": "hello", + "bucket": { + "name": "mybucket", + "ownerIdentity": { + "principalId": "KMGP12GHKKH" + }, + "arn": "arn:aws:s3:::mybucket" + }, + "object": { + "key": "my_deleted_file", + "sequencer": "GKHOFLGKHSALFK0" + } + } + } + ] + }"#; + let result = + uri_from_s3_notification(&OwnedBytes::new(invalid_message.as_bytes()), "myackid"); + assert!(matches!( + result, + Err(PreProcessingError::Discardable { .. }) + )); + } + + #[test] + fn test_uri_from_s3_notification_discardable() { + let invalid_message = r#"{ + "Service":"Amazon S3", + "Event":"s3:TestEvent", + "Time":"2014-10-13T15:57:02.089Z", + "Bucket":"bucketname", + "RequestId":"5582815E1AEA5ADF", + "HostId":"8cLeGAmw098X5cv4Zkwcmo8vvZa3eH3eKxsPzbB9wrR+YstdA6Knx4Ip8EXAMPLE" + }"#; + let result = + uri_from_s3_notification(&OwnedBytes::new(invalid_message.as_bytes()), "myackid"); + if let Err(PreProcessingError::Discardable { ack_id }) = result { + assert_eq!(ack_id, "myackid"); + } else { + panic!("Expected skippable error"); + } + } +} diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/mod.rs b/quickwit/quickwit-indexing/src/source/queue_sources/mod.rs new file mode 100644 index 00000000000..dc94a53d7f8 --- /dev/null +++ b/quickwit/quickwit-indexing/src/source/queue_sources/mod.rs @@ -0,0 +1,89 @@ +// Copyright (C) 2024 Quickwit, Inc. +// +// Quickwit is offered under the AGPL v3.0 and as commercial software. +// For commercial licensing, contact us at hello@quickwit.io. +// +// AGPL: +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as +// published by the Free Software Foundation, either version 3 of the +// License, or (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +pub mod coordinator; +mod helpers; +mod local_state; +#[cfg(test)] +mod memory_queue; +mod message; +mod shared_state; +#[cfg(feature = "sqs")] +pub mod sqs_queue; +mod visibility; + +use std::fmt; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use async_trait::async_trait; +use message::RawMessage; + +/// The queue abstraction is based on the AWS SQS and Google Pubsub APIs. The +/// only requirement of the underlying implementation is that messages exposed +/// to a given consumer are hidden to other consumers for a configurable period +/// of time. Retries are handled by the implementation because queues might +/// behave differently (throttling, deduplication...). +#[async_trait] +pub trait Queue: fmt::Debug + Send + Sync + 'static { + /// Polls the queue to receive messages. + /// + /// The implementation is in charge of choosing the wait strategy when there + /// are no messages in the queue. It will typically use long polling to do + /// this efficiently. On the other hand, when there is a message available + /// in the queue, it should be returned as quickly as possible, regardless + /// of the `max_messages` parameter. The `max_messages` paramater should + /// always be clamped by the implementation to not violate the maximum value + /// supported by the backing queue (e.g 10 messages for AWS SQS). + /// + /// As soon as the message is received, the caller is responsible for + /// maintaining the message visibility in a timely fashion. Failing to do so + /// implies that duplicates will be received by other indexing pipelines, + /// thus increasing competition for the commit lock. + async fn receive( + // `Arc` to make the resulting future `'static` and thus easily + // wrappable by the `QueueReceiver` + self: Arc, + max_messages: usize, + suggested_deadline: Duration, + ) -> anyhow::Result>; + + /// Tries to acknowledge (delete) the messages. + /// + /// The call returns `Ok(())` if at the message level: + /// - the acknowledgement failed due to a transient failure + /// - the message was already acknowledged + /// - the message was not acknowledged in time and is back to the queue + /// + /// If an empty list of ack_ids is provided, the call should be a no-op. + async fn acknowledge(&self, ack_ids: &[String]) -> anyhow::Result<()>; + + /// Modifies the visibility deadline of the messages. + /// + /// We try to set the initial visibility large enough to avoid having to + /// call this too often. The implementation can retry as long as desired, + /// it's the caller's responsibility to cancel the future if the deadline is + /// getting to close to the expiration. The returned `Instant` is a + /// conservative estimate of the new deadline expiration time. + async fn modify_deadlines( + &self, + ack_id: &str, + suggested_deadline: Duration, + ) -> anyhow::Result; +} diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/shared_state.rs b/quickwit/quickwit-indexing/src/source/queue_sources/shared_state.rs new file mode 100644 index 00000000000..cdd0ade05b3 --- /dev/null +++ b/quickwit/quickwit-indexing/src/source/queue_sources/shared_state.rs @@ -0,0 +1,371 @@ +// Copyright (C) 2024 Quickwit, Inc. +// +// Quickwit is offered under the AGPL v3.0 and as commercial software. +// For commercial licensing, contact us at hello@quickwit.io. +// +// AGPL: +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as +// published by the Free Software Foundation, either version 3 of the +// License, or (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::collections::BTreeMap; + +use anyhow::{bail, Context}; +use quickwit_metastore::checkpoint::PartitionId; +use quickwit_proto::metastore::{ + AcquireShardsRequest, MetastoreService, MetastoreServiceClient, OpenShardSubrequest, + OpenShardsRequest, +}; +use quickwit_proto::types::{DocMappingUid, IndexUid, Position, ShardId}; +use tracing::info; + +use super::message::PreProcessedMessage; + +#[derive(Clone)] +pub struct QueueSharedState { + pub metastore: MetastoreServiceClient, + pub index_uid: IndexUid, + pub source_id: String, +} + +impl QueueSharedState { + /// Tries to acquire the ownership for the provided messages from the global + /// shared context. For each partition id, if the ownership was successfully + /// acquired or the partition was already successfully indexed, the position + /// is returned along with the partition id, otherwise the partition id is + /// dropped. + async fn acquire_partitions( + &self, + publish_token: &str, + partitions: Vec, + ) -> anyhow::Result> { + let open_shard_subrequests = partitions + .iter() + .enumerate() + .map(|(idx, partition_id)| OpenShardSubrequest { + subrequest_id: idx as u32, + index_uid: Some(self.index_uid.clone()), + source_id: self.source_id.clone(), + leader_id: String::new(), + follower_id: None, + shard_id: Some(ShardId::from(partition_id.as_str())), + doc_mapping_uid: Some(DocMappingUid::default()), + publish_token: Some(publish_token.to_string()), + }) + .collect(); + + let open_shard_resp = self + .metastore + .open_shards(OpenShardsRequest { + subrequests: open_shard_subrequests, + }) + .await?; + + let mut shards = Vec::new(); + let mut re_acquired_shards = Vec::new(); + for sub in open_shard_resp.subresponses { + // we could also just cast the shard_id back to a partition_id + let partition_id = partitions[sub.subrequest_id as usize].clone(); + let shard = sub.open_shard(); + let position = shard.publish_position_inclusive.clone().unwrap_or_default(); + let is_owned = sub.open_shard().publish_token.as_deref() == Some(publish_token); + if position.is_eof() || (is_owned && position.is_beginning()) { + shards.push((partition_id, position)); + } else if !is_owned { + // TODO: Add logic to only re-acquire shards that have a token that is not + // the local token when they haven't been updated recently + info!(previous_token = shard.publish_token, "shard re-acquired"); + re_acquired_shards.push(shard.shard_id().clone()); + } else if is_owned && !position.is_beginning() { + bail!("Partition is owned by this indexing pipeline but is not at the beginning. This should never happen! Please, report on https://github.com/quickwit-oss/quickwit/issues.") + } + } + + if re_acquired_shards.is_empty() { + return Ok(shards); + } + + // re-acquire shards that have a token that is not the local token + let acquire_shard_resp = self + .metastore + .acquire_shards(AcquireShardsRequest { + index_uid: Some(self.index_uid.clone()), + source_id: self.source_id.clone(), + shard_ids: re_acquired_shards, + publish_token: publish_token.to_string(), + }) + .await + .unwrap(); + for shard in acquire_shard_resp.acquired_shards { + let partition_id = PartitionId::from(shard.shard_id().as_str()); + let position = shard.publish_position_inclusive.unwrap_or_default(); + shards.push((partition_id, position)); + } + + Ok(shards) + } +} + +/// Acquires shards from the shared state for the provided list of messages and +/// maps results to that same list +pub async fn checkpoint_messages( + shared_state: &QueueSharedState, + publish_token: &str, + messages: Vec, +) -> anyhow::Result> { + let mut message_map = + BTreeMap::from_iter(messages.into_iter().map(|msg| (msg.partition_id(), msg))); + let partition_ids = message_map.keys().cloned().collect(); + + let shards = shared_state + .acquire_partitions(publish_token, partition_ids) + .await?; + + shards + .into_iter() + .map(|(partition_id, position)| { + let content = message_map.remove(&partition_id).context("Unexpected partition ID. This should never happen! Please, report on https://github.com/quickwit-oss/quickwit/issues.")?; + Ok(( + content, + position, + )) + }) + .collect::>() +} + +#[cfg(test)] +pub mod shared_state_for_tests { + use std::sync::{Arc, Mutex}; + + use quickwit_proto::ingest::{Shard, ShardState}; + use quickwit_proto::metastore::{ + AcquireShardsResponse, MockMetastoreService, OpenShardSubresponse, OpenShardsResponse, + }; + + use super::*; + + pub(super) fn mock_metastore( + initial_state: &[(PartitionId, (String, Position))], + open_shard_times: Option, + acquire_times: Option, + ) -> MetastoreServiceClient { + let mut mock_metastore = MockMetastoreService::new(); + let inner_state = Arc::new(Mutex::new(BTreeMap::from_iter( + initial_state.iter().cloned(), + ))); + let inner_state_ref = Arc::clone(&inner_state); + let open_shards_expectation = + mock_metastore + .expect_open_shards() + .returning(move |request| { + let subresponses = request + .subrequests + .into_iter() + .map(|sub_req| { + let partition_id: PartitionId = sub_req.shard_id().to_string().into(); + let (token, position) = inner_state_ref + .lock() + .unwrap() + .get(&partition_id) + .cloned() + .unwrap_or((sub_req.publish_token.unwrap(), Position::Beginning)); + inner_state_ref + .lock() + .unwrap() + .insert(partition_id, (token.clone(), position.clone())); + OpenShardSubresponse { + subrequest_id: sub_req.subrequest_id, + open_shard: Some(Shard { + shard_id: sub_req.shard_id, + source_id: sub_req.source_id, + publish_token: Some(token), + index_uid: sub_req.index_uid, + follower_id: sub_req.follower_id, + leader_id: sub_req.leader_id, + doc_mapping_uid: sub_req.doc_mapping_uid, + publish_position_inclusive: Some(position), + shard_state: ShardState::Open as i32, + }), + } + }) + .collect(); + Ok(OpenShardsResponse { subresponses }) + }); + if let Some(times) = open_shard_times { + open_shards_expectation.times(times); + } + let acquire_shards_expectation = mock_metastore + .expect_acquire_shards() + // .times(acquire_times) + .returning(move |request| { + let acquired_shards = request + .shard_ids + .into_iter() + .map(|shard_id| { + let partition_id: PartitionId = shard_id.to_string().into(); + let (existing_token, position) = inner_state + .lock() + .unwrap() + .get(&partition_id) + .cloned() + .expect("we should never try to acquire a shard that doesn't exist"); + inner_state.lock().unwrap().insert( + partition_id, + (request.publish_token.clone(), position.clone()), + ); + assert_ne!(existing_token, request.publish_token); + Shard { + shard_id: Some(shard_id), + source_id: "dummy".to_string(), + publish_token: Some(request.publish_token.clone()), + index_uid: None, + follower_id: None, + leader_id: "dummy".to_string(), + doc_mapping_uid: None, + publish_position_inclusive: Some(position), + shard_state: ShardState::Open as i32, + } + }) + .collect(); + Ok(AcquireShardsResponse { acquired_shards }) + }); + if let Some(times) = acquire_times { + acquire_shards_expectation.times(times); + } + MetastoreServiceClient::from_mock(mock_metastore) + } + + pub fn shared_state_for_tests( + index_id: &str, + initial_state: &[(PartitionId, (String, Position))], + ) -> QueueSharedState { + let index_uid = IndexUid::new_with_random_ulid(index_id); + let metastore = mock_metastore(initial_state, None, None); + QueueSharedState { + metastore, + index_uid, + source_id: "test-queue-src".to_string(), + } + } +} + +#[cfg(test)] +mod tests { + use std::str::FromStr; + use std::time::Instant; + use std::vec; + + use quickwit_common::uri::Uri; + use shared_state_for_tests::mock_metastore; + + use super::*; + use crate::source::queue_sources::message::{MessageMetadata, PreProcessedPayload}; + + fn test_messages(message_number: usize) -> Vec { + (0..message_number) + .map(|i| PreProcessedMessage { + metadata: MessageMetadata { + ack_id: format!("ackid{}", i), + delivery_attempts: 0, + initial_deadline: Instant::now(), + message_id: format!("mid{}", i), + }, + payload: PreProcessedPayload::ObjectUri( + Uri::from_str(&format!("s3://bucket/key{}", i)).unwrap(), + ), + }) + .collect() + } + + #[tokio::test] + async fn test_acquire_shards_with_completed() { + let index_id = "test-sqs-index"; + let index_uid = IndexUid::new_with_random_ulid(index_id); + let init_state = &[("p1".into(), ("token2".to_string(), Position::eof(100usize)))]; + let metastore = mock_metastore(init_state, Some(1), Some(0)); + + let shared_state = QueueSharedState { + metastore, + index_uid, + source_id: "test-sqs-source".to_string(), + }; + + let aquired = shared_state + .acquire_partitions("token1", vec!["p1".into(), "p2".into()]) + .await + .unwrap(); + assert!(aquired.contains(&("p1".into(), Position::eof(100usize)))); + assert!(aquired.contains(&("p2".into(), Position::Beginning))); + } + + #[tokio::test] + async fn test_re_acquire_shards() { + let index_id = "test-sqs-index"; + let index_uid = IndexUid::new_with_random_ulid(index_id); + let init_state = &[( + "p1".into(), + ("token2".to_string(), Position::offset(100usize)), + )]; + let metastore = mock_metastore(init_state, Some(1), Some(1)); + + let shared_state = QueueSharedState { + metastore, + index_uid, + source_id: "test-sqs-source".to_string(), + }; + + let aquired = shared_state + .acquire_partitions("token1", vec!["p1".into(), "p2".into()]) + .await + .unwrap(); + // TODO: this test should fail once we implement the grace + // period before a partition can be re-acquired + assert!(aquired.contains(&("p1".into(), Position::offset(100usize)))); + assert!(aquired.contains(&("p2".into(), Position::Beginning))); + } + + #[tokio::test] + async fn test_checkpoint_with_completed() { + let index_id = "test-sqs-index"; + let index_uid = IndexUid::new_with_random_ulid(index_id); + + let source_messages = test_messages(2); + let completed_partition_id = source_messages[0].partition_id(); + let new_partition_id = source_messages[1].partition_id(); + + let init_state = &[( + completed_partition_id.clone(), + ("token2".to_string(), Position::eof(100usize)), + )]; + let metastore = mock_metastore(init_state, Some(1), Some(0)); + let shared_state = QueueSharedState { + metastore, + index_uid, + source_id: "test-sqs-source".to_string(), + }; + + let checkpointed_msg = checkpoint_messages(&shared_state, "token1", source_messages) + .await + .unwrap(); + assert_eq!(checkpointed_msg.len(), 2); + let completed_msg = checkpointed_msg + .iter() + .find(|(msg, _)| msg.partition_id() == completed_partition_id) + .unwrap(); + assert_eq!(completed_msg.1, Position::eof(100usize)); + let new_msg = checkpointed_msg + .iter() + .find(|(msg, _)| msg.partition_id() == new_partition_id) + .unwrap(); + assert_eq!(new_msg.1, Position::Beginning); + } +} diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/sqs_queue.rs b/quickwit/quickwit-indexing/src/source/queue_sources/sqs_queue.rs new file mode 100644 index 00000000000..49d5923a9e1 --- /dev/null +++ b/quickwit/quickwit-indexing/src/source/queue_sources/sqs_queue.rs @@ -0,0 +1,468 @@ +// Copyright (C) 2024 Quickwit, Inc. +// +// Quickwit is offered under the AGPL v3.0 and as commercial software. +// For commercial licensing, contact us at hello@quickwit.io. +// +// AGPL: +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as +// published by the Free Software Foundation, either version 3 of the +// License, or (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use anyhow::{bail, Context}; +use async_trait::async_trait; +use aws_sdk_sqs::config::{BehaviorVersion, Builder, Region, SharedAsyncSleep}; +use aws_sdk_sqs::types::{DeleteMessageBatchRequestEntry, MessageSystemAttributeName}; +use aws_sdk_sqs::{Client, Config}; +use itertools::Itertools; +use quickwit_aws::retry::{aws_retry, AwsRetryable}; +use quickwit_aws::{get_aws_config, DEFAULT_AWS_REGION}; +use quickwit_common::rate_limited_error; +use quickwit_common::retry::RetryParams; +use quickwit_storage::OwnedBytes; +use regex::Regex; + +use super::message::MessageMetadata; +use super::{Queue, RawMessage}; + +#[derive(Debug)] +pub struct SqsQueue { + sqs_client: Client, + queue_url: String, + receive_retries: RetryParams, + acknowledge_retries: RetryParams, + modify_deadline_retries: RetryParams, +} + +impl SqsQueue { + pub async fn try_new(queue_url: String) -> anyhow::Result { + let sqs_client = get_sqs_client(&queue_url).await?; + Ok(SqsQueue { + sqs_client, + queue_url, + receive_retries: RetryParams::standard(), + // Acknowledgment is retried when the message is received again + acknowledge_retries: RetryParams::no_retries(), + // Retry aggressively to avoid loosing the ownership of the message + modify_deadline_retries: RetryParams::aggressive(), + }) + } +} + +#[async_trait] +impl Queue for SqsQueue { + async fn receive( + self: Arc, + max_messages: usize, + suggested_deadline: Duration, + ) -> anyhow::Result> { + // TODO: We estimate the message deadline using the start of the + // ReceiveMessage request. This might be overly pessimistic: the docs + // state that it starts when the message is returned. + let initial_deadline = Instant::now() + suggested_deadline; + let clamped_max_messages = std::cmp::min(max_messages, 10) as i32; + let receive_output = aws_retry(&self.receive_retries, || async { + self.sqs_client + .receive_message() + .queue_url(&self.queue_url) + .message_system_attribute_names(MessageSystemAttributeName::ApproximateReceiveCount) + .wait_time_seconds(20) + .set_max_number_of_messages(Some(clamped_max_messages)) + .visibility_timeout(suggested_deadline.as_secs() as i32) + .send() + .await + }) + .await?; + + let received_messages = receive_output.messages.unwrap_or_default(); + let mut resulting_raw_messages = Vec::with_capacity(received_messages.len()); + for received_message in received_messages { + let delivery_attempts: usize = received_message + .attributes + .as_ref() + .and_then(|attrs| attrs.get(&MessageSystemAttributeName::ApproximateReceiveCount)) + .and_then(|s| s.parse().ok()) + .unwrap_or(0); + let ack_id = received_message + .receipt_handle + .context("missing receipt_handle in received message")?; + let message_id = received_message + .message_id + .context("missing message_id in received message")?; + let raw_message = RawMessage { + metadata: MessageMetadata { + ack_id, + message_id, + initial_deadline, + delivery_attempts, + }, + payload: OwnedBytes::new(received_message.body.unwrap_or_default().into_bytes()), + }; + resulting_raw_messages.push(raw_message); + } + Ok(resulting_raw_messages) + } + + async fn acknowledge(&self, ack_ids: &[String]) -> anyhow::Result<()> { + if ack_ids.is_empty() { + return Ok(()); + } + let entry_batches: Vec> = ack_ids + .iter() + .dedup() + .enumerate() + .map(|(i, id)| { + DeleteMessageBatchRequestEntry::builder() + .id(i.to_string()) + .receipt_handle(id.to_string()) + .build() + .unwrap() + }) + .chunks(10) + .into_iter() + .map(|chunk| chunk.collect()) + .collect(); + + // TODO: parallelization + let mut batch_errors = Vec::new(); + let mut message_errors = Vec::new(); + for batch in entry_batches { + let res = aws_retry(&self.acknowledge_retries, || { + self.sqs_client + .delete_message_batch() + .queue_url(&self.queue_url) + .set_entries(Some(batch.clone())) + .send() + }) + .await; + match res { + Ok(res) => { + message_errors.extend(res.failed.into_iter()); + } + Err(err) => { + batch_errors.push(err); + } + } + } + if batch_errors.iter().any(|err| !err.is_retryable()) { + let fatal_error = batch_errors + .into_iter() + .find(|err| !err.is_retryable()) + .unwrap(); + bail!(fatal_error); + } else if !batch_errors.is_empty() { + rate_limited_error!( + limit_per_min = 10, + count = batch_errors.len(), + first_err = ?batch_errors.into_iter().next().unwrap(), + "failed to acknowledge some message batches", + ); + } + // The documentation is unclear about these partial failures. We assume + // it is either: + // - a transient failure + // - the message is already acknowledged + // - the message is expired + if !message_errors.is_empty() { + rate_limited_error!( + limit_per_min = 10, + count = message_errors.len(), + first_err = ?message_errors.into_iter().next().unwrap(), + "failed to acknowledge individual messages", + ); + } + Ok(()) + } + + async fn modify_deadlines( + &self, + ack_id: &str, + suggested_deadline: Duration, + ) -> anyhow::Result { + let visibility_timeout = std::cmp::min(suggested_deadline.as_secs() as i32, 43200); + let new_deadline = Instant::now() + suggested_deadline; + aws_retry(&self.modify_deadline_retries, || { + self.sqs_client + .change_message_visibility() + .queue_url(&self.queue_url) + .visibility_timeout(visibility_timeout) + .receipt_handle(ack_id) + .send() + }) + .await?; + Ok(new_deadline) + } +} + +async fn preconfigured_builder() -> anyhow::Result { + let aws_config = get_aws_config().await; + + let mut sqs_config = Config::builder().behavior_version(BehaviorVersion::v2024_03_28()); + sqs_config.set_retry_config(aws_config.retry_config().cloned()); + sqs_config.set_credentials_provider(aws_config.credentials_provider()); + sqs_config.set_http_client(aws_config.http_client()); + sqs_config.set_timeout_config(aws_config.timeout_config().cloned()); + + if let Some(identity_cache) = aws_config.identity_cache() { + sqs_config.set_identity_cache(identity_cache); + } + sqs_config.set_sleep_impl(Some(SharedAsyncSleep::new( + quickwit_aws::TokioSleep::default(), + ))); + + Ok(sqs_config) +} + +fn queue_url_region(queue_url: &str) -> Option { + let re = Regex::new(r"^https?://sqs\.(.*?)\.amazonaws\.com").unwrap(); + let caps = re.captures(queue_url)?; + let region_str = caps.get(1)?.as_str(); + Some(Region::new(region_str.to_string())) +} + +fn queue_url_endpoint(queue_url: &str) -> anyhow::Result { + let re = Regex::new(r"(^https?://[^/]+)").unwrap(); + let caps = re.captures(queue_url).context("Invalid queue URL")?; + let endpoint_str = caps.get(1).context("Invalid queue URL")?.as_str(); + Ok(endpoint_str.to_string()) +} + +pub async fn get_sqs_client(queue_url: &str) -> anyhow::Result { + let mut sqs_config = preconfigured_builder().await?; + // region is required by the SDK to work + let inferred_region = queue_url_region(queue_url).unwrap_or(DEFAULT_AWS_REGION); + let inferred_endpoint = queue_url_endpoint(queue_url)?; + sqs_config.set_region(Some(inferred_region)); + sqs_config.set_endpoint_url(Some(inferred_endpoint)); + Ok(Client::from_conf(sqs_config.build())) +} + +/// Checks whether we can establish a connection to the SQS service and we can +/// access the provided queue_url +pub(crate) async fn check_connectivity(queue_url: &str) -> anyhow::Result<()> { + let client = get_sqs_client(queue_url).await?; + client + .get_queue_attributes() + .queue_url(queue_url) + .send() + .await?; + + Ok(()) +} + +#[cfg(feature = "sqs-localstack-tests")] +pub mod test_helpers { + use aws_sdk_sqs::types::QueueAttributeName; + use ulid::Ulid; + + use super::*; + + pub async fn get_localstack_sqs_client() -> anyhow::Result { + let mut sqs_config = preconfigured_builder().await?; + sqs_config.set_endpoint_url(Some("http://localhost:4566".to_string())); + sqs_config.set_region(Some(DEFAULT_AWS_REGION)); + Ok(Client::from_conf(sqs_config.build())) + } + + pub async fn create_queue(sqs_client: &Client, queue_name_prefix: &str) -> String { + let queue_name = format!("{}-{}", queue_name_prefix, Ulid::new()); + sqs_client + .create_queue() + .queue_name(queue_name) + .send() + .await + .unwrap() + .queue_url + .unwrap() + } + + pub async fn send_message(sqs_client: &Client, queue_url: &str, payload: &str) { + sqs_client + .send_message() + .queue_url(queue_url) + .message_body(payload.to_string()) + .send() + .await + .unwrap(); + } + + pub async fn get_queue_attribute( + sqs_client: &Client, + queue_url: &str, + attribute: QueueAttributeName, + ) -> String { + let queue_attributes = sqs_client + .get_queue_attributes() + .queue_url(queue_url) + .attribute_names(attribute.clone()) + .send() + .await + .unwrap(); + queue_attributes + .attributes + .unwrap() + .get(&attribute) + .unwrap() + .to_string() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_queue_url_region() { + let url = "https://sqs.eu-west-2.amazonaws.com/12345678910/test"; + let region = queue_url_region(url); + assert_eq!(region, Some(Region::from_static("eu-west-2"))); + + let url = "https://sqs.ap-south-1.amazonaws.com/12345678910/test"; + let region = queue_url_region(url); + assert_eq!(region, Some(Region::from_static("ap-south-1"))); + + let url = "http://localhost:4566/000000000000/test-queue"; + let region = queue_url_region(url); + assert_eq!(region, None); + } + + #[test] + fn test_queue_url_endpoint() { + let url = "https://sqs.eu-west-2.amazonaws.com/12345678910/test"; + let endpoint = queue_url_endpoint(url).unwrap(); + assert_eq!(endpoint, "https://sqs.eu-west-2.amazonaws.com"); + + let url = "https://sqs.ap-south-1.amazonaws.com/12345678910/test"; + let endpoint = queue_url_endpoint(url).unwrap(); + assert_eq!(endpoint, "https://sqs.ap-south-1.amazonaws.com"); + + let url = "http://localhost:4566/000000000000/test-queue"; + let endpoint = queue_url_endpoint(url).unwrap(); + assert_eq!(endpoint, "http://localhost:4566"); + + let url = "http://localhost:4566/000000000000/test-queue"; + let endpoint = queue_url_endpoint(url).unwrap(); + assert_eq!(endpoint, "http://localhost:4566"); + } +} + +#[cfg(all(test, feature = "sqs-localstack-tests"))] +mod localstack_tests { + use aws_sdk_sqs::types::QueueAttributeName; + + use super::*; + use crate::source::queue_sources::helpers::QueueReceiver; + use crate::source::queue_sources::sqs_queue::test_helpers::{ + create_queue, get_localstack_sqs_client, + }; + + #[tokio::test] + async fn test_check_connectivity() { + let sqs_client = get_localstack_sqs_client().await.unwrap(); + let queue_url = create_queue(&sqs_client, "check-connectivity").await; + check_connectivity(&queue_url).await.unwrap(); + } + + #[tokio::test] + async fn test_receive_existing_msg_quickly() { + let client = test_helpers::get_localstack_sqs_client().await.unwrap(); + let queue_url = test_helpers::create_queue(&client, "test-receive-existing-msg").await; + let message = "hello world"; + test_helpers::send_message(&client, &queue_url, message).await; + + let queue = Arc::new(SqsQueue::try_new(queue_url).await.unwrap()); + let messages = tokio::time::timeout( + Duration::from_millis(500), + queue.clone().receive(5, Duration::from_secs(60)), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].payload.as_slice(), message.as_bytes()); + + // just assess that there are no errors for now + queue + .modify_deadlines(&messages[0].metadata.ack_id, Duration::from_secs(10)) + .await + .unwrap(); + queue + .acknowledge(&[messages[0].metadata.ack_id.clone()]) + .await + .unwrap(); + } + + #[tokio::test] + async fn test_acknowledge_larger_batch() { + let client = test_helpers::get_localstack_sqs_client().await.unwrap(); + let queue_url = test_helpers::create_queue(&client, "test-ack-large").await; + let message = "hello world"; + for _ in 0..20 { + test_helpers::send_message(&client, &queue_url, message).await; + } + + let queue: Arc = Arc::new(SqsQueue::try_new(queue_url.clone()).await.unwrap()); + let mut queue_receiver = QueueReceiver::new(queue.clone(), Duration::from_millis(200)); + let mut messages = Vec::new(); + for _ in 0..5 { + let new_messages = queue_receiver + .receive(20, Duration::from_secs(60)) + .await + .unwrap(); + messages.extend(new_messages.into_iter()); + } + assert_eq!(messages.len(), 20); + let in_flight_count: usize = test_helpers::get_queue_attribute( + &client, + &queue_url, + QueueAttributeName::ApproximateNumberOfMessagesNotVisible, + ) + .await + .parse() + .unwrap(); + assert_eq!(in_flight_count, 20); + + let ack_ids = messages + .iter() + .map(|msg| msg.metadata.ack_id.clone()) + .collect::>(); + + queue.acknowledge(&ack_ids).await.unwrap(); + + let in_flight_count: usize = test_helpers::get_queue_attribute( + &client, + &queue_url, + QueueAttributeName::ApproximateNumberOfMessagesNotVisible, + ) + .await + .parse() + .unwrap(); + assert_eq!(in_flight_count, 0); + } + + #[tokio::test] + async fn test_receive_wrong_queue() { + let client = test_helpers::get_localstack_sqs_client().await.unwrap(); + let queue_url = test_helpers::create_queue(&client, "test-receive-existing-msg").await; + let bad_queue_url = format!("{}wrong", queue_url); + let queue = Arc::new(SqsQueue::try_new(bad_queue_url).await.unwrap()); + tokio::time::timeout( + Duration::from_millis(500), + queue.clone().receive(5, Duration::from_secs(60)), + ) + .await + .unwrap() + .unwrap_err(); + } +} diff --git a/quickwit/quickwit-indexing/src/source/queue_sources/visibility.rs b/quickwit/quickwit-indexing/src/source/queue_sources/visibility.rs new file mode 100644 index 00000000000..340a6c05b95 --- /dev/null +++ b/quickwit/quickwit-indexing/src/source/queue_sources/visibility.rs @@ -0,0 +1,341 @@ +// Copyright (C) 2024 Quickwit, Inc. +// +// Quickwit is offered under the AGPL v3.0 and as commercial software. +// For commercial licensing, contact us at hello@quickwit.io. +// +// AGPL: +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as +// published by the Free Software Foundation, either version 3 of the +// License, or (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::sync::{Arc, Weak}; +use std::time::{Duration, Instant}; + +use anyhow::{anyhow, Context}; +use async_trait::async_trait; +use quickwit_actors::{ + Actor, ActorContext, ActorExitStatus, ActorHandle, ActorState, Handler, Mailbox, +}; +use serde_json::{json, Value as JsonValue}; + +use super::Queue; +use crate::source::SourceContext; + +#[derive(Debug, Clone)] +pub(super) struct VisibilitySettings { + /// The original deadline asked from the queue when polling the messages + pub deadline_for_receive: Duration, + /// The last deadline extension when the message reading is completed + pub deadline_for_last_extension: Duration, + /// The extension applied why the VisibilityTask to maintain the message visibility + pub deadline_for_default_extension: Duration, + /// Rhe timeout for the visibility extension request + pub request_timeout: Duration, + /// an extra margin that is substracted from the expected deadline when + /// asserting whether we are still in time to extend the visibility + pub request_margin: Duration, +} + +impl VisibilitySettings { + /// The commit timeout gives us a first estimate on how long the processing + /// will take for the messages. We could include other factors such as the + /// message size. + pub(super) fn from_commit_timeout(commit_timeout_secs: usize) -> Self { + let commit_timeout = Duration::from_secs(commit_timeout_secs as u64); + Self { + deadline_for_receive: Duration::from_secs(120) + commit_timeout, + deadline_for_last_extension: 2 * commit_timeout, + deadline_for_default_extension: Duration::from_secs(60), + request_timeout: Duration::from_secs(3), + request_margin: Duration::from_secs(1), + } + } +} + +#[derive(Debug)] +struct VisibilityTask { + queue: Arc, + ack_id: String, + extension_count: u64, + current_deadline: Instant, + last_extension_requested: bool, + visibility_settings: VisibilitySettings, + ref_count: Weak<()>, +} + +// A handle to the visibility actor. When dropped, the actor exits and the +// visibility isn't maintained anymore. +pub(super) struct VisibilityTaskHandle { + mailbox: Mailbox, + actor_handle: ActorHandle, + ack_id: String, + _ref_count: Arc<()>, +} + +/// Spawns actor that ensures that the visibility of a given message +/// (represented by its ack_id) is extended when required. We prefer applying +/// ample margins in the extension process to avoid missing deadlines while also +/// keeping the number of extension requests (and associated cost) small. +pub(super) fn spawn_visibility_task( + ctx: &SourceContext, + queue: Arc, + ack_id: String, + current_deadline: Instant, + visibility_settings: VisibilitySettings, +) -> VisibilityTaskHandle { + let ref_count = Arc::new(()); + let weak_ref = Arc::downgrade(&ref_count); + let task = VisibilityTask { + queue, + ack_id: ack_id.clone(), + extension_count: 0, + current_deadline, + last_extension_requested: false, + visibility_settings, + ref_count: weak_ref, + }; + let (mailbox, actor_handle) = ctx.spawn_actor().spawn(task); + VisibilityTaskHandle { + mailbox, + actor_handle, + ack_id, + _ref_count: ref_count, + } +} + +impl VisibilityTask { + async fn extend_visibility( + &mut self, + ctx: &ActorContext, + extension: Duration, + ) -> anyhow::Result<()> { + let _zone = ctx.protect_zone(); + self.current_deadline = tokio::time::timeout( + self.visibility_settings.request_timeout, + self.queue.modify_deadlines(&self.ack_id, extension), + ) + .await + .context("deadline extension timed out")??; + self.extension_count += 1; + Ok(()) + } + + fn next_extension(&self) -> Duration { + (self.current_deadline - Instant::now()) + - self.visibility_settings.request_timeout + - self.visibility_settings.request_margin + } +} + +impl VisibilityTaskHandle { + pub fn extension_failed(&self) -> bool { + self.actor_handle.state() == ActorState::Failure + } + + pub fn ack_id(&self) -> &str { + &self.ack_id + } + + pub async fn request_last_extension(self, extension: Duration) -> anyhow::Result<()> { + self.mailbox + .ask_for_res(RequestLastExtension { extension }) + .await + .map_err(|e| anyhow!(e))?; + Ok(()) + } +} + +#[async_trait] +impl Actor for VisibilityTask { + type ObservableState = JsonValue; + + fn name(&self) -> String { + "QueueVisibilityTask".to_string() + } + + async fn initialize(&mut self, ctx: &ActorContext) -> Result<(), ActorExitStatus> { + let first_extension = self.next_extension(); + if first_extension.is_zero() { + return Err(anyhow!("initial visibility deadline insufficient").into()); + } + ctx.schedule_self_msg(first_extension, Loop); + Ok(()) + } + + fn yield_after_each_message(&self) -> bool { + false + } + + fn observable_state(&self) -> Self::ObservableState { + json!({ + "ack_id": self.ack_id, + "extension_count": self.extension_count, + }) + } +} + +#[derive(Debug)] +struct Loop; + +#[async_trait] +impl Handler for VisibilityTask { + type Reply = (); + + async fn handle( + &mut self, + _message: Loop, + ctx: &ActorContext, + ) -> Result<(), ActorExitStatus> { + if self.ref_count.strong_count() == 0 { + return Ok(()); + } + if self.last_extension_requested { + return Ok(()); + } + self.extend_visibility(ctx, self.visibility_settings.deadline_for_default_extension) + .await?; + ctx.schedule_self_msg(self.next_extension(), Loop); + Ok(()) + } +} + +/// Ensures that the visibility of the message is extended until the given +/// deadline and then stops the extension loop. +#[derive(Debug)] +struct RequestLastExtension { + extension: Duration, +} + +#[async_trait] +impl Handler for VisibilityTask { + type Reply = anyhow::Result<()>; + + async fn handle( + &mut self, + message: RequestLastExtension, + ctx: &ActorContext, + ) -> Result { + let last_deadline = Instant::now() + message.extension; + self.last_extension_requested = true; + if last_deadline > self.current_deadline { + Ok(self.extend_visibility(ctx, message.extension).await) + } else { + Ok(Ok(())) + } + } +} + +#[cfg(test)] +mod tests { + use quickwit_actors::Universe; + use tokio::sync::watch; + + use super::*; + use crate::source::queue_sources::memory_queue::MemoryQueueForTests; + + #[tokio::test] + async fn test_visibility_task_request_last_extension() { + // actor context + let universe = Universe::with_accelerated_time(); + let (source_mailbox, _source_inbox) = universe.create_test_mailbox(); + let (observable_state_tx, _observable_state_rx) = watch::channel(serde_json::Value::Null); + let ctx: SourceContext = + ActorContext::for_test(&universe, source_mailbox, observable_state_tx); + // queue with test message + let ack_id = "ack_id".to_string(); + let queue = Arc::new(MemoryQueueForTests::new()); + queue.send_message("test message".to_string(), &ack_id); + let initial_deadline = queue + .clone() + .receive(1, Duration::from_secs(1)) + .await + .unwrap()[0] + .metadata + .initial_deadline; + // spawn task + let visibility_settings = VisibilitySettings { + deadline_for_default_extension: Duration::from_secs(1), + deadline_for_last_extension: Duration::from_secs(20), + deadline_for_receive: Duration::from_secs(1), + request_timeout: Duration::from_millis(100), + request_margin: Duration::from_millis(100), + }; + let handle = spawn_visibility_task( + &ctx, + queue.clone(), + ack_id.clone(), + initial_deadline, + visibility_settings.clone(), + ); + // assert that the background task performs extensions + assert!(!handle.extension_failed()); + tokio::time::sleep_until(initial_deadline.into()).await; + let next_deadline = queue.next_visibility_deadline(&ack_id).unwrap(); + assert!(initial_deadline < next_deadline); + assert!(!handle.extension_failed()); + // request last extension + handle + .request_last_extension(Duration::from_secs(5)) + .await + .unwrap(); + assert!( + Instant::now() + Duration::from_secs(4) + < queue.next_visibility_deadline(&ack_id).unwrap() + ); + universe.assert_quit().await; + } + + #[tokio::test] + async fn test_visibility_task_stop_on_drop() { + // actor context + let universe = Universe::with_accelerated_time(); + let (source_mailbox, _source_inbox) = universe.create_test_mailbox(); + let (observable_state_tx, _observable_state_rx) = watch::channel(serde_json::Value::Null); + let ctx: SourceContext = + ActorContext::for_test(&universe, source_mailbox, observable_state_tx); + // queue with test message + let ack_id = "ack_id".to_string(); + let queue = Arc::new(MemoryQueueForTests::new()); + queue.send_message("test message".to_string(), &ack_id); + let initial_deadline = queue + .clone() + .receive(1, Duration::from_secs(1)) + .await + .unwrap()[0] + .metadata + .initial_deadline; + // spawn task + let visibility_settings = VisibilitySettings { + deadline_for_default_extension: Duration::from_secs(1), + deadline_for_last_extension: Duration::from_secs(20), + deadline_for_receive: Duration::from_secs(1), + request_timeout: Duration::from_millis(100), + request_margin: Duration::from_millis(100), + }; + let handle = spawn_visibility_task( + &ctx, + queue.clone(), + ack_id.clone(), + initial_deadline, + visibility_settings.clone(), + ); + // assert that visibility is not extended after drop + drop(handle); + tokio::time::sleep_until(initial_deadline.into()).await; + // the message is either already expired or about to expire + if let Some(next_deadline) = queue.next_visibility_deadline(&ack_id) { + assert_eq!(next_deadline, initial_deadline); + } + // assert_eq!(q, None); + universe.assert_quit().await; + } +} diff --git a/quickwit/quickwit-indexing/src/source/stdin_source.rs b/quickwit/quickwit-indexing/src/source/stdin_source.rs new file mode 100644 index 00000000000..81e9f40b3e6 --- /dev/null +++ b/quickwit/quickwit-indexing/src/source/stdin_source.rs @@ -0,0 +1,133 @@ +// Copyright (C) 2024 Quickwit, Inc. +// +// Quickwit is offered under the AGPL v3.0 and as commercial software. +// For commercial licensing, contact us at hello@quickwit.io. +// +// AGPL: +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as +// published by the Free Software Foundation, either version 3 of the +// License, or (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::fmt; +use std::time::Duration; + +use async_trait::async_trait; +use quickwit_actors::{ActorExitStatus, Mailbox}; +use quickwit_common::Progress; +use quickwit_proto::metastore::SourceType; +use tokio::io::{AsyncBufReadExt, BufReader}; + +use super::{BatchBuilder, BATCH_NUM_BYTES_LIMIT}; +use crate::actors::DocProcessor; +use crate::source::{Source, SourceContext, SourceRuntime, TypedSourceFactory}; + +pub struct StdinBatchReader { + reader: BufReader, + is_eof: bool, +} + +impl StdinBatchReader { + pub fn new() -> Self { + Self { + reader: BufReader::new(tokio::io::stdin()), + is_eof: false, + } + } + + async fn read_batch(&mut self, source_progress: &Progress) -> anyhow::Result { + let mut batch_builder = BatchBuilder::new(SourceType::Stdin); + while batch_builder.num_bytes < BATCH_NUM_BYTES_LIMIT { + let mut buf = String::new(); + // stdin might be slow because it's depending on external + // input (e.g. user typing on a keyboard) + let bytes_read = source_progress + .protect_future(self.reader.read_line(&mut buf)) + .await?; + if bytes_read > 0 { + batch_builder.add_doc(buf.into()); + } else { + self.is_eof = true; + break; + } + } + + Ok(batch_builder) + } + + fn is_eof(&self) -> bool { + self.is_eof + } +} + +pub struct StdinSource { + reader: StdinBatchReader, + num_bytes_processed: u64, + num_lines_processed: u64, +} + +impl fmt::Debug for StdinSource { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "StdinSource") + } +} + +#[async_trait] +impl Source for StdinSource { + async fn emit_batches( + &mut self, + doc_processor_mailbox: &Mailbox, + ctx: &SourceContext, + ) -> Result { + let batch_builder = self.reader.read_batch(ctx.progress()).await?; + self.num_bytes_processed += batch_builder.num_bytes; + self.num_lines_processed += batch_builder.docs.len() as u64; + doc_processor_mailbox + .send_message(batch_builder.build()) + .await?; + if self.reader.is_eof() { + ctx.send_exit_with_success(doc_processor_mailbox).await?; + return Err(ActorExitStatus::Success); + } + + Ok(Duration::ZERO) + } + + fn name(&self) -> String { + format!("{:?}", self) + } + + fn observable_state(&self) -> serde_json::Value { + serde_json::json!({ + "num_bytes_processed": self.num_bytes_processed, + "num_lines_processed": self.num_lines_processed, + }) + } +} + +pub struct FileSourceFactory; + +#[async_trait] +impl TypedSourceFactory for FileSourceFactory { + type Source = StdinSource; + type Params = (); + + async fn typed_create_source( + _source_runtime: SourceRuntime, + _params: (), + ) -> anyhow::Result { + Ok(StdinSource { + reader: StdinBatchReader::new(), + num_bytes_processed: 0, + num_lines_processed: 0, + }) + } +} diff --git a/quickwit/quickwit-integration-tests/Cargo.toml b/quickwit/quickwit-integration-tests/Cargo.toml index 7030d2ba004..aa6a692c236 100644 --- a/quickwit/quickwit-integration-tests/Cargo.toml +++ b/quickwit/quickwit-integration-tests/Cargo.toml @@ -10,10 +10,17 @@ repository.workspace = true authors.workspace = true license.workspace = true +[features] +sqs-localstack-tests = [ + "quickwit-indexing/sqs", + "quickwit-indexing/sqs-localstack-tests" +] + [dependencies] [dev-dependencies] anyhow = { workspace = true } +aws-sdk-sqs = { workspace = true } futures-util = { workspace = true } hyper = { workspace = true } itertools = { workspace = true } @@ -23,11 +30,13 @@ tempfile = { workspace = true } tokio = { workspace = true } tonic = { workspace = true } tracing = { workspace = true } +tracing-subscriber = { workspace = true } quickwit-actors = { workspace = true, features = ["testsuite"] } quickwit-cli = { workspace = true } quickwit-common = { workspace = true, features = ["testsuite"] } quickwit-config = { workspace = true, features = ["testsuite"] } +quickwit-indexing = { workspace = true, features = ["testsuite"] } quickwit-metastore = { workspace = true, features = ["testsuite"] } quickwit-proto = { workspace = true, features = ["testsuite"] } quickwit-rest-client = { workspace = true } diff --git a/quickwit/quickwit-integration-tests/src/test_utils/cluster_sandbox.rs b/quickwit/quickwit-integration-tests/src/test_utils/cluster_sandbox.rs index 824e24988e1..17eed1116fb 100644 --- a/quickwit/quickwit-integration-tests/src/test_utils/cluster_sandbox.rs +++ b/quickwit/quickwit-integration-tests/src/test_utils/cluster_sandbox.rs @@ -24,6 +24,7 @@ use std::path::PathBuf; use std::str::FromStr; use std::time::{Duration, Instant}; +use anyhow::Context; use futures_util::future; use itertools::Itertools; use quickwit_actors::ActorExitStatus; @@ -419,7 +420,12 @@ impl ClusterSandbox { input_format: quickwit_config::SourceInputFormat::Json, overwrite: false, vrl_script: None, - input_path_opt: Some(tmp_data_file.path().to_path_buf()), + input_path_opt: Some(QuickwitUri::from_str( + tmp_data_file + .path() + .to_str() + .context("temp path could not be converted to URI")?, + )?), }) .await?; Ok(()) diff --git a/quickwit/quickwit-integration-tests/src/tests/mod.rs b/quickwit/quickwit-integration-tests/src/tests/mod.rs index 9c50db45bbc..103fd9a9f1f 100644 --- a/quickwit/quickwit-integration-tests/src/tests/mod.rs +++ b/quickwit/quickwit-integration-tests/src/tests/mod.rs @@ -19,4 +19,6 @@ mod basic_tests; mod index_tests; +#[cfg(feature = "sqs-localstack-tests")] +mod sqs_tests; mod update_tests; diff --git a/quickwit/quickwit-integration-tests/src/tests/sqs_tests.rs b/quickwit/quickwit-integration-tests/src/tests/sqs_tests.rs new file mode 100644 index 00000000000..d2f86b00ff6 --- /dev/null +++ b/quickwit/quickwit-integration-tests/src/tests/sqs_tests.rs @@ -0,0 +1,167 @@ +// Copyright (C) 2024 Quickwit, Inc. +// +// Quickwit is offered under the AGPL v3.0 and as commercial software. +// For commercial licensing, contact us at hello@quickwit.io. +// +// AGPL: +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as +// published by the Free Software Foundation, either version 3 of the +// License, or (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::io::Write; +use std::iter; +use std::str::FromStr; +use std::time::Duration; + +use aws_sdk_sqs::types::QueueAttributeName; +use quickwit_common::test_utils::wait_until_predicate; +use quickwit_common::uri::Uri; +use quickwit_config::ConfigFormat; +use quickwit_indexing::source::sqs_queue::test_helpers as sqs_test_helpers; +use quickwit_metastore::SplitState; +use quickwit_serve::SearchRequestQueryString; +use tempfile::NamedTempFile; +use tracing::info; + +use crate::test_utils::ClusterSandbox; + +fn create_mock_data_file(num_lines: usize) -> (NamedTempFile, Uri) { + let mut temp_file = tempfile::NamedTempFile::new().unwrap(); + for i in 0..num_lines { + writeln!(temp_file, "{{\"body\": \"hello {}\"}}", i).unwrap() + } + temp_file.flush().unwrap(); + let path = temp_file.path().to_str().unwrap(); + let uri = Uri::from_str(path).unwrap(); + (temp_file, uri) +} + +#[tokio::test] +async fn test_sqs_single_node_cluster() { + tracing_subscriber::fmt::init(); + let sandbox = ClusterSandbox::start_standalone_node().await.unwrap(); + let index_id = "test-sqs-source-single-node-cluster"; + let index_config = format!( + r#" + version: 0.8 + index_id: {} + doc_mapping: + field_mappings: + - name: body + type: text + indexing_settings: + commit_timeout_secs: 1 + "#, + index_id + ); + + info!("create SQS queue"); + let sqs_client = sqs_test_helpers::get_localstack_sqs_client().await.unwrap(); + let queue_url = sqs_test_helpers::create_queue(&sqs_client, "test-single-node-cluster").await; + + sandbox.wait_for_cluster_num_ready_nodes(1).await.unwrap(); + + info!("create index"); + sandbox + .indexer_rest_client + .indexes() + .create(index_config.clone(), ConfigFormat::Yaml, false) + .await + .unwrap(); + + let source_id: &str = "test-sqs-single-node-cluster"; + let source_config_input = format!( + r#" + version: 0.7 + source_id: {} + desired_num_pipelines: 1 + max_num_pipelines_per_indexer: 1 + source_type: file + params: + notifications: + - type: sqs + queue_url: {} + message_type: raw_uri + input_format: plain_text + "#, + source_id, queue_url + ); + + info!("create file source with SQS notification"); + sandbox + .indexer_rest_client + .sources(index_id) + .create(source_config_input, ConfigFormat::Yaml) + .await + .unwrap(); + + // Send messages with duplicates + let tmp_mock_data_files: Vec<_> = iter::repeat_with(|| create_mock_data_file(1000)) + .take(10) + .collect(); + for (_, uri) in &tmp_mock_data_files { + sqs_test_helpers::send_message(&sqs_client, &queue_url, uri.as_str()).await; + } + sqs_test_helpers::send_message(&sqs_client, &queue_url, tmp_mock_data_files[0].1.as_str()) + .await; + sqs_test_helpers::send_message(&sqs_client, &queue_url, tmp_mock_data_files[5].1.as_str()) + .await; + + info!("wait for split to be published"); + sandbox + .wait_for_splits(index_id, Some(vec![SplitState::Published]), 1) + .await + .unwrap(); + + info!("count docs using search"); + let search_result = sandbox + .indexer_rest_client + .search( + index_id, + SearchRequestQueryString { + query: "".to_string(), + max_hits: 0, + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(search_result.num_hits, 10 * 1000); + + wait_until_predicate( + || async { + let in_flight_count: usize = sqs_test_helpers::get_queue_attribute( + &sqs_client, + &queue_url, + QueueAttributeName::ApproximateNumberOfMessagesNotVisible, + ) + .await + .parse() + .unwrap(); + in_flight_count == 2 + }, + Duration::from_secs(5), + Duration::from_millis(100), + ) + .await + .expect("Number of in-flight messages didn't reach 2 within the timeout"); + + info!("delete index"); + sandbox + .indexer_rest_client + .indexes() + .delete(index_id, false) + .await + .unwrap(); + + sandbox.shutdown().await.unwrap(); +} diff --git a/quickwit/quickwit-lambda/src/indexer/handler.rs b/quickwit/quickwit-lambda/src/indexer/handler.rs index 1282b0e54a4..c2027b25955 100644 --- a/quickwit/quickwit-lambda/src/indexer/handler.rs +++ b/quickwit/quickwit-lambda/src/indexer/handler.rs @@ -34,7 +34,7 @@ async fn indexer_handler(event: LambdaEvent) -> Result { let payload = serde_json::from_value::(event.payload)?; let ingest_res = ingest(IngestArgs { - input_path: payload.uri(), + input_path: payload.uri()?, input_format: quickwit_config::SourceInputFormat::Json, overwrite: false, vrl_script: None, diff --git a/quickwit/quickwit-lambda/src/indexer/ingest/helpers.rs b/quickwit/quickwit-lambda/src/indexer/ingest/helpers.rs index 176dc8bbc9a..282eb6622c7 100644 --- a/quickwit/quickwit-lambda/src/indexer/ingest/helpers.rs +++ b/quickwit/quickwit-lambda/src/indexer/ingest/helpers.rs @@ -19,7 +19,7 @@ use std::collections::HashSet; use std::num::NonZeroUsize; -use std::path::{Path, PathBuf}; +use std::path::Path; use anyhow::{bail, Context}; use chitchat::transport::ChannelTransport; @@ -138,12 +138,12 @@ pub(super) async fn send_telemetry() { /// Convert the incoming file path to a source config pub(super) async fn configure_source( - input_path: PathBuf, + input_uri: Uri, input_format: SourceInputFormat, vrl_script: Option, ) -> anyhow::Result { let transform_config = vrl_script.map(|vrl_script| TransformConfig::new(vrl_script, None)); - let source_params = SourceParams::file(input_path.clone()); + let source_params = SourceParams::file_from_uri(input_uri); Ok(SourceConfig { source_id: LAMBDA_SOURCE_ID.to_owned(), num_pipelines: NonZeroUsize::new(1).expect("1 is always non-zero."), diff --git a/quickwit/quickwit-lambda/src/indexer/ingest/mod.rs b/quickwit/quickwit-lambda/src/indexer/ingest/mod.rs index 6faf495b85f..13dd7f9d1b5 100644 --- a/quickwit/quickwit-lambda/src/indexer/ingest/mod.rs +++ b/quickwit/quickwit-lambda/src/indexer/ingest/mod.rs @@ -20,7 +20,6 @@ mod helpers; use std::collections::HashSet; -use std::path::PathBuf; use anyhow::bail; use helpers::{ @@ -31,6 +30,7 @@ use quickwit_actors::Universe; use quickwit_cli::start_actor_runtimes; use quickwit_cli::tool::start_statistics_reporting_loop; use quickwit_common::runtimes::RuntimesConfig; +use quickwit_common::uri::Uri; use quickwit_config::service::QuickwitService; use quickwit_config::SourceInputFormat; use quickwit_index_management::clear_cache_directory; @@ -43,7 +43,7 @@ use crate::utils::load_node_config; #[derive(Debug, Eq, PartialEq)] pub struct IngestArgs { - pub input_path: PathBuf, + pub input_path: Uri, pub input_format: SourceInputFormat, pub overwrite: bool, pub vrl_script: Option, diff --git a/quickwit/quickwit-lambda/src/indexer/model.rs b/quickwit/quickwit-lambda/src/indexer/model.rs index fe6ae14aea4..2cf785ca178 100644 --- a/quickwit/quickwit-lambda/src/indexer/model.rs +++ b/quickwit/quickwit-lambda/src/indexer/model.rs @@ -17,9 +17,10 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::path::PathBuf; +use std::str::FromStr; use aws_lambda_events::event::s3::S3Event; +use quickwit_common::uri::Uri; use serde::{Deserialize, Serialize}; #[derive(Clone, Debug, Deserialize, Serialize)] @@ -31,17 +32,17 @@ pub enum IndexerEvent { } impl IndexerEvent { - pub fn uri(&self) -> PathBuf { - match &self { - IndexerEvent::Custom { source_uri } => PathBuf::from(source_uri), + pub fn uri(&self) -> anyhow::Result { + let path: String = match self { + IndexerEvent::Custom { source_uri } => source_uri.clone(), IndexerEvent::S3(event) => [ "s3://", event.records[0].s3.bucket.name.as_ref().unwrap(), event.records[0].s3.object.key.as_ref().unwrap(), ] - .iter() - .collect(), - } + .join(""), + }; + Uri::from_str(&path) } } @@ -58,14 +59,14 @@ mod tests { }); let parsed_cust_event: IndexerEvent = serde_json::from_value(cust_event).unwrap(); assert_eq!( - parsed_cust_event.uri(), - PathBuf::from("s3://quickwit-test/test.json"), + parsed_cust_event.uri().unwrap(), + Uri::from_str("s3://quickwit-test/test.json").unwrap(), ); } #[test] fn test_s3_event_uri() { - let cust_event = json!({ + let s3_event = json!({ "Records": [ { "eventVersion": "2.0", @@ -103,10 +104,10 @@ mod tests { } ] }); - let parsed_cust_event: IndexerEvent = serde_json::from_value(cust_event).unwrap(); + let s3_event: IndexerEvent = serde_json::from_value(s3_event).unwrap(); assert_eq!( - parsed_cust_event.uri(), - PathBuf::from("s3://quickwit-test/test.json"), + s3_event.uri().unwrap(), + Uri::from_str("s3://quickwit-test/test.json").unwrap(), ); } } diff --git a/quickwit/quickwit-metastore/src/checkpoint.rs b/quickwit/quickwit-metastore/src/checkpoint.rs index 5627af53a0e..bd8f2b7bd97 100644 --- a/quickwit/quickwit-metastore/src/checkpoint.rs +++ b/quickwit/quickwit-metastore/src/checkpoint.rs @@ -33,7 +33,7 @@ use thiserror::Error; use tracing::{debug, warn}; /// A `PartitionId` uniquely identifies a partition for a given source. -#[derive(Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)] +#[derive(Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, Hash)] pub struct PartitionId(pub Arc); impl PartitionId { diff --git a/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/mod.rs b/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/mod.rs index 120c0574831..d7c75e9a6f3 100644 --- a/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/mod.rs +++ b/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/mod.rs @@ -32,7 +32,6 @@ use itertools::Itertools; use quickwit_common::pretty::PrettySample; use quickwit_config::{ DocMapping, IndexingSettings, RetentionPolicy, SearchSettings, SourceConfig, - INGEST_V2_SOURCE_ID, }; use quickwit_proto::metastore::{ AcquireShardsRequest, AcquireShardsResponse, DeleteQuery, DeleteShardsRequest, @@ -48,6 +47,7 @@ use tracing::{info, warn}; use super::MutationOccurred; use crate::checkpoint::IndexCheckpointDelta; +use crate::metastore::use_shard_api; use crate::{split_tag_filter, IndexMetadata, ListSplitsQuery, Split, SplitMetadata, SplitState}; /// A `FileBackedIndex` object carries an index metadata and its split metadata. @@ -82,6 +82,7 @@ pub(crate) struct FileBackedIndex { #[cfg(any(test, feature = "testsuite"))] impl quickwit_config::TestableForRegression for FileBackedIndex { fn sample_for_regression() -> Self { + use quickwit_config::INGEST_V2_SOURCE_ID; use quickwit_proto::ingest::{Shard, ShardState}; use quickwit_proto::types::{DocMappingUid, Position, ShardId}; @@ -381,8 +382,14 @@ impl FileBackedIndex { ) -> MetastoreResult<()> { if let Some(checkpoint_delta) = checkpoint_delta_opt { let source_id = checkpoint_delta.source_id.clone(); + let source = self.metadata.sources.get(&source_id).ok_or_else(|| { + MetastoreError::NotFound(EntityKind::Source { + index_id: self.index_id().to_string(), + source_id: source_id.clone(), + }) + })?; - if source_id == INGEST_V2_SOURCE_ID { + if use_shard_api(&source.source_params) { let publish_token = publish_token_opt.ok_or_else(|| { let message = format!( "publish token is required for publishing splits for source `{source_id}`" diff --git a/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/shards.rs b/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/shards.rs index 236a4dcd9c7..c9246b1aacf 100644 --- a/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/shards.rs +++ b/quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/shards.rs @@ -131,7 +131,7 @@ impl Shards { follower_id: subrequest.follower_id, doc_mapping_uid: subrequest.doc_mapping_uid, publish_position_inclusive: Some(Position::Beginning), - publish_token: None, + publish_token: subrequest.publish_token.clone(), }; mutation_occurred = true; entry.insert(shard.clone()); @@ -335,6 +335,7 @@ mod tests { leader_id: "leader_id".to_string(), follower_id: None, doc_mapping_uid: Some(DocMappingUid::default()), + publish_token: None, }; let MutationOccurred::Yes(subresponse) = shards.open_shard(subrequest.clone()).unwrap() else { @@ -349,6 +350,7 @@ mod tests { assert_eq!(shard.shard_state(), ShardState::Open); assert_eq!(shard.leader_id, "leader_id"); assert_eq!(shard.follower_id, None); + assert_eq!(shard.publish_token, None); assert_eq!(shard.publish_position_inclusive(), Position::Beginning); let MutationOccurred::No(subresponse) = shards.open_shard(subrequest).unwrap() else { @@ -367,6 +369,7 @@ mod tests { leader_id: "leader_id".to_string(), follower_id: Some("follower_id".to_string()), doc_mapping_uid: Some(DocMappingUid::default()), + publish_token: Some("publish_token".to_string()), }; let MutationOccurred::Yes(subresponse) = shards.open_shard(subrequest).unwrap() else { panic!("Expected `MutationOccurred::No`"); diff --git a/quickwit/quickwit-metastore/src/metastore/mod.rs b/quickwit/quickwit-metastore/src/metastore/mod.rs index cbddbbd4c1f..89822b549cc 100644 --- a/quickwit/quickwit-metastore/src/metastore/mod.rs +++ b/quickwit/quickwit-metastore/src/metastore/mod.rs @@ -33,7 +33,8 @@ pub use index_metadata::IndexMetadata; use itertools::Itertools; use quickwit_common::thread_pool::run_cpu_intensive; use quickwit_config::{ - DocMapping, IndexConfig, IndexingSettings, RetentionPolicy, SearchSettings, SourceConfig, + DocMapping, FileSourceParams, IndexConfig, IndexingSettings, RetentionPolicy, SearchSettings, + SourceConfig, SourceParams, }; use quickwit_doc_mapper::tag_pruning::TagFilterAst; use quickwit_proto::metastore::{ @@ -903,6 +904,25 @@ impl Default for FilterRange { } } +/// Maps the given source params to whether checkpoints should be stored in the index metadata +/// (false) or the shard table (true) +fn use_shard_api(params: &SourceParams) -> bool { + match params { + SourceParams::File(FileSourceParams::Filepath(_)) => false, + SourceParams::File(FileSourceParams::Notifications(_)) => true, + SourceParams::Ingest => true, + SourceParams::IngestApi => false, + SourceParams::IngestCli => false, + SourceParams::Kafka(_) => false, + SourceParams::Kinesis(_) => false, + SourceParams::PubSub(_) => false, + SourceParams::Pulsar(_) => false, + SourceParams::Stdin => panic!("stdin cannot be checkpointed"), + SourceParams::Vec(_) => false, + SourceParams::Void(_) => false, + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/quickwit/quickwit-metastore/src/metastore/postgres/metastore.rs b/quickwit/quickwit-metastore/src/metastore/postgres/metastore.rs index 0289b10867e..133f1413b1e 100644 --- a/quickwit/quickwit-metastore/src/metastore/postgres/metastore.rs +++ b/quickwit/quickwit-metastore/src/metastore/postgres/metastore.rs @@ -28,7 +28,6 @@ use quickwit_common::uri::Uri; use quickwit_common::ServiceStream; use quickwit_config::{ validate_index_id_pattern, IndexTemplate, IndexTemplateId, PostgresMetastoreConfig, - INGEST_V2_SOURCE_ID, }; use quickwit_proto::ingest::{Shard, ShardState}; use quickwit_proto::metastore::{ @@ -68,7 +67,7 @@ use crate::file_backed::MutationOccurred; use crate::metastore::postgres::model::Shards; use crate::metastore::postgres::utils::split_maturity_timestamp; use crate::metastore::{ - IndexesMetadataResponseExt, PublishSplitsRequestExt, STREAM_SPLITS_CHUNK_SIZE, + use_shard_api, IndexesMetadataResponseExt, PublishSplitsRequestExt, STREAM_SPLITS_CHUNK_SIZE, }; use crate::{ AddSourceRequestExt, CreateIndexRequestExt, IndexMetadata, IndexMetadataResponseExt, @@ -684,8 +683,14 @@ impl MetastoreService for PostgresqlMetastore { } if let Some(checkpoint_delta) = checkpoint_delta_opt { let source_id = checkpoint_delta.source_id.clone(); + let source = index_metadata.sources.get(&source_id).ok_or_else(|| { + MetastoreError::NotFound(EntityKind::Source { + index_id: index_uid.index_id.to_string(), + source_id: source_id.to_string(), + }) + })?; - if source_id == INGEST_V2_SOURCE_ID { + if use_shard_api(&source.source_params) { let publish_token = request.publish_token_opt.ok_or_else(|| { let message = format!( "publish token is required for publishing splits for source \ @@ -1622,6 +1627,7 @@ async fn open_or_fetch_shard<'e>( .bind(&subrequest.leader_id) .bind(&subrequest.follower_id) .bind(subrequest.doc_mapping_uid) + .bind(&subrequest.publish_token) .fetch_optional(executor.clone()) .await?; diff --git a/quickwit/quickwit-metastore/src/metastore/postgres/queries/shards/open.sql b/quickwit/quickwit-metastore/src/metastore/postgres/queries/shards/open.sql index d6747fcb9f1..bd9e9240688 100644 --- a/quickwit/quickwit-metastore/src/metastore/postgres/queries/shards/open.sql +++ b/quickwit/quickwit-metastore/src/metastore/postgres/queries/shards/open.sql @@ -1,5 +1,5 @@ -INSERT INTO shards(index_uid, source_id, shard_id, leader_id, follower_id, doc_mapping_uid) - VALUES ($1, $2, $3, $4, $5, $6) +INSERT INTO shards(index_uid, source_id, shard_id, leader_id, follower_id, doc_mapping_uid, publish_token) + VALUES ($1, $2, $3, $4, $5, $6, $7) ON CONFLICT DO NOTHING RETURNING diff --git a/quickwit/quickwit-metastore/src/tests/shard.rs b/quickwit/quickwit-metastore/src/tests/shard.rs index 1bbbd1e010a..ed3327774ec 100644 --- a/quickwit/quickwit-metastore/src/tests/shard.rs +++ b/quickwit/quickwit-metastore/src/tests/shard.rs @@ -134,6 +134,7 @@ pub async fn test_metastore_open_shards< leader_id: "test-ingester-foo".to_string(), follower_id: Some("test-ingester-bar".to_string()), doc_mapping_uid: Some(DocMappingUid::default()), + publish_token: None, }], }; let open_shards_response = metastore.open_shards(open_shards_request).await.unwrap(); @@ -163,6 +164,7 @@ pub async fn test_metastore_open_shards< leader_id: "test-ingester-foo".to_string(), follower_id: Some("test-ingester-bar".to_string()), doc_mapping_uid: Some(DocMappingUid::default()), + publish_token: Some("publish-token-baz".to_string()), }], }; let open_shards_response = metastore.open_shards(open_shards_request).await.unwrap(); @@ -181,6 +183,35 @@ pub async fn test_metastore_open_shards< assert_eq!(shard.publish_position_inclusive(), Position::Beginning); assert!(shard.publish_token.is_none()); + // Test open shard #2. + let open_shards_request = OpenShardsRequest { + subrequests: vec![OpenShardSubrequest { + subrequest_id: 0, + index_uid: Some(test_index.index_uid.clone()), + source_id: test_index.source_id.clone(), + shard_id: Some(ShardId::from(2)), + leader_id: "test-ingester-foo".to_string(), + follower_id: None, + doc_mapping_uid: Some(DocMappingUid::default()), + publish_token: Some("publish-token-open".to_string()), + }], + }; + let open_shards_response = metastore.open_shards(open_shards_request).await.unwrap(); + assert_eq!(open_shards_response.subresponses.len(), 1); + + let subresponse = &open_shards_response.subresponses[0]; + assert_eq!(subresponse.subrequest_id, 0); + + let shard = subresponse.open_shard(); + assert_eq!(shard.index_uid(), &test_index.index_uid); + assert_eq!(shard.source_id, test_index.source_id); + assert_eq!(shard.shard_id(), ShardId::from(2)); + assert_eq!(shard.shard_state(), ShardState::Open); + assert_eq!(shard.leader_id, "test-ingester-foo"); + assert!(shard.follower_id.is_none()); + assert_eq!(shard.publish_position_inclusive(), Position::Beginning); + assert_eq!(shard.publish_token(), "publish-token-open"); + cleanup_index(&mut metastore, test_index.index_uid).await; } diff --git a/quickwit/quickwit-metastore/src/tests/split.rs b/quickwit/quickwit-metastore/src/tests/split.rs index eb8865731f6..ed9d21a36ed 100644 --- a/quickwit/quickwit-metastore/src/tests/split.rs +++ b/quickwit/quickwit-metastore/src/tests/split.rs @@ -21,7 +21,7 @@ use std::time::Duration; use futures::future::try_join_all; use quickwit_common::rand::append_random_suffix; -use quickwit_config::IndexConfig; +use quickwit_config::{IndexConfig, SourceConfig, SourceParams}; use quickwit_proto::metastore::{ CreateIndexRequest, DeleteSplitsRequest, EntityKind, IndexMetadataRequest, ListSplitsRequest, ListStaleSplitsRequest, MarkSplitsForDeletionRequest, MetastoreError, PublishSplitsRequest, @@ -77,8 +77,10 @@ pub async fn test_metastore_publish_splits_empty_splits_array_is_allowed< // checkpoint. This operation is allowed and used in the Indexer. { let index_config = IndexConfig::for_test(&index_id, &index_uri); + let source_configs = &[SourceConfig::for_test(&source_id, SourceParams::void())]; let create_index_request = - CreateIndexRequest::try_from_index_config(&index_config).unwrap(); + CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs) + .unwrap(); let index_uid: IndexUid = metastore .create_index(create_index_request) .await @@ -133,6 +135,7 @@ pub async fn test_metastore_publish_splits< let index_config = IndexConfig::for_test(&index_id, &index_uri); let source_id = format!("{index_id}--source"); + let source_configs = &[SourceConfig::for_test(&source_id, SourceParams::void())]; let split_id_1 = format!("{index_id}--split-1"); let split_metadata_1 = SplitMetadata { @@ -199,7 +202,8 @@ pub async fn test_metastore_publish_splits< // Publish a non-existent split on an index { let create_index_request = - CreateIndexRequest::try_from_index_config(&index_config).unwrap(); + CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs) + .unwrap(); let index_uid: IndexUid = metastore .create_index(create_index_request) .await @@ -226,7 +230,8 @@ pub async fn test_metastore_publish_splits< // Publish a staged split on an index { let create_index_request = - CreateIndexRequest::try_from_index_config(&index_config).unwrap(); + CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs) + .unwrap(); let index_uid: IndexUid = metastore .create_index(create_index_request) .await @@ -255,7 +260,8 @@ pub async fn test_metastore_publish_splits< // Publish a published split on an index { let create_index_request = - CreateIndexRequest::try_from_index_config(&index_config).unwrap(); + CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs) + .unwrap(); let index_uid: IndexUid = metastore .create_index(create_index_request) .await @@ -306,7 +312,8 @@ pub async fn test_metastore_publish_splits< // Publish a non-staged split on an index { let create_index_request = - CreateIndexRequest::try_from_index_config(&index_config).unwrap(); + CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs) + .unwrap(); let index_uid: IndexUid = metastore .create_index(create_index_request) .await @@ -369,7 +376,8 @@ pub async fn test_metastore_publish_splits< // Publish a staged split and non-existent split on an index { let create_index_request = - CreateIndexRequest::try_from_index_config(&index_config).unwrap(); + CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs) + .unwrap(); let index_uid: IndexUid = metastore .create_index(create_index_request) .await @@ -407,7 +415,8 @@ pub async fn test_metastore_publish_splits< // Publish a published split and non-existent split on an index { let create_index_request = - CreateIndexRequest::try_from_index_config(&index_config).unwrap(); + CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs) + .unwrap(); let index_uid: IndexUid = metastore .create_index(create_index_request) .await @@ -460,7 +469,8 @@ pub async fn test_metastore_publish_splits< // Publish a non-staged split and non-existent split on an index { let create_index_request = - CreateIndexRequest::try_from_index_config(&index_config).unwrap(); + CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs) + .unwrap(); let index_uid: IndexUid = metastore .create_index(create_index_request) .await @@ -520,7 +530,8 @@ pub async fn test_metastore_publish_splits< // Publish staged splits on an index { let create_index_request = - CreateIndexRequest::try_from_index_config(&index_config).unwrap(); + CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs) + .unwrap(); let index_uid: IndexUid = metastore .create_index(create_index_request) .await @@ -559,7 +570,8 @@ pub async fn test_metastore_publish_splits< // Publish a staged split and published split on an index { let create_index_request = - CreateIndexRequest::try_from_index_config(&index_config).unwrap(); + CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs) + .unwrap(); let index_uid: IndexUid = metastore .create_index(create_index_request) .await @@ -617,7 +629,8 @@ pub async fn test_metastore_publish_splits< // Publish published splits on an index { let create_index_request = - CreateIndexRequest::try_from_index_config(&index_config).unwrap(); + CreateIndexRequest::try_from_index_and_source_configs(&index_config, source_configs) + .unwrap(); let index_uid: IndexUid = metastore .create_index(create_index_request) .await @@ -681,7 +694,12 @@ pub async fn test_metastore_publish_splits_concurrency< let index_id = append_random_suffix("test-publish-concurrency"); let index_uri = format!("ram:///indexes/{index_id}"); let index_config = IndexConfig::for_test(&index_id, &index_uri); - let create_index_request = CreateIndexRequest::try_from_index_config(&index_config).unwrap(); + let source_id = format!("{index_id}--source"); + + let source_config = SourceConfig::for_test(&source_id, SourceParams::void()); + let create_index_request = + CreateIndexRequest::try_from_index_and_source_configs(&index_config, &[source_config]) + .unwrap(); let index_uid: IndexUid = metastore .create_index(create_index_request) @@ -690,8 +708,6 @@ pub async fn test_metastore_publish_splits_concurrency< .index_uid() .clone(); - let source_id = format!("{index_id}--source"); - let mut join_handles = Vec::with_capacity(10); for partition_id in 0..10 { @@ -1430,6 +1446,7 @@ pub async fn test_metastore_split_update_timestamp< let index_config = IndexConfig::for_test(&index_id, &index_uri); let source_id = format!("{index_id}--source"); + let source_config = SourceConfig::for_test(&source_id, SourceParams::void()); let split_id = format!("{index_id}--split"); let split_metadata = SplitMetadata { @@ -1440,7 +1457,9 @@ pub async fn test_metastore_split_update_timestamp< }; // Create an index - let create_index_request = CreateIndexRequest::try_from_index_config(&index_config).unwrap(); + let create_index_request = + CreateIndexRequest::try_from_index_and_source_configs(&index_config, &[source_config]) + .unwrap(); let index_uid: IndexUid = metastore .create_index(create_index_request) .await diff --git a/quickwit/quickwit-proto/protos/quickwit/metastore.proto b/quickwit/quickwit-proto/protos/quickwit/metastore.proto index 4daf5b6d171..b3cd3a7898d 100644 --- a/quickwit/quickwit-proto/protos/quickwit/metastore.proto +++ b/quickwit/quickwit-proto/protos/quickwit/metastore.proto @@ -41,6 +41,7 @@ enum SourceType { SOURCE_TYPE_PULSAR = 9; SOURCE_TYPE_VEC = 10; SOURCE_TYPE_VOID = 11; + SOURCE_TYPE_STDIN = 13; } // Metastore meant to manage Quickwit's indexes, their splits and delete tasks. @@ -406,6 +407,7 @@ message OpenShardSubrequest { string leader_id = 5; optional string follower_id = 6; quickwit.common.DocMappingUid doc_mapping_uid = 7; + optional string publish_token = 8; } message OpenShardsResponse { diff --git a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.metastore.rs b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.metastore.rs index 4b3d88ee7ed..1f0f36db21c 100644 --- a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.metastore.rs +++ b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.metastore.rs @@ -338,6 +338,8 @@ pub struct OpenShardSubrequest { pub follower_id: ::core::option::Option<::prost::alloc::string::String>, #[prost(message, optional, tag = "7")] pub doc_mapping_uid: ::core::option::Option, + #[prost(string, optional, tag = "8")] + pub publish_token: ::core::option::Option<::prost::alloc::string::String>, } #[derive(serde::Serialize, serde::Deserialize, utoipa::ToSchema)] #[allow(clippy::derive_partial_eq_without_eq)] @@ -529,6 +531,7 @@ pub enum SourceType { Pulsar = 9, Vec = 10, Void = 11, + Stdin = 13, } impl SourceType { /// String value of the enum field names used in the ProtoBuf definition. @@ -549,6 +552,7 @@ impl SourceType { SourceType::Pulsar => "SOURCE_TYPE_PULSAR", SourceType::Vec => "SOURCE_TYPE_VEC", SourceType::Void => "SOURCE_TYPE_VOID", + SourceType::Stdin => "SOURCE_TYPE_STDIN", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -566,6 +570,7 @@ impl SourceType { "SOURCE_TYPE_PULSAR" => Some(Self::Pulsar), "SOURCE_TYPE_VEC" => Some(Self::Vec), "SOURCE_TYPE_VOID" => Some(Self::Void), + "SOURCE_TYPE_STDIN" => Some(Self::Stdin), _ => None, } } diff --git a/quickwit/quickwit-proto/src/metastore/mod.rs b/quickwit/quickwit-proto/src/metastore/mod.rs index 6caba6f7f1d..4782dac03c2 100644 --- a/quickwit/quickwit-proto/src/metastore/mod.rs +++ b/quickwit/quickwit-proto/src/metastore/mod.rs @@ -306,6 +306,7 @@ impl SourceType { SourceType::Nats => "nats", SourceType::PubSub => "pubsub", SourceType::Pulsar => "pulsar", + SourceType::Stdin => "stdin", SourceType::Unspecified => "unspecified", SourceType::Vec => "vec", SourceType::Void => "void", @@ -325,6 +326,7 @@ impl fmt::Display for SourceType { SourceType::Nats => "NATS", SourceType::PubSub => "Google Cloud Pub/Sub", SourceType::Pulsar => "Apache Pulsar", + SourceType::Stdin => "Stdin", SourceType::Unspecified => "unspecified", SourceType::Vec => "vec", SourceType::Void => "void", diff --git a/quickwit/quickwit-serve/src/index_api/rest_handler.rs b/quickwit/quickwit-serve/src/index_api/rest_handler.rs index 778911e8d10..827c2b027ca 100644 --- a/quickwit/quickwit-serve/src/index_api/rest_handler.rs +++ b/quickwit/quickwit-serve/src/index_api/rest_handler.rs @@ -23,7 +23,8 @@ use bytes::Bytes; use quickwit_common::uri::Uri; use quickwit_config::{ load_index_config_update, load_source_config_from_user_config, validate_index_id_pattern, - ConfigFormat, NodeConfig, SourceConfig, SourceParams, CLI_SOURCE_ID, INGEST_API_SOURCE_ID, + ConfigFormat, FileSourceParams, NodeConfig, SourceConfig, SourceParams, CLI_SOURCE_ID, + INGEST_API_SOURCE_ID, }; use quickwit_doc_mapper::{analyze_text, TokenizerConfig}; use quickwit_index_management::{IndexService, IndexServiceError}; @@ -708,12 +709,15 @@ async fn create_source( let source_config: SourceConfig = load_source_config_from_user_config(config_format, &source_config_bytes) .map_err(IndexServiceError::InvalidConfig)?; - if let SourceParams::File(_) = &source_config.source_params { - return Err(IndexServiceError::OperationNotAllowed( - "file sources are limited to a local usage. please use the CLI command `quickwit tool \ - local-ingest` to ingest data from a file" - .to_string(), - )); + // Note: This check is performed here instead of the source config serde + // because many tests use the file source, and can't store that config in + // the metastore without going through the validation. + if let SourceParams::File(FileSourceParams::Filepath(_)) = &source_config.source_params { + return Err(IndexServiceError::InvalidConfig(anyhow::anyhow!( + "path based file sources are limited to a local usage, please use the CLI command \ + `quickwit tool local-ingest` to ingest data from a specific file or setup a \ + notification based file source" + ))); } let index_metadata_request = IndexMetadataRequest::for_index_id(index_id.to_string()); let index_uid: IndexUid = index_service @@ -1641,28 +1645,6 @@ mod tests { assert!(indexes.is_empty()); } - #[tokio::test] - async fn test_create_file_source_returns_403() { - let metastore = metastore_for_test(); - let index_service = IndexService::new(metastore.clone(), StorageResolver::unconfigured()); - let mut node_config = NodeConfig::for_test(); - node_config.default_index_root_uri = Uri::for_test("file:///default-index-root-uri"); - let index_management_handler = - super::index_management_handlers(index_service, Arc::new(node_config)) - .recover(recover_fn); - let source_config_body = r#"{"version": "0.7", "source_id": "file-source", "source_type": - "file", "params": {"filepath": "FILEPATH"}}"#; - let resp = warp::test::request() - .path("/indexes/hdfs-logs/sources") - .method("POST") - .body(source_config_body) - .reply(&index_management_handler) - .await; - assert_eq!(resp.status(), 403); - let response_body = std::str::from_utf8(resp.body()).unwrap(); - assert!(response_body.contains("limited to a local usage")) - } - #[tokio::test] async fn test_create_index_with_yaml() { let metastore = metastore_for_test(); @@ -1886,6 +1868,34 @@ mod tests { sources" )); } + { + let resp = warp::test::request() + .path("/indexes/hdfs-logs/sources") + .method("POST") + .body( + r#"{"version": "0.8", "source_id": "my-stdin-source", "source_type": "stdin"}"#, + ) + .reply(&index_management_handler) + .await; + assert_eq!(resp.status(), 400); + let response_body = std::str::from_utf8(resp.body()).unwrap(); + assert!( + response_body.contains("stdin can only be used as source through the CLI command") + ) + } + { + let resp = warp::test::request() + .path("/indexes/hdfs-logs/sources") + .method("POST") + .body( + r#"{"version": "0.8", "source_id": "my-local-file-source", "source_type": "file", "params": {"filepath": "localfile"}}"#, + ) + .reply(&index_management_handler) + .await; + assert_eq!(resp.status(), 400); + let response_body = std::str::from_utf8(resp.body()).unwrap(); + assert!(response_body.contains("limited to a local usage")) + } } #[tokio::test] diff --git a/quickwit/quickwit-storage/src/storage.rs b/quickwit/quickwit-storage/src/storage.rs index 9ca1dd2041f..9e8f1c54000 100644 --- a/quickwit/quickwit-storage/src/storage.rs +++ b/quickwit/quickwit-storage/src/storage.rs @@ -103,7 +103,9 @@ pub trait Storage: fmt::Debug + Send + Sync + 'static { /// Downloads a slice of a file from the storage, and returns an in memory buffer async fn get_slice(&self, path: &Path, range: Range) -> StorageResult; - /// Open a stream handle on the file from the storage + /// Opens a stream handle on the file from the storage. + /// + /// Might panic, return an error or an empty stream if the range is empty. async fn get_slice_stream( &self, path: &Path,