From 7312c4e006035e6cadaafb2ecc0a3880425ca7f0 Mon Sep 17 00:00:00 2001 From: kris927b Date: Fri, 20 Jun 2025 15:47:06 +0200 Subject: [PATCH] Major refactoring of worker and producer logic. Added many new tests. coverage is now 67.5% --- Cargo.lock | 1 + Cargo.toml | 1 + src/bin/producer.rs | 40 ++- src/bin/worker.rs | 373 +-------------------- src/config/mod.rs | 1 + src/config/producer.rs | 4 + src/config/worker.rs | 39 +++ src/lib.rs | 1 + src/pipeline/writers/parquet_writer.rs | 2 +- src/producer_logic.rs | 183 ++--------- src/utils/common.rs | 79 ++++- src/worker_logic.rs | 283 ++++++++++++++++ tests/full_pipeline_test.rs | 6 +- tests/parquet_io_test.rs | 6 +- tests/producer_tests.rs | 438 ++++++++++--------------- tests/worker_tests.rs | 156 +++++++++ 16 files changed, 800 insertions(+), 813 deletions(-) create mode 100644 src/config/worker.rs create mode 100644 src/worker_logic.rs create mode 100644 tests/worker_tests.rs diff --git a/Cargo.lock b/Cargo.lock index 0dccb05..ee2bcaa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -42,6 +42,7 @@ dependencies = [ "tracing", "tracing-appender", "tracing-subscriber", + "uuid", "whatlang", ] diff --git a/Cargo.toml b/Cargo.toml index 6e1a286..9f707e0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -63,5 +63,6 @@ serde_json = "1.0" async-stream = "0.3.6" tokio-stream = "0.1.17" polars = { version = "0.49.1", features = ["lazy", "parquet"] } +uuid = "1.17.0" # {{ No need for the old default main.rs binary now }} diff --git a/src/bin/producer.rs b/src/bin/producer.rs index 68c6d89..d2ea577 100644 --- a/src/bin/producer.rs +++ b/src/bin/producer.rs @@ -16,7 +16,7 @@ use tracing_subscriber::{fmt, layer::SubscriberExt, util::SubscriberInitExt, Env // TextBlaster::error::{PipelineError, Result} is used use TextBlaster::error::{PipelineError, Result}; // TextBlaster::pipeline::writers::parquet_writer::ParquetWriter is used by producer_logic -use TextBlaster::utils::common::connect_rabbitmq; +use TextBlaster::utils::common::{connect_rabbitmq, setup_channels_and_queues}; use TextBlaster::utils::prometheus_metrics::setup_prometheus_metrics; // chrono::Utc is used by producer_logic // TextBlaster::utils::prometheus_metrics::* is used by producer_logic @@ -99,8 +99,15 @@ async fn main() -> Result<()> { let conn = connect_rabbitmq(&args.amqp_addr).await?; // This returns a lapin::Connection // Create channels for publishing and consuming results - let task_publish_channel = conn.create_channel().await.map_err(PipelineError::from)?; - let results_consume_channel = conn.create_channel().await.map_err(PipelineError::from)?; + let (publish_channel, consumer) = setup_channels_and_queues( + &conn, + &args.task_queue, + &args.results_queue, + args.prefetch_count, + "producer".to_string(), + ) + .await + .unwrap(); // Optionally, set task_publish_channel to confirm mode if desired for all publishes // task_publish_channel.confirm_select(lapin::options::ConfirmSelectOptions::default()).await @@ -112,20 +119,17 @@ async fn main() -> Result<()> { let publishing_pb = create_progress_bar(0, "Publishing tasks", publishing_pb_template); // 2. Publish Tasks - now passing the channel directly - let published_count = match TextBlaster::producer_logic::publish_tasks( - &args, - &task_publish_channel, - &publishing_pb, - ) - .await - { - Ok(count) => count, - Err(e) => { - error!("Failed during task publishing: {}", e); - publishing_pb.finish_with_message(format!("Publishing failed: {}", e)); - return Err(e); - } - }; + let published_count = + match TextBlaster::producer_logic::publish_tasks(&args, &publish_channel, &publishing_pb) + .await + { + Ok(count) => count, + Err(e) => { + error!("Failed during task publishing: {}", e); + publishing_pb.finish_with_message(format!("Publishing failed: {}", e)); + return Err(e); + } + }; // Early exit if no tasks were published (nothing to wait for) if published_count == 0 { @@ -147,7 +151,7 @@ async fn main() -> Result<()> { let (outcomes_received_count, success_count, filtered_count) = match TextBlaster::producer_logic::aggregate_results( &args, - &results_consume_channel, + consumer, published_count, &aggregation_pb, ) diff --git a/src/bin/worker.rs b/src/bin/worker.rs index 6f7a5a3..3d301fe 100644 --- a/src/bin/worker.rs +++ b/src/bin/worker.rs @@ -1,382 +1,27 @@ // src/bin/worker.rs use clap::Parser; -use futures::StreamExt; // For processing the consumer stream use indicatif::{ProgressBar, ProgressStyle}; -use std::time::{Duration, Instant}; // Added for progress bar // Added for progress bar speed calculation - // {{ Use the new load_pipeline_config function }} -use TextBlaster::config::pipeline::{load_pipeline_config, PipelineConfig, StepConfig}; // Added config imports and load_pipeline_config -use TextBlaster::data_model::{ProcessingOutcome, TextDocument}; // Updated import +use std::time::{Duration, Instant}; +use TextBlaster::config::worker::Args; +use TextBlaster::worker_logic::{build_pipeline_from_config, process_tasks_with_executor}; // Added for progress bar // Added for progress bar speed calculation + // {{ Use the new load_pipeline_config function }} +use TextBlaster::config::pipeline::{load_pipeline_config, PipelineConfig}; use TextBlaster::error::{PipelineError, Result}; // Use the library's Result type -use TextBlaster::executor::{PipelineExecutor, ProcessingStep}; +use TextBlaster::executor::PipelineExecutor; // Import necessary filters (adjust if steps change) // {{ Remove ArrowNativeType import if no longer needed directly here }} // use arrow::datatypes::ArrowNativeType; -use lapin::{ - options::{ - BasicAckOptions, - BasicConsumeOptions, - BasicPublishOptions, - BasicQosOptions, // Added BasicPublishOptions - QueueDeclareOptions, - }, - protocol::basic::AMQPProperties, // Added AMQPProperties - types::FieldTable, -}; -use TextBlaster::pipeline::filters::{ - C4BadWordsFilter, - C4QualityFilter, - FineWebQualityFilter, - GopherQualityFilter, // Updated import - GopherRepetitionFilter, - LanguageDetectionFilter, -}; -use TextBlaster::pipeline::token::TokenCounter; + use TextBlaster::utils::common::connect_rabbitmq; // Updated for shared functions use TextBlaster::utils::prometheus_metrics::*; -use std::path::PathBuf; use std::sync::Arc; // To share the executor across potential concurrent tasks // {{ Add serde_json for result serialization }} -use tracing::{debug, error, info, info_span, instrument, warn}; // Added tracing +use tracing::{error, info}; // Added tracing use tracing_appender::{non_blocking, rolling}; use tracing_subscriber::{fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer}; // Added tracing_subscriber // Added for file logging -// Define command-line arguments -#[derive(Parser, Debug)] -#[command(author, version, about, long_about = None)] -struct Args { - /// RabbitMQ connection string (e.g., amqp://guest:guest@localhost:5672/%2f) - #[arg(short, long, default_value = "amqp://guest:guest@localhost:5672/%2f")] - amqp_addr: String, - - /// Name of the queue to consume tasks from - #[arg(short = 'q', long, default_value = "task_queue")] - // Use short arg 'q' consistent with producer - task_queue: String, - - /// Name of the queue to publish results/outcomes to - #[arg(short = 'r', long, default_value = "results_queue")] - // Use short arg 'r' consistent with producer - results_queue: String, - - /// Prefetch count (how many messages to buffer locally) - #[arg(long, default_value_t = 10)] // Adjust based on task duration/resources - prefetch_count: u16, - - // {{ Add argument for pipeline configuration file }} - /// Path to the pipeline configuration YAML file. - #[arg(short = 'c', long, default_value = "config/pipeline_config.yaml")] - pipeline_config: PathBuf, - - /// Optional: Port for the Prometheus metrics HTTP endpoint - #[arg(long)] - metrics_port: Option, - - /// Validate the pipeline configuration and exit - #[arg(long)] - validate_config: bool, -} - -// --- Prometheus Metrics (now imported from TextBlaster::utils::prometheus_metrics) --- -// The local static definitions and specific prometheus imports are removed. - -// {{ Add the new function to build pipeline from configuration }} -/// Builds the processing pipeline based on the configuration read from YAML. -#[instrument(skip(config), fields(num_steps = config.pipeline.len()))] -fn build_pipeline_from_config(config: &PipelineConfig) -> Result>> { - let mut steps: Vec> = Vec::new(); - info!("Building pipeline from configuration..."); - - for (i, step_config) in config.pipeline.iter().enumerate() { - let step_span = info_span!("pipeline_step", index = i, type = step_config.name()); - let _enter = step_span.enter(); - - let step: Box = match step_config { - StepConfig::C4QualityFilter(params) => { - debug!(params = ?params, "Adding C4QualityFilter"); - Box::new(C4QualityFilter::new( - params.split_paragraph, - params.remove_citations, - params.filter_no_terminal_punct, - params.min_num_sentences, - params.min_words_per_line, - params.max_word_length, - params.filter_lorem_ipsum, - params.filter_javascript, - params.filter_curly_bracket, - params.filter_policy, - )) - } - StepConfig::GopherRepetitionFilter(params) => { - debug!(params = ?params, "Adding GopherRepetitionFilter"); - Box::new(GopherRepetitionFilter::new( - params.dup_line_frac, - params.dup_para_frac, - params.dup_line_char_frac, - params.dup_para_char_frac, - params.top_n_grams.clone(), // Clone the vec - params.dup_n_grams.clone(), // Clone the vec - )) - } - StepConfig::GopherQualityFilter(params) => { - debug!(params = ?params, "Adding GopherQualityFilter"); - Box::new(GopherQualityFilter::new( - params.min_doc_words, - params.max_doc_words, - params.min_avg_word_length, - params.max_avg_word_length, - params.max_symbol_word_ratio, - params.max_bullet_lines_ratio, - params.max_ellipsis_lines_ratio, - params.max_non_alpha_words_ratio, - params.min_stop_words, - params.stop_words.clone(), // Clone the Option> - )) - } - StepConfig::LanguageDetectionFilter(params) => { - debug!(params = ?params, "Adding LanguageDetectionFilter"); - Box::new(LanguageDetectionFilter::new( - params.min_confidence, - params.allowed_languages.clone(), - )) - } - StepConfig::C4BadWordsFilter(params) => { - debug!(params = ?params, "Adding C4BadWordsFilter"); - Box::new(C4BadWordsFilter::new(params.clone())) - } - StepConfig::FineWebQualityFilter(params) => { - // Updated variant name - debug!(params = ?params, "Adding FineWebQualityFilter"); - - Box::new(FineWebQualityFilter::new( - params.line_punct_thr, - params.line_punct_exclude_zero, - params.short_line_thr, - params.short_line_length, - params.char_duplicates_ratio, - params.new_line_ratio, - // params.language.clone(), - params.stop_chars.clone(), - )) - } - StepConfig::TokenCounter(params) => { - debug!(params=?params, "Adding TokenCounter"); - let token_step = TokenCounter::new(¶ms.tokenizer_name); - if let Err(e) = token_step { - panic!("{}", e); - } - Box::new(token_step.unwrap()) - } - }; - steps.push(step); - info!("Added step: {}", step_config.name()); - } - - if steps.is_empty() { - warn!("Warning: Building an empty pipeline from configuration!"); - } else { - info!("Pipeline built successfully with {} steps.", steps.len()); - } - Ok(steps) -} - -async fn process_tasks( - args: &Args, - conn: &lapin::Connection, - executor: Arc, -) -> Result<()> { - // Create two channels: one for consuming tasks, one for publishing results/outcomes - let consume_channel = conn.create_channel().await.map_err(|e| { - PipelineError::QueueError(format!("Worker failed to create consume channel: {}", e)) - })?; - let publish_channel = conn.create_channel().await.map_err(|e| { - PipelineError::QueueError(format!("Worker failed to create publish channel: {}", e)) - })?; - - // Declare the task queue (ensure durability matches producer) - let _task_queue = consume_channel - .queue_declare( - &args.task_queue, - QueueDeclareOptions { - durable: true, // MUST match the producer's declaration - ..Default::default() - }, - Default::default(), - ) - .await - .map_err(|e| { - PipelineError::QueueError(format!("Worker failed to declare task queue: {}", e)) - })?; - - // Declare the results queue (also durable) - let _results_queue = publish_channel - .queue_declare( - &args.results_queue, - QueueDeclareOptions { - durable: true, // Results/Outcomes should also survive restarts - ..Default::default() - }, - Default::default(), - ) - .await - .map_err(|e| { - PipelineError::QueueError(format!("Worker failed to declare results queue: {}", e)) - })?; - - // Set Quality of Service (Prefetch Count) on the consume channel - consume_channel - .basic_qos(args.prefetch_count, BasicQosOptions::default()) - .await - .map_err(|e| PipelineError::QueueError(format!("Failed to set QoS: {}", e)))?; - - // Start Consuming Messages from the task queue - let consumer_tag = format!( - "worker-{}-{}", - std::process::id(), - chrono::Utc::now().timestamp() // Ensure chrono::Utc is imported - ); - info!(consumer_tag = %consumer_tag, "Worker started consuming tasks. Waiting for messages..."); - - let mut consumer = consume_channel - .basic_consume( - &args.task_queue, - &consumer_tag, - BasicConsumeOptions::default(), // auto_ack: false (default) - FieldTable::default(), - ) - .await - .map_err(|e| PipelineError::QueueError(format!("Failed to start consuming: {}", e)))?; - - // Process messages from the stream - while let Some(delivery_result) = consumer.next().await { - match delivery_result { - Ok(delivery) => { - let executor_clone = Arc::clone(&executor); - let publish_channel_clone = publish_channel.clone(); - let results_queue_name = args.results_queue.clone(); - // let worker_id_tag = consumer_tag.clone(); // Use the specific consumer_tag for this worker instance - - tokio::spawn(async move { - ACTIVE_PROCESSING_TASKS.inc(); - let processing_timer = TASK_PROCESSING_DURATION_SECONDS.start_timer(); - - let result: Option = match serde_json::from_slice::< - TextDocument, - >( - &delivery.data - ) { - Ok(doc) => { - let original_doc_id = doc.id.clone(); - let task_span = info_span!("process_task", doc_id = %original_doc_id, delivery_tag = %delivery.delivery_tag); - let _enter = task_span.enter(); - debug!("Processing document"); - - match executor_clone.run_single_async(doc.clone()).await { - Ok(processed_doc) => { - debug!(processed_doc_id = %processed_doc.id, "Successfully processed document"); - TASKS_PROCESSED_TOTAL.inc(); - Some(ProcessingOutcome::Success(processed_doc)) - } - Err(pipeline_error) => { - if let PipelineError::StepError { step_name, source } = - pipeline_error - { - match *source { - PipelineError::DocumentFiltered { - document, - reason, - } => { - info!(filtered_doc_id = %document.id, %step_name, %reason, "Document was filtered"); - TASKS_FILTERED_TOTAL.inc(); // Increment filtered counter - Some(ProcessingOutcome::Filtered { - document: *document, - reason, - }) - } - other_error => { - error!(%step_name, error = %other_error, "Pipeline step failed"); - TASKS_FAILED_TOTAL.inc(); // Increment failed counter - None // Don't send outcome for pipeline errors - } - } - } else { - error!(error = %pipeline_error, "Unexpected pipeline error"); - TASKS_FAILED_TOTAL.inc(); // Increment failed counter - None // Don't send outcome - } - } - } - } - Err(e) => { - error!( - delivery_tag = %delivery.delivery_tag, - error = %e, - payload = %String::from_utf8_lossy(&delivery.data), - "Failed to deserialize task message" - ); - TASK_DESERIALIZATION_ERRORS_TOTAL.inc(); - None - } - }; - - if let Some(actual_outcome) = result { - match serde_json::to_vec(&actual_outcome) { - Ok(payload) => { - let publish_confirm = publish_channel_clone - .basic_publish( - "", - &results_queue_name, - BasicPublishOptions::default(), - &payload, - AMQPProperties::default().with_delivery_mode(2), - ) - .await; - - match publish_confirm { - Ok(confirmation) => match confirmation.await { - Ok(_) => debug!("Published outcome"), - Err(e) => { - error!(error = %e, "Failed publish confirmation for outcome"); - OUTCOME_PUBLISH_ERRORS_TOTAL.inc(); - } - }, - Err(e) => { - error!(error = %e, "Failed to initiate publish for outcome"); - OUTCOME_PUBLISH_ERRORS_TOTAL.inc(); - } - } - } - Err(e) => { - error!(error = %e, "Failed to serialize outcome"); - } - } - } - - if let Err(ack_err) = delivery.ack(BasicAckOptions::default()).await { - error!(error = %ack_err, "Failed to ack task message"); - } - - processing_timer.observe_duration(); - ACTIVE_PROCESSING_TASKS.dec(); - }); - } - Err(e) => { - error!(error = %e, "Error receiving task message from consumer stream. Worker will stop."); - // This error will propagate up from process_tasks if the loop breaks - return Err(PipelineError::QueueError(format!( - "Consumer stream error: {}", - e - ))); - } - } - } - // If the loop finishes (e.g. queue deleted, channel closed gracefully), it's not necessarily an error. - // Specific errors during consumption (like connection loss) would break the loop and return Err. - info!("Worker stopped consuming tasks (consumer stream ended)."); - Ok(()) -} - #[tokio::main] async fn main() -> Result<()> { let args = Args::parse(); @@ -509,7 +154,7 @@ async fn main() -> Result<()> { .map_err(|e| PipelineError::QueueError(format!("Worker failed to connect: {}", e)))?; // Process tasks - let task_processing_result = process_tasks(&args, &conn, executor).await; + let task_processing_result = process_tasks_with_executor(&args, &conn, executor).await; // Stop the progress bar and the updater task pb.finish_with_message("Processing finished or interrupted."); // Updated message diff --git a/src/config/mod.rs b/src/config/mod.rs index d62ca31..997dd36 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -3,3 +3,4 @@ pub mod parquet; pub mod pipeline; pub mod producer; +pub mod worker; diff --git a/src/config/producer.rs b/src/config/producer.rs index f4b69b7..ba7a905 100644 --- a/src/config/producer.rs +++ b/src/config/producer.rs @@ -29,6 +29,10 @@ pub struct Args { #[arg(short = 'r', long, default_value = "results_queue")] pub results_queue: String, + /// Prefetch count (how many messages to buffer locally) + #[arg(long, default_value_t = 10)] // Adjust based on task duration/resources + pub prefetch_count: u16, + /// Path to the output Parquet file #[arg(short = 'o', long, default_value = "output_processed.parquet")] pub output_file: String, diff --git a/src/config/worker.rs b/src/config/worker.rs new file mode 100644 index 0000000..7175f7a --- /dev/null +++ b/src/config/worker.rs @@ -0,0 +1,39 @@ +use std::path::PathBuf; + +use clap::Parser; + +// Define command-line arguments +#[derive(Parser, Debug)] +#[command(author, version, about, long_about = None)] +pub struct Args { + /// RabbitMQ connection string (e.g., amqp://guest:guest@localhost:5672/%2f) + #[arg(short, long, default_value = "amqp://guest:guest@localhost:5672/%2f")] + pub amqp_addr: String, + + /// Name of the queue to consume tasks from + #[arg(short = 'q', long, default_value = "task_queue")] + // Use short arg 'q' consistent with producer + pub task_queue: String, + + /// Name of the queue to publish results/outcomes to + #[arg(short = 'r', long, default_value = "results_queue")] + // Use short arg 'r' consistent with producer + pub results_queue: String, + + /// Prefetch count (how many messages to buffer locally) + #[arg(long, default_value_t = 10)] // Adjust based on task duration/resources + pub prefetch_count: u16, + + // {{ Add argument for pipeline configuration file }} + /// Path to the pipeline configuration YAML file. + #[arg(short = 'c', long, default_value = "config/pipeline_config.yaml")] + pub pipeline_config: PathBuf, + + /// Optional: Port for the Prometheus metrics HTTP endpoint + #[arg(long)] + pub metrics_port: Option, + + /// Validate the pipeline configuration and exit + #[arg(long)] + pub validate_config: bool, +} diff --git a/src/lib.rs b/src/lib.rs index e0c553a..dc31073 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,6 +16,7 @@ pub mod utils; // pub use executor::{PipelineExecutor, ProcessingStep}; pub mod producer_logic; // Declare the new module +pub mod worker_logic; // The AmqpConnectionManager trait and its implementation have been removed from here // as per the new strategy focusing on TaskPublisherChannel defined in producer_logic.rs. diff --git a/src/pipeline/writers/parquet_writer.rs b/src/pipeline/writers/parquet_writer.rs index e56981b..c9c8c1c 100644 --- a/src/pipeline/writers/parquet_writer.rs +++ b/src/pipeline/writers/parquet_writer.rs @@ -17,7 +17,7 @@ fn create_schema() -> SchemaRef { Arc::new(Schema::new(vec![ Field::new("id", DataType::Utf8, false), Field::new("source", DataType::Utf8, false), - Field::new("content", DataType::Utf8, false), + Field::new("text", DataType::Utf8, false), // Store metadata as a single JSON string. Nullable if metadata might be empty. Field::new("metadata", DataType::Utf8, true), ])) diff --git a/src/producer_logic.rs b/src/producer_logic.rs index d804068..d3fe11a 100644 --- a/src/producer_logic.rs +++ b/src/producer_logic.rs @@ -7,21 +7,12 @@ use crate::error::{PipelineError, Result as AppResult}; use crate::pipeline::readers::ParquetReader; use crate::pipeline::writers::parquet_writer::ParquetWriter; use crate::utils::prometheus_metrics::*; -use async_trait::async_trait; -use chrono::Utc; use futures::{pin_mut, Stream, StreamExt}; use indicatif::ProgressBar; use lapin::{ - options::{ - BasicAckOptions, BasicConsumeOptions, BasicPublishOptions, ConfirmSelectOptions, - QueueDeclareOptions, - }, + options::{BasicAckOptions, BasicPublishOptions}, protocol::basic::AMQPProperties, - publisher_confirm::Confirmation, - types::FieldTable, - Channel as LapinChannel, // Alias lapin::Channel to avoid confusion - Consumer, - Result as LapinResult, + Channel, Consumer, }; use serde_json; use std::time::Instant; @@ -29,69 +20,11 @@ use tracing::{error, info, info_span, warn}; // For aggregate_results consumer t pub const PARQUET_WRITE_BATCH_SIZE: usize = 500; -#[async_trait] -pub trait TaskPublisherChannel: Send + Sync { - async fn queue_declare( - &self, - name: &str, - options: QueueDeclareOptions, - arguments: FieldTable, - ) -> LapinResult<()>; - async fn basic_publish( - &self, - exchange: &str, - routing_key: &str, - options: BasicPublishOptions, - payload: &[u8], - properties: AMQPProperties, - ) -> LapinResult; - async fn confirm_select(&self, options: ConfirmSelectOptions) -> LapinResult<()>; -} - -#[async_trait] -impl TaskPublisherChannel for LapinChannel { - async fn queue_declare( - &self, - name: &str, - options: QueueDeclareOptions, - arguments: FieldTable, - ) -> LapinResult<()> { - LapinChannel::queue_declare(self, name, options, arguments).await?; - Ok(()) - } - async fn basic_publish( - &self, - exchange: &str, - routing_key: &str, - options: BasicPublishOptions, - payload: &[u8], - properties: AMQPProperties, - ) -> LapinResult { - let publisher_confirmation = - LapinChannel::basic_publish(self, exchange, routing_key, options, payload, properties) - .await?; - publisher_confirmation.await - } - async fn confirm_select(&self, options: ConfirmSelectOptions) -> LapinResult<()> { - LapinChannel::confirm_select(self, options).await - } -} - -pub async fn publish_tasks( +pub async fn publish_tasks( args: &Args, - publish_channel: &CH, + publish_channel: &Channel, publishing_pb: &ProgressBar, ) -> AppResult { - publish_channel - .queue_declare( - &args.task_queue, - QueueDeclareOptions { - durable: true, - ..Default::default() - }, - FieldTable::default(), - ) - .await?; info!("Declared durable task queue '{}'", args.task_queue); let parquet_config = ParquetInputConfig { @@ -126,27 +59,23 @@ pub async fn publish_tasks( &payload, AMQPProperties::default().with_delivery_mode(2), ) - .await?; + .await; task_publish_timer.observe_duration(); match confirmation { - Confirmation::Ack(_) | Confirmation::NotRequested => { + Ok(_) => { published_count += 1; TASKS_PUBLISHED_TOTAL.inc(); ACTIVE_TASKS_IN_FLIGHT.inc(); publishing_pb.inc(1); - if matches!(confirmation, Confirmation::Ack(_)) { - info!("Successfully published task and received ACK."); - } else { - info!("Successfully published task (no confirmation requested/received)."); - } + info!("Successfully published task"); } - Confirmation::Nack(_) => { + Err(e) => { TASK_PUBLISH_ERRORS_TOTAL.inc(); error!(doc_id = %doc.id, "FATAL: Broker NACKed task. Stopping."); return Err(PipelineError::QueueError(format!( - "Publish confirmation failed (NACK) for doc {}", - doc.id + "Publish confirmation failed (NACK) for doc {}. Error {:?}", + doc.id, e ))); } } @@ -174,45 +103,6 @@ pub async fn publish_tasks( Ok(published_count) } -#[async_trait] -pub trait ResultConsumerChannel: Send + Sync { - async fn queue_declare( - &self, - name: &str, - options: QueueDeclareOptions, - arguments: FieldTable, - ) -> LapinResult<()>; - async fn basic_consume( - &self, - queue: &str, - consumer_tag: &str, - options: BasicConsumeOptions, - arguments: FieldTable, - ) -> LapinResult; -} - -#[async_trait] -impl ResultConsumerChannel for LapinChannel { - async fn queue_declare( - &self, - name: &str, - options: QueueDeclareOptions, - arguments: FieldTable, - ) -> LapinResult<()> { - let _ = LapinChannel::queue_declare(self, name, options, arguments).await; - Ok(()) - } - async fn basic_consume( - &self, - queue: &str, - consumer_tag: &str, - options: BasicConsumeOptions, - arguments: FieldTable, - ) -> LapinResult { - LapinChannel::basic_consume(self, queue, consumer_tag, options, arguments).await - } -} - pub async fn aggregate_results_from_stream( args: &Args, stream: S, @@ -302,53 +192,30 @@ where Ok((outcomes_received_count, success_count, filtered_count)) } -pub async fn aggregate_results( +pub async fn aggregate_results( args: &Args, - consume_channel: &CH, + consumer: Consumer, published_count: u64, aggregation_pb: &ProgressBar, ) -> AppResult<(u64, u64, u64)> { info!("Starting results aggregation phase..."); - - consume_channel - .queue_declare( - &args.results_queue, - QueueDeclareOptions { - durable: true, - ..Default::default() - }, - FieldTable::default(), - ) - .await?; - - let consumer_tag = format!( - "producer-aggregator-{}-{}", - std::process::id(), - Utc::now().timestamp() - ); - - let consumer = consume_channel - .basic_consume( - &args.results_queue, - &consumer_tag, - BasicConsumeOptions::default(), - FieldTable::default(), - ) - .await?; - // Map deliveries into ProcessingOutcome values let mapped_stream = consumer.filter_map(|delivery_result| async { match delivery_result { - Ok(delivery) => match serde_json::from_slice::(&delivery.data) { - Ok(outcome) => { - let _ = delivery.ack(BasicAckOptions::default()).await; - Some(outcome) - } - Err(err) => { - warn!(error = %err, "Failed to deserialize outcome."); - None + Ok(delivery) => { + let s = std::str::from_utf8(&delivery.data).unwrap(); + warn!("{}", s); + match serde_json::from_slice::(&delivery.data) { + Ok(outcome) => { + let _ = delivery.ack(BasicAckOptions::default()).await; + Some(outcome) + } + Err(err) => { + warn!(error = %err, "Failed to deserialize outcome."); + None + } } - }, + } Err(err) => { error!(error = %err, "Failed to receive delivery."); None diff --git a/src/utils/common.rs b/src/utils/common.rs index dd910d4..bcdad7d 100644 --- a/src/utils/common.rs +++ b/src/utils/common.rs @@ -1,10 +1,16 @@ // src/utils/utils.rs -use lapin::{Connection, ConnectionProperties, Result as LapinResult}; +use lapin::{ + options::{BasicConsumeOptions, BasicQosOptions, QueueDeclareOptions}, + types::FieldTable, + Channel, Connection, ConnectionProperties, Consumer, Result as LapinResult, +}; use std::time::Duration; use tokio::time::sleep; use tracing::{error, info}; +use crate::error::{PipelineError, Result}; + // Helper function to connect to RabbitMQ with retry (already here) pub async fn connect_rabbitmq(addr: &str) -> LapinResult { let options = ConnectionProperties::default() @@ -33,3 +39,74 @@ pub async fn connect_rabbitmq(addr: &str) -> LapinResult { } } } + +pub async fn setup_channels_and_queues( + conn: &Connection, + publish_queue: &str, + consume_queue: &str, + prefetch_count: u16, + binary: String, // For now just string, but could be enum of worker | producer +) -> Result<(Channel, Consumer)> { + let consume_channel = conn.create_channel().await.map_err(|e| { + PipelineError::QueueError(format!( + "{} failed to create consume channel: {}", + binary, e + )) + })?; + let publish_channel = conn.create_channel().await.map_err(|e| { + PipelineError::QueueError(format!( + "{} failed to create publish channel: {}", + binary, e + )) + })?; + + publish_channel + .queue_declare( + publish_queue, + QueueDeclareOptions { + durable: true, + ..Default::default() + }, + FieldTable::default(), + ) + .await + .map_err(|e| { + PipelineError::QueueError(format!("{} failed to declare task queue: {}", binary, e)) + })?; + + consume_channel + .queue_declare( + consume_queue, + QueueDeclareOptions { + durable: true, + ..Default::default() + }, + FieldTable::default(), + ) + .await + .map_err(|e| { + PipelineError::QueueError(format!("{} failed to declare results queue: {}", binary, e)) + })?; + + consume_channel + .basic_qos(prefetch_count, BasicQosOptions::default()) + .await + .map_err(|e| PipelineError::QueueError(format!("Failed to set QoS: {}", e)))?; + + let consumer_tag = format!( + "{}-{}-{}", + binary, + std::process::id(), + chrono::Utc::now().timestamp() + ); + let consumer = consume_channel + .basic_consume( + consume_queue, + &consumer_tag, + BasicConsumeOptions::default(), + FieldTable::default(), + ) + .await?; + + Ok((publish_channel, consumer)) +} diff --git a/src/worker_logic.rs b/src/worker_logic.rs new file mode 100644 index 0000000..91bda40 --- /dev/null +++ b/src/worker_logic.rs @@ -0,0 +1,283 @@ +// src/worker_logic.rs + +use crate::config::pipeline::{PipelineConfig, StepConfig}; +use crate::error::{PipelineError, Result}; // Use the library's Result type +use crate::executor::{PipelineExecutor, ProcessingStep}; +use crate::utils::common::setup_channels_and_queues; +use futures::StreamExt; // For processing the consumer stream +use lapin::{Channel, Connection}; +// Import necessary filters (adjust if steps change) +// {{ Remove ArrowNativeType import if no longer needed directly here }} +// use arrow::datatypes::ArrowNativeType; +use crate::config::worker::Args; +use crate::pipeline::filters::{ + C4BadWordsFilter, + C4QualityFilter, + FineWebQualityFilter, + GopherQualityFilter, // Updated import + GopherRepetitionFilter, + LanguageDetectionFilter, +}; +use crate::pipeline::token::TokenCounter; + +use std::sync::Arc; // To share the executor across potential concurrent tasks + // {{ Add serde_json for result serialization }} +use tracing::{debug, error, info, info_span, instrument, warn}; // Added tracing + +use lapin::{ + message::Delivery, + options::{BasicAckOptions, BasicPublishOptions}, + protocol::basic::AMQPProperties, +}; + +use crate::data_model::{ProcessingOutcome, TextDocument}; +use crate::utils::prometheus_metrics::*; + +// {{ Add the new function to build pipeline from configuration }} +/// Builds the processing pipeline based on the configuration read from YAML. +#[instrument(skip(config), fields(num_steps = config.pipeline.len()))] +pub fn build_pipeline_from_config(config: &PipelineConfig) -> Result>> { + let mut steps: Vec> = Vec::new(); + info!("Building pipeline from configuration..."); + + for (i, step_config) in config.pipeline.iter().enumerate() { + let step_span = info_span!("pipeline_step", index = i, type = step_config.name()); + let _enter = step_span.enter(); + + let step: Box = match step_config { + StepConfig::C4QualityFilter(params) => { + debug!(params = ?params, "Adding C4QualityFilter"); + Box::new(C4QualityFilter::new( + params.split_paragraph, + params.remove_citations, + params.filter_no_terminal_punct, + params.min_num_sentences, + params.min_words_per_line, + params.max_word_length, + params.filter_lorem_ipsum, + params.filter_javascript, + params.filter_curly_bracket, + params.filter_policy, + )) + } + StepConfig::GopherRepetitionFilter(params) => { + debug!(params = ?params, "Adding GopherRepetitionFilter"); + Box::new(GopherRepetitionFilter::new( + params.dup_line_frac, + params.dup_para_frac, + params.dup_line_char_frac, + params.dup_para_char_frac, + params.top_n_grams.clone(), // Clone the vec + params.dup_n_grams.clone(), // Clone the vec + )) + } + StepConfig::GopherQualityFilter(params) => { + debug!(params = ?params, "Adding GopherQualityFilter"); + Box::new(GopherQualityFilter::new( + params.min_doc_words, + params.max_doc_words, + params.min_avg_word_length, + params.max_avg_word_length, + params.max_symbol_word_ratio, + params.max_bullet_lines_ratio, + params.max_ellipsis_lines_ratio, + params.max_non_alpha_words_ratio, + params.min_stop_words, + params.stop_words.clone(), // Clone the Option> + )) + } + StepConfig::LanguageDetectionFilter(params) => { + debug!(params = ?params, "Adding LanguageDetectionFilter"); + Box::new(LanguageDetectionFilter::new( + params.min_confidence, + params.allowed_languages.clone(), + )) + } + StepConfig::C4BadWordsFilter(params) => { + debug!(params = ?params, "Adding C4BadWordsFilter"); + Box::new(C4BadWordsFilter::new(params.clone())) + } + StepConfig::FineWebQualityFilter(params) => { + // Updated variant name + debug!(params = ?params, "Adding FineWebQualityFilter"); + + Box::new(FineWebQualityFilter::new( + params.line_punct_thr, + params.line_punct_exclude_zero, + params.short_line_thr, + params.short_line_length, + params.char_duplicates_ratio, + params.new_line_ratio, + // params.language.clone(), + params.stop_chars.clone(), + )) + } + StepConfig::TokenCounter(params) => { + debug!(params=?params, "Adding TokenCounter"); + let token_step = TokenCounter::new(¶ms.tokenizer_name); + if let Err(e) = token_step { + panic!("{}", e); + } + Box::new(token_step.unwrap()) + } + }; + steps.push(step); + info!("Added step: {}", step_config.name()); + } + + if steps.is_empty() { + warn!("Warning: Building an empty pipeline from configuration!"); + } else { + info!("Pipeline built successfully with {} steps.", steps.len()); + } + Ok(steps) +} + +/// Decoupled core processing logic for a single task. +/// +/// Returns `Some(ProcessingOutcome)` if the document was successfully processed +/// or filtered. Returns `None` if there was an unrecoverable error. +pub async fn execute_processing_pipeline( + data: &[u8], + executor: Arc, +) -> Option { + ACTIVE_PROCESSING_TASKS.inc(); + let processing_timer = TASK_PROCESSING_DURATION_SECONDS.start_timer(); + + let result = match serde_json::from_slice::(data) { + Ok(doc) => { + let original_doc_id = doc.id.clone(); + let task_span = info_span!("process_task", doc_id = %original_doc_id); + let _enter = task_span.enter(); + debug!("Processing document"); + + match executor.run_single_async(doc.clone()).await { + Ok(processed_doc) => { + debug!(processed_doc_id = %processed_doc.id, "Successfully processed document"); + TASKS_PROCESSED_TOTAL.inc(); + Some(ProcessingOutcome::Success(processed_doc)) + } + Err(PipelineError::StepError { step_name, source }) => match *source { + PipelineError::DocumentFiltered { document, reason } => { + debug!(filtered_doc_id = %document.id, %step_name, %reason, "Document was filtered"); + TASKS_FILTERED_TOTAL.inc(); + Some(ProcessingOutcome::Filtered { + document: *document, + reason, + }) + } + other => { + error!(%step_name, error = %other, "Pipeline step failed"); + TASKS_FAILED_TOTAL.inc(); + None + } + }, + Err(e) => { + error!(error = %e, "Unexpected pipeline error"); + TASKS_FAILED_TOTAL.inc(); + None + } + } + } + Err(e) => { + error!(error = %e, payload = %String::from_utf8_lossy(data), "Failed to deserialize task message"); + TASK_DESERIALIZATION_ERRORS_TOTAL.inc(); + None + } + }; + + processing_timer.observe_duration(); + ACTIVE_PROCESSING_TASKS.dec(); + + result +} + +pub async fn process_single_task( + delivery: Delivery, + executor: Arc, + publish_channel: Channel, + results_queue_name: &str, +) { + ACTIVE_PROCESSING_TASKS.inc(); + let processing_timer = TASK_PROCESSING_DURATION_SECONDS.start_timer(); + + let result: Option = + execute_processing_pipeline(&delivery.data, executor).await; + + if let Some(outcome) = result { + match serde_json::to_vec(&outcome) { + Ok(payload) => { + match publish_channel + .basic_publish( + "", + results_queue_name, + BasicPublishOptions::default(), + &payload, + AMQPProperties::default().with_delivery_mode(2), + ) + .await + { + Ok(_) => debug!("Published outcome"), + Err(e) => { + error!(error = %e, "Failed publish confirmation for outcome"); + OUTCOME_PUBLISH_ERRORS_TOTAL.inc(); + } + } + } + Err(e) => { + error!(error = %e, "Failed to serialize outcome"); + } + } + } + + if let Err(ack_err) = delivery.ack(BasicAckOptions::default()).await { + error!(error = %ack_err, "Failed to ack task message"); + } + + processing_timer.observe_duration(); + ACTIVE_PROCESSING_TASKS.dec(); +} + +pub async fn process_tasks_with_executor( + args: &Args, + conn: &Connection, + executor: Arc, +) -> Result<()> { + let (publish_channel, consumer) = setup_channels_and_queues( + conn, + &args.results_queue, + &args.task_queue, + args.prefetch_count, + "worker".to_string(), + ) + .await + .unwrap(); + + let mut consumer = consumer; + + while let Some(delivery_result) = consumer.next().await { + if let Ok(delivery) = delivery_result { + let executor_clone = Arc::clone(&executor); + let publish_channel_clone = publish_channel.clone(); + let results_queue_name = args.results_queue.clone(); + + tokio::spawn(async move { + process_single_task( + delivery, + executor_clone, + publish_channel_clone, + &results_queue_name, + ) + .await; + }); + } else { + error!("Error receiving task message. Worker stopping."); + return Err(PipelineError::QueueError( + "Consumer stream error".to_string(), + )); + } + } + + info!("Consumer stream ended."); + Ok(()) +} diff --git a/tests/full_pipeline_test.rs b/tests/full_pipeline_test.rs index 53931f4..673b265 100644 --- a/tests/full_pipeline_test.rs +++ b/tests/full_pipeline_test.rs @@ -210,7 +210,7 @@ async fn test_full_pipeline_e2e() -> Result<()> { .arg("--input-file") .arg(input_parquet_file.to_str().unwrap()) .arg("--text-column") - .arg("content") + .arg("text") .arg("--id-column") // Add this line .arg("id") // Add this line .arg("--amqp-addr") @@ -316,7 +316,7 @@ async fn test_full_pipeline_e2e() -> Result<()> { ); let output_reader_config = ParquetInputConfig { path: output_parquet_file.to_str().unwrap().to_string(), - text_column: "content".to_string(), + text_column: "text".to_string(), id_column: Some("id".to_string()), batch_size: Some(10), }; @@ -336,7 +336,7 @@ async fn test_full_pipeline_e2e() -> Result<()> { ); let excluded_reader_config = ParquetInputConfig { path: excluded_parquet_file.to_str().unwrap().to_string(), - text_column: "content".to_string(), + text_column: "text".to_string(), id_column: Some("id".to_string()), batch_size: Some(10), }; diff --git a/tests/parquet_io_test.rs b/tests/parquet_io_test.rs index 63e7859..f71b9fb 100644 --- a/tests/parquet_io_test.rs +++ b/tests/parquet_io_test.rs @@ -72,9 +72,9 @@ fn test_parquet_read_write_roundtrip() -> Result<()> { // and its fields are public. let reader_config = ParquetInputConfig { path: file_path_str.to_string(), - text_column: "content".to_string(), // Matches the field name in TextDocument and schema in ParquetWriter - id_column: Some("id".to_string()), // Matches the field name - batch_size: Some(10), // Optional, can be tested + text_column: "text".to_string(), // Matches the field name in TextDocument and schema in ParquetWriter + id_column: Some("id".to_string()), // Matches the field name + batch_size: Some(10), // Optional, can be tested }; let reader = ParquetReader::new(reader_config); diff --git a/tests/producer_tests.rs b/tests/producer_tests.rs index 767a51b..21dced1 100644 --- a/tests/producer_tests.rs +++ b/tests/producer_tests.rs @@ -98,314 +98,221 @@ mod args_tests { } #[cfg(test)] -mod publish_task_tests { - use arrow::array::StringArray; - use arrow::datatypes::{DataType, Field, Schema}; - use arrow::record_batch::RecordBatch; - use async_trait::async_trait; - use indicatif::ProgressBar; - use lapin::options::{BasicPublishOptions, ConfirmSelectOptions, QueueDeclareOptions}; - use lapin::protocol::basic::AMQPProperties; - use lapin::publisher_confirm::Confirmation; - use lapin::types::FieldTable; - use lapin::Result as LapinResult; - use parquet::arrow::arrow_writer::ArrowWriter; - use parquet::file::properties::WriterProperties; +mod publish_tasks_tests { use std::collections::HashMap; - use std::fs::File; - use std::sync::{Arc, Mutex}; + + use indicatif::ProgressBar; + use lapin::{options::*, types::FieldTable, Channel, Connection, ConnectionProperties}; use tempfile::NamedTempFile; + use testcontainers::{ + core::{IntoContainerPort, WaitFor}, + runners::AsyncRunner, + ContainerAsync, GenericImage, + }; + use tokio::time::{sleep, Duration, Instant}; + use uuid::Uuid; + // Added AsyncRunner use TextBlaster::config::producer::Args; use TextBlaster::data_model::TextDocument; - use TextBlaster::error::PipelineError; - use TextBlaster::producer_logic::*; - use TextBlaster::utils::prometheus_metrics::{ - TASKS_PUBLISHED_TOTAL, TASK_PUBLISH_ERRORS_TOTAL, - }; - - //=============== MOCK SETUP ===============// - - /// Defines the behavior of our mock publisher channel. - #[derive(Clone, Copy)] - enum MockBehavior { - /// Always return a successful ACK. - AlwaysAck, - /// Return a NACK on the first publish attempt. - NackOnFirstPublish, - /// Return a generic LapinError on publish. - FailOnPublish, - } - - /// A mock implementation of the TaskPublisherChannel trait. - /// It allows us to simulate RabbitMQ behavior without a real connection. - struct MockTaskPublisherChannel { - /// Shared state to inspect after the test runs. - state: Arc>, - } - - struct MockState { - /// Stores the payloads that were "published". - published_payloads: Vec>, - /// Controls how the mock responds to publish calls. - behavior: MockBehavior, - } - - impl MockTaskPublisherChannel { - fn new(behavior: MockBehavior) -> Self { - Self { - state: Arc::new(Mutex::new(MockState { - published_payloads: Vec::new(), - behavior, - })), - } - } - } - - #[async_trait] - impl TaskPublisherChannel for MockTaskPublisherChannel { - async fn queue_declare( - &self, - _name: &str, - _options: QueueDeclareOptions, - _arguments: FieldTable, - ) -> LapinResult<()> { - // Return a dummy queue. Its properties don't matter for this test. - Ok(()) - } - - async fn basic_publish( - &self, - _exchange: &str, - _routing_key: &str, - _options: BasicPublishOptions, - payload: &[u8], - _properties: AMQPProperties, - ) -> LapinResult { - let mut state = self.state.lock().unwrap(); - state.published_payloads.push(payload.to_vec()); - - match state.behavior { - MockBehavior::AlwaysAck => Ok(Confirmation::Ack(Default::default())), - MockBehavior::NackOnFirstPublish => Ok(Confirmation::Nack(Default::default())), - MockBehavior::FailOnPublish => { - // CORRECTED LINE: - // We must construct the full AMQPError with a code and text. - let amqp_error = lapin::protocol::AMQPError::new( - lapin::protocol::AMQPErrorKind::Hard( - lapin::protocol::AMQPHardError::INTERNALERROR, - ), - "mock failure".into(), - ); - Err(lapin::Error::ProtocolError(amqp_error)) - } - } - } - - async fn confirm_select(&self, _options: ConfirmSelectOptions) -> LapinResult<()> { - Ok(()) - } - } + use TextBlaster::error::Result; + use TextBlaster::pipeline::writers::parquet_writer::ParquetWriter; + use TextBlaster::producer_logic::publish_tasks; // replace `my_crate` with your actual crate - //=============== TEST HELPER FUNCTIONS ===============// - - /// Helper to create a temporary Parquet file with a specified number of documents. - /// Returns the temp file handle (to prevent deletion), the file path, and the original docs. - fn create_test_parquet_file(num_records: usize) -> (NamedTempFile, Vec) { - let temp_file = NamedTempFile::new().unwrap(); - let file_path = temp_file.path().to_str().unwrap().to_string(); - - let ids: Vec = (0..num_records).map(|i| format!("doc_{}", i)).collect(); - let texts: Vec = (0..num_records) - .map(|i| format!("This is text for doc {}.", i)) - .collect(); - - let original_docs: Vec = ids - .iter() - .zip(texts.iter()) - .map(|(id, text)| TextDocument { - id: id.clone(), - source: "test".to_string(), - content: text.clone(), - metadata: HashMap::new(), - }) - .collect(); - - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Utf8, false), - Field::new("text", DataType::Utf8, false), - ])); - - let id_array = StringArray::from_iter_values(ids.iter()); - let text_array = StringArray::from_iter_values(texts.iter()); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(id_array), Arc::new(text_array)], - ) - .unwrap(); - - let file = File::create(&file_path).unwrap(); - let mut writer = - ArrowWriter::try_new(file, schema, Some(WriterProperties::builder().build())).unwrap(); - writer.write(&batch).unwrap(); - writer.close().unwrap(); - - (temp_file, original_docs) - } - - fn create_mock_args(input_path: String) -> Args { + fn create_mock_args(input_path: String, queue_name: String) -> Args { Args { input_file: input_path, text_column: "text".to_string(), id_column: Some("id".to_string()), amqp_addr: "amqp://guest:guest@localhost:5672/%2f".to_string(), - task_queue: "test_task_queue".to_string(), + task_queue: queue_name, results_queue: "result_queue".to_string(), + prefetch_count: 10, output_file: "output".to_string(), excluded_file: "excluded".to_string(), metrics_port: Some(1234), } } - //=============== TEST CASES ===============// - - #[tokio::test] - async fn test_publish_tasks_happy_path() { - // ARRANGE - let num_docs = 5; - let (_temp_file, original_docs) = create_test_parquet_file(num_docs); - let args = create_mock_args(_temp_file.path().to_str().unwrap().to_string()); - let mock_channel = MockTaskPublisherChannel::new(MockBehavior::AlwaysAck); - let pb = ProgressBar::hidden(); + // Helper function to start a RabbitMQ container + async fn start_rabbitmq_container() -> (ContainerAsync, Channel, String) { + let image = GenericImage::new("rabbitmq", "3.13-management") + .with_wait_for(WaitFor::message_on_stdout( + "Server startup complete".to_string(), + )) + .with_exposed_port(5672.tcp()); // Default AMQP port + + // Use AsyncRunner for async test environments + let container = image + .start() + .await + .expect("Failed to start RabbitMQ container"); - // Reset metrics for a clean slate - TASKS_PUBLISHED_TOTAL.reset(); + let host_ip = container + .get_host() + .await + .expect("Failed to get container host IP"); + let host_port = container + .get_host_port_ipv4(5672) + .await + .expect("Failed to get mapped port"); - // ACT - let result = publish_tasks(&args, &mock_channel, &pb).await; + let amqp_addr = format!("amqp://guest:guest@{}:{}/%2f", host_ip, host_port); - // ASSERT - assert!(result.is_ok(), "Function should succeed"); - assert_eq!( - result.unwrap(), - num_docs as u64, - "Should report all documents as published" - ); + let conn = Connection::connect(&amqp_addr, ConnectionProperties::default()) + .await + .expect("connection failed"); + let channel = conn.create_channel().await.expect("channel failed"); - // Check metrics - assert_eq!( - TASKS_PUBLISHED_TOTAL.get(), - num_docs as f64, - "Prometheus metric for published tasks should be correct" - ); + let queue_name = format!("test_q_{}", Uuid::new_v4()); - // Check mock state - let state = mock_channel.state.lock().unwrap(); - assert_eq!( - state.published_payloads.len(), - num_docs, - "Exactly 5 messages should have been published" - ); + channel + .queue_declare( + &queue_name, + QueueDeclareOptions::default(), + FieldTable::default(), + ) + .await + .expect("queue declare failed"); - // Verify content of a published message - let first_payload = &state.published_payloads[0]; - let deserialized_doc: TextDocument = serde_json::from_slice(first_payload).unwrap(); - assert_eq!( - deserialized_doc.content, original_docs[0].content, - "The content of the published message should match the source document" - ); + (container, channel, queue_name) } - #[tokio::test] - async fn test_publish_tasks_stops_on_nack() { - // ARRANGE - let (_temp_file, _) = create_test_parquet_file(5); - let args = create_mock_args(_temp_file.path().to_str().unwrap().to_string()); - let mock_channel = MockTaskPublisherChannel::new(MockBehavior::NackOnFirstPublish); - let pb = ProgressBar::hidden(); + async fn fetch_message(channel: &Channel, queue: &str) -> lapin::message::BasicGetMessage { + let timeout = Duration::from_secs(3); + let start = Instant::now(); - // Reset metrics - TASK_PUBLISH_ERRORS_TOTAL.reset(); + loop { + if start.elapsed() > timeout { + panic!("Timed out waiting for message"); + } - // ACT - let result = publish_tasks(&args, &mock_channel, &pb).await; + let result = channel + .basic_get(queue, BasicGetOptions::default()) + .await + .expect("basic_get failed"); - // ASSERT - assert!(result.is_err(), "Function should fail on NACK"); - let err = result.unwrap_err(); - assert!( - matches!(err, PipelineError::QueueError(_)), - "Error should be of type QueueError" - ); - assert!( - err.to_string() - .contains("Publish confirmation failed (NACK)"), - "Error message should indicate a NACK" - ); + if let Some(delivery) = result { + return delivery; + } - // Check metrics - assert_eq!( - TASK_PUBLISH_ERRORS_TOTAL.get(), - 1.0, - "Prometheus metric for publish errors should be incremented" - ); + sleep(Duration::from_millis(100)).await; + } + } - // Check mock state: The message was still sent before the NACK was received. - let state = mock_channel.state.lock().unwrap(); - assert_eq!( - state.published_payloads.len(), - 1, - "Only one message should have been attempted before stopping" - ); + fn create_test_parquet_file(docs: &[TextDocument]) -> Result { + let file = NamedTempFile::new().unwrap(); + let mut writer = ParquetWriter::new(file.path().to_str().unwrap())?; + writer.write_batch(docs)?; + writer.close()?; + Ok(file) } + #[ignore] #[tokio::test] - async fn test_publish_tasks_propagates_lapin_error() { - // ARRANGE - let (_temp_file, _) = create_test_parquet_file(5); - let args = create_mock_args(_temp_file.path().to_str().unwrap().to_string()); - let mock_channel = MockTaskPublisherChannel::new(MockBehavior::FailOnPublish); - let pb = ProgressBar::hidden(); + async fn test_publish_tasks_single_document() -> Result<()> { + let (_container, channel, queue_name) = start_rabbitmq_container().await; - // ACT - let result = publish_tasks(&args, &mock_channel, &pb).await; + let doc = TextDocument { + id: "doc-1".into(), + source: "test".into(), + content: "Simple content".into(), + metadata: [("lang".into(), "en".into())].into(), + }; + let parquet: NamedTempFile = create_test_parquet_file(&[doc])?; - // ASSERT - assert!( - result.is_err(), - "Function should fail if basic_publish returns an error" - ); - let err = result.unwrap_err(); - assert!( - matches!(err, PipelineError::QueueError(_)), - "Error should be a wrapped QueueError" + let args = create_mock_args( + parquet.path().to_str().unwrap().to_string(), + queue_name.clone(), ); + + let pb = ProgressBar::hidden(); + let result = publish_tasks(&args, &channel, &pb).await?; + assert_eq!(result, 1); + + let delivery = fetch_message(&channel, &queue_name).await; + let value: serde_json::Value = serde_json::from_slice(&delivery.data)?; + assert_eq!(value["id"], "doc-1"); + assert_eq!(value["content"], "Simple content"); + Ok(()) } #[tokio::test] - async fn test_publish_tasks_handles_nonexistent_input_file() { - // ARRANGE - let args = create_mock_args("does_not_exist".to_string()); - // The mock won't even be used, as the failure happens before publishing. - let mock_channel = MockTaskPublisherChannel::new(MockBehavior::AlwaysAck); + #[ignore] + async fn test_publish_tasks_multiple_documents() -> Result<()> { + let (_container, channel, queue_name) = start_rabbitmq_container().await; + + let docs = vec![ + TextDocument { + id: "a".into(), + source: "s".into(), + content: "1".into(), + metadata: HashMap::new(), + }, + TextDocument { + id: "b".into(), + source: "s".into(), + content: "2".into(), + metadata: HashMap::new(), + }, + TextDocument { + id: "c".into(), + source: "s".into(), + content: "3".into(), + metadata: HashMap::new(), + }, + ]; + let parquet = create_test_parquet_file(&docs)?; + + let args = create_mock_args( + parquet.path().to_str().unwrap().to_string(), + queue_name.clone(), + ); + let pb = ProgressBar::hidden(); + let result = publish_tasks(&args, &channel, &pb).await?; + assert_eq!(result, 3); + + let mut seen_ids = vec![]; + for _ in 0..3 { + let d = fetch_message(&channel, &queue_name).await; + let val: serde_json::Value = serde_json::from_slice(&d.data)?; + seen_ids.push(val["id"].as_str().unwrap().to_string()); + } - // ACT - // Note: The error here is synchronous, as it happens during ParquetReader setup, - // but publish_tasks wraps it in the AppResult. - let result = publish_tasks(&args, &mock_channel, &pb).await; + assert_eq!(seen_ids.len(), 3); + assert!(seen_ids.contains(&"a".to_string())); + assert!(seen_ids.contains(&"b".to_string())); + assert!(seen_ids.contains(&"c".to_string())); + Ok(()) + } - // ASSERT - assert!( - result.is_err(), - "Function should fail if input file doesn't exist" - ); - let err = result.unwrap_err(); - assert!( - matches!(err, PipelineError::IoError { source: _ }), - "{}", - format!("Error should be of type ParquetError not {}", err) + #[tokio::test] + #[ignore] + async fn test_publish_tasks_empty_metadata() -> Result<()> { + let (_container, channel, queue_name) = start_rabbitmq_container().await; + + let doc = TextDocument { + id: "empty-meta".into(), + source: "unit".into(), + content: "Testing empty metadata".into(), + metadata: HashMap::new(), + }; + + let parquet = create_test_parquet_file(&[doc])?; + + let args = create_mock_args( + parquet.path().to_str().unwrap().to_string(), + queue_name.clone(), ); + + let pb = ProgressBar::hidden(); + let result = publish_tasks(&args, &channel, &pb).await?; + assert_eq!(result, 1); + + let delivery = fetch_message(&channel, &queue_name).await; + let val: serde_json::Value = serde_json::from_slice(&delivery.data)?; + assert_eq!(val["id"], "empty-meta"); + Ok(()) } } @@ -427,6 +334,7 @@ mod aggregate_results_tests { amqp_addr: "amqp://guest:guest@localhost:5672/%2f".to_string(), task_queue: "test_task_queue".to_string(), results_queue: "result_queue".to_string(), + prefetch_count: 10, output_file: output_path, excluded_file: excluded_path, metrics_port: Some(1234), diff --git a/tests/worker_tests.rs b/tests/worker_tests.rs new file mode 100644 index 0000000..bd58118 --- /dev/null +++ b/tests/worker_tests.rs @@ -0,0 +1,156 @@ +// tests/worker_logic_tests.rs + +use std::collections::HashMap; +use std::sync::Arc; + +use axum::async_trait; +use TextBlaster::data_model::{ProcessingOutcome, TextDocument}; +use TextBlaster::error::{PipelineError, Result}; +use TextBlaster::executor::{PipelineExecutor, ProcessingStep}; +use TextBlaster::worker_logic::execute_processing_pipeline; + +struct FailingStep; + +#[async_trait] +impl ProcessingStep for FailingStep { + fn name(&self) -> &'static str { + "FailingStep" + } + + async fn process(&self, document: TextDocument) -> Result { + Err(PipelineError::StepError { + step_name: "FailingStep".to_string(), + source: Box::new(PipelineError::DocumentFiltered { + document: Box::new(document), + reason: "FailingStep".to_string(), + }), + }) + } +} + +struct FilteringStep; + +#[async_trait] +impl ProcessingStep for FilteringStep { + fn name(&self) -> &'static str { + "FilteringStep" + } + + async fn process(&self, doc: TextDocument) -> Result { + Err(PipelineError::DocumentFiltered { + document: Box::new(doc), + reason: "Blocked by test filter".to_string(), + }) + } +} + +struct IdentityStep; + +#[async_trait] +impl ProcessingStep for IdentityStep { + fn name(&self) -> &'static str { + "IdentityStep" + } + + async fn process(&self, doc: TextDocument) -> Result { + Ok(doc) + } +} + +#[tokio::test] +async fn test_success_path() { + let steps: Vec> = vec![Box::new(IdentityStep)]; + let executor = Arc::new(PipelineExecutor::new(steps)); + + let input_doc = TextDocument { + id: "doc-1".into(), + content: "This is a test".into(), + source: "test-suite".into(), + metadata: HashMap::new(), + }; + + let raw = serde_json::to_vec(&input_doc).unwrap(); + let outcome = execute_processing_pipeline(&raw, executor).await; + + match outcome.unwrap() { + ProcessingOutcome::Success(doc) => assert_eq!(doc.id, "doc-1"), + _ => panic!("Expected success outcome"), + } +} + +#[tokio::test] +async fn test_deserialization_failure() { + let steps: Vec> = vec![Box::new(IdentityStep)]; + let executor = Arc::new(PipelineExecutor::new(steps)); + + let raw = b"not valid json".to_vec(); + let result = execute_processing_pipeline(&raw, executor).await; + assert!(result.is_none()); +} + +#[tokio::test] +async fn test_filtered_outcome() { + let steps: Vec> = vec![Box::new(FilteringStep)]; + let executor = Arc::new(PipelineExecutor::new(steps)); + + let input_doc = TextDocument { + id: "doc-filtered".into(), + content: "Block me".into(), + source: "filter-test".into(), + metadata: HashMap::new(), + }; + + let raw = serde_json::to_vec(&input_doc).unwrap(); + let outcome = execute_processing_pipeline(&raw, executor).await; + + match outcome.unwrap() { + ProcessingOutcome::Filtered { document, reason } => { + assert_eq!(document.id, "doc-filtered"); + assert_eq!(reason, "Blocked by test filter"); + } + _ => panic!("Expected filtered outcome"), + } +} + +#[tokio::test] +async fn test_failing_step_returns_none() { + let steps: Vec> = vec![Box::new(FailingStep)]; + let executor = Arc::new(PipelineExecutor::new(steps)); + + let input_doc = TextDocument { + id: "doc-fail".into(), + content: "Cause failure".into(), + source: "fail-test".into(), + metadata: HashMap::new(), + }; + + let raw = serde_json::to_vec(&input_doc).unwrap(); + let outcome = execute_processing_pipeline(&raw, executor).await; + + assert!(outcome.is_none()); +} + +#[tokio::test] +async fn test_chained_steps_success() { + let steps: Vec> = vec![ + Box::new(IdentityStep), + Box::new(IdentityStep), + Box::new(IdentityStep), + ]; + let executor = Arc::new(PipelineExecutor::new(steps)); + + let input_doc = TextDocument { + id: "doc-chain".into(), + content: "Chain of steps".into(), + source: "chained".into(), + metadata: HashMap::new(), + }; + + let raw = serde_json::to_vec(&input_doc).unwrap(); + let outcome = execute_processing_pipeline(&raw, executor).await; + + match outcome.unwrap() { + ProcessingOutcome::Success(doc) => assert_eq!(doc.id, "doc-chain"), + _ => panic!("Expected success outcome"), + } +}