From e0e7a02b72b21bf8b90935af661491b07c6ce40e Mon Sep 17 00:00:00 2001 From: Rohan Dubey Date: Mon, 10 Aug 2026 21:23:15 +0530 Subject: [PATCH] feat(connectors): add source batch acknowledgments --- Cargo.lock | 2 +- Cargo.toml | 2 +- core/connectors/runtime/src/main.rs | 12 +- core/connectors/runtime/src/manager/source.rs | 6 +- core/connectors/runtime/src/source.rs | 150 +++++- core/connectors/sdk/Cargo.toml | 2 +- core/connectors/sdk/README.md | 27 ++ core/connectors/sdk/src/lib.rs | 11 + core/connectors/sdk/src/source.rs | 445 +++++++++++++++++- 9 files changed, 627 insertions(+), 30 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a56f39bb0f..1f0b6ac92b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7190,7 +7190,7 @@ dependencies = [ [[package]] name = "iggy_connector_sdk" -version = "0.3.1-edge.1" +version = "0.4.0-edge.1" dependencies = [ "anyhow", "apache-avro", diff --git a/Cargo.toml b/Cargo.toml index f33c0eddc7..da05e4ba61 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -204,7 +204,7 @@ iggy = { path = "core/sdk", version = "0.11.0-edge.1" } iggy-cli = { path = "core/cli", version = "0.14.0-edge.1" } iggy_binary_protocol = { path = "core/binary_protocol", version = "0.11.0-edge.1" } iggy_common = { path = "core/common", version = "0.11.0-edge.1" } -iggy_connector_sdk = { path = "core/connectors/sdk", version = "0.3.1-edge.1" } +iggy_connector_sdk = { path = "core/connectors/sdk", version = "0.4.0-edge.1" } indexmap = "2.14.0" integration = { path = "core/integration" } ipnet = "2.12.0" diff --git a/core/connectors/runtime/src/main.rs b/core/connectors/runtime/src/main.rs index 5c5ebe7774..6ff2c472f1 100644 --- a/core/connectors/runtime/src/main.rs +++ b/core/connectors/runtime/src/main.rs @@ -29,7 +29,7 @@ use iggy_connector_sdk::{ StreamDecoder, StreamEncoder, api::ConnectorStatus, sink::ConsumeCallback, - source::{HandleCallback, SendCallback}, + source::{BatchResultCallback, HandleCallback, SendCallback}, transforms::Transform, }; use mimalloc::MiMalloc; @@ -81,6 +81,7 @@ pub(crate) struct SourceApi { log_callback: iggy_connector_sdk::LogCallback, ) -> i32, iggy_source_handle: extern "C" fn(id: u32, callback: SendCallback) -> i32, + iggy_source_batch_result: extern "C" fn(plugin_id: u32, batch_id: u64, result: u8) -> i32, iggy_source_close: extern "C" fn(id: u32) -> i32, iggy_source_version: extern "C" fn() -> *const std::ffi::c_char, } @@ -185,12 +186,14 @@ async fn main() -> Result<(), RuntimeError> { let mut source_containers_by_key: HashMap>> = HashMap::new(); for (_path, source) in sources { let container = Arc::new(source.container); - let callback = container.iggy_source_handle; + let handle_callback = container.iggy_source_handle; + let batch_result_callback = container.iggy_source_batch_result; for plugin in &source.plugins { source_containers_by_key.insert(plugin.key.clone(), container.clone()); } source_wrappers.push(SourceConnectorWrapper { - callback, + handle_callback, + batch_result_callback, plugins: source.plugins, }); } @@ -460,7 +463,8 @@ struct SourceConnectorProducer { } struct SourceConnectorWrapper { - callback: HandleCallback, + handle_callback: HandleCallback, + batch_result_callback: BatchResultCallback, plugins: Vec, } diff --git a/core/connectors/runtime/src/manager/source.rs b/core/connectors/runtime/src/manager/source.rs index b011472fe7..2da39c9baa 100644 --- a/core/connectors/runtime/src/manager/source.rs +++ b/core/connectors/runtime/src/manager/source.rs @@ -223,7 +223,8 @@ impl SourceManager { let (producer, encoder, transforms) = source::setup_source_producer(key, config, iggy_client).await?; - let callback = container.iggy_source_handle; + let handle_callback = container.iggy_source_handle; + let batch_result_callback = container.iggy_source_batch_result; let handler_tasks = source::spawn_source_handler( plugin_id, key, @@ -233,7 +234,8 @@ impl SourceManager { encoder, transforms, state_storage, - callback, + handle_callback, + batch_result_callback, context.clone(), ); diff --git a/core/connectors/runtime/src/source.rs b/core/connectors/runtime/src/source.rs index e259e212cc..1b896db97e 100644 --- a/core/connectors/runtime/src/source.rs +++ b/core/connectors/runtime/src/source.rs @@ -25,7 +25,8 @@ use iggy::prelude::{ use iggy_connector_sdk::encoders::avro::{AvroEncoderConfig, AvroStreamEncoder}; use iggy_connector_sdk::{ ConnectorState, DecodedMessage, ProducedMessages, Schema, StreamEncoder, TopicMetadata, - source::HandleCallback, transforms::Transform, + source::{BatchResultCallback, HandleCallback, SourceBatchResult}, + transforms::Transform, }; use std::{ collections::{BTreeMap, HashMap}, @@ -51,12 +52,18 @@ use prometheus_client::metrics::counter::Counter; use tokio::task::JoinHandle; pub(crate) struct SourceSenderEntry { - pub(crate) sender: Sender, + pub(crate) sender: Sender, // Owned errors counter (Arc inside) so the FFI callback bumps // it with one relaxed atomic - no Family RwLock + HashMap lookup per call. pub(crate) error_counter: Counter, } +#[derive(Debug)] +pub(crate) struct ProducedBatch { + id: u64, + messages: ProducedMessages, +} + pub(crate) static SOURCE_SENDERS: LazyLock> = LazyLock::new(DashMap::new); @@ -364,7 +371,8 @@ pub(crate) async fn source_forwarding_loop( encoder: Arc, transforms: Vec>, state_storage: StateStorage, - receiver: Receiver, + receiver: Receiver, + batch_result_callback: BatchResultCallback, context: Arc, labels: Arc, ) { @@ -390,8 +398,10 @@ pub(crate) async fn source_forwarding_loop( topic: producer.topic().to_string(), }; - while let Ok(produced_messages) = receiver.recv_async().await { + while let Ok(produced_batch) = receiver.recv_async().await { let total_start = Instant::now(); + let batch_id = produced_batch.id; + let produced_messages = produced_batch.messages; let count = produced_messages.messages.len(); context .metrics @@ -461,6 +471,7 @@ pub(crate) async fn source_forwarding_loop( // Total histogram + emit (below) run regardless of send outcome. let mut state_save_us: Option = None; + let mut batch_result = SourceBatchResult::Nack; if let Err(error) = send_result { let error_msg = format!( "Failed to send {sent_count} messages to stream: {}, topic: {} by source connector with ID: {plugin_id}. {error}", @@ -489,11 +500,13 @@ pub(crate) async fn source_forwarding_loop( ); } + let mut state_saved = true; if let Some(state) = produced_messages.state { let state_save_start = Instant::now(); match &state_storage { StateStorage::File(file) => { if let Err(error) = file.save(state).await { + state_saved = false; let error_msg = format!( "Failed to save state for source connector with ID: {plugin_id}. {error}" ); @@ -514,6 +527,20 @@ pub(crate) async fn source_forwarding_loop( } else { debug!("No state provided for source connector with ID: {plugin_id}"); } + + if state_saved { + batch_result = SourceBatchResult::Ack; + } + } + + let result_code = batch_result_callback(plugin_id, batch_id, batch_result as u8); + if result_code != 0 { + let error_msg = format!( + "Failed to deliver {batch_result:?} for source connector with ID: {plugin_id}, batch ID: {batch_id}. Plugin returned: {result_code}" + ); + error!("{error_msg}"); + context.metrics.inc_errors_with_labels(&labels.counter); + context.sources.set_error(&plugin_key, &error_msg).await; } let total_elapsed = total_start.elapsed(); @@ -558,7 +585,8 @@ pub(crate) fn spawn_source_handler( encoder: Arc, transforms: Vec>, state_storage: StateStorage, - callback: HandleCallback, + handle_callback: HandleCallback, + batch_result_callback: BatchResultCallback, context: Arc, ) -> Vec> { let (sender, receiver) = flume::unbounded(); @@ -573,7 +601,7 @@ pub(crate) fn spawn_source_handler( ); let blocking_handle = tokio::task::spawn_blocking(move || { - callback(plugin_id, handle_produced_messages); + handle_callback(plugin_id, handle_produced_messages); }); let handler_task = tokio::spawn(async move { source_forwarding_loop( @@ -586,6 +614,7 @@ pub(crate) fn spawn_source_handler( transforms, state_storage, receiver, + batch_result_callback, context, labels, ) @@ -627,7 +656,8 @@ pub fn handle( producer_wrapper.encoder, plugin.transforms, plugin.state_storage, - source.callback, + source.handle_callback, + source.batch_result_callback, context.clone(), ); @@ -713,9 +743,10 @@ fn process_messages( pub(crate) extern "C" fn handle_produced_messages( plugin_id: u32, + batch_id: u64, messages_ptr: *const u8, messages_len: usize, -) { +) -> i32 { unsafe { // Entry missing = SOURCE_SENDERS cleaned up at shutdown; benign race // expected on stop/restart. No metric (would conflate with real failures). @@ -724,23 +755,29 @@ pub(crate) extern "C" fn handle_produced_messages( plugin_id, "dropping produced batch: sender already cleaned up" ); - return; + return -1; }; let messages = std::slice::from_raw_parts(messages_ptr, messages_len); match postcard::from_bytes::(messages) { Ok(messages) => { - if let Err(send_error) = entry.sender.send(messages) { + if let Err(send_error) = entry.sender.send(ProducedBatch { + id: batch_id, + messages, + }) { error!( "Failed to send messages for source connector with ID: {plugin_id}. Channel closed: {send_error}" ); entry.error_counter.inc(); + return -1; } + 0 } Err(err) => { error!( "Failed to deserialize produced messages for source connector with ID: {plugin_id}. {err}" ); entry.error_counter.inc(); + -1 } } } @@ -768,3 +805,96 @@ fn build_iggy_message( (None, None) => IggyMessage::builder().payload(payload.into()).build(), } } + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicU32, Ordering}; + + static TEST_PLUGIN_ID: AtomicU32 = AtomicU32::new(u32::MAX / 2); + + fn next_plugin_id() -> u32 { + TEST_PLUGIN_ID.fetch_add(1, Ordering::Relaxed) + } + + #[test] + fn given_serialized_batch_when_callback_runs_should_forward_batch_id() { + let plugin_id = next_plugin_id(); + let batch_id = 73; + let (sender, receiver) = flume::unbounded(); + SOURCE_SENDERS.insert( + plugin_id, + SourceSenderEntry { + sender, + error_counter: Counter::default(), + }, + ); + let messages = ProducedMessages { + schema: Schema::Raw, + messages: Vec::new(), + state: Some(ConnectorState(vec![1, 2, 3])), + }; + let serialized = postcard::to_allocvec(&messages).expect("failed to serialize batch"); + + assert_eq!( + handle_produced_messages(plugin_id, batch_id, serialized.as_ptr(), serialized.len()), + 0 + ); + let forwarded = receiver.recv().expect("batch was not forwarded"); + assert_eq!(forwarded.id, batch_id); + assert_eq!( + forwarded + .messages + .state + .expect("state should be preserved") + .0, + vec![1, 2, 3] + ); + + cleanup_sender(plugin_id); + } + + #[test] + fn given_invalid_payload_when_callback_runs_should_reject_batch() { + let plugin_id = next_plugin_id(); + let (sender, _receiver) = flume::unbounded(); + let error_counter = Counter::default(); + SOURCE_SENDERS.insert( + plugin_id, + SourceSenderEntry { + sender, + error_counter: error_counter.clone(), + }, + ); + let invalid_payload = [0xff]; + + assert_eq!( + handle_produced_messages( + plugin_id, + 1, + invalid_payload.as_ptr(), + invalid_payload.len(), + ), + -1 + ); + assert_eq!(error_counter.get(), 1); + + cleanup_sender(plugin_id); + } + + #[test] + fn given_missing_sender_when_callback_runs_should_reject_batch() { + let plugin_id = next_plugin_id(); + let serialized = postcard::to_allocvec(&ProducedMessages { + schema: Schema::Raw, + messages: Vec::new(), + state: None, + }) + .expect("failed to serialize batch"); + + assert_eq!( + handle_produced_messages(plugin_id, 1, serialized.as_ptr(), serialized.len()), + -1 + ); + } +} diff --git a/core/connectors/sdk/Cargo.toml b/core/connectors/sdk/Cargo.toml index fcd2b6f2ae..74b6e7d098 100644 --- a/core/connectors/sdk/Cargo.toml +++ b/core/connectors/sdk/Cargo.toml @@ -17,7 +17,7 @@ [package] name = "iggy_connector_sdk" -version = "0.3.1-edge.1" +version = "0.4.0-edge.1" description = "Iggy is the persistent message streaming platform written in Rust, supporting QUIC, TCP and HTTP transport protocols, capable of processing millions of messages per second." edition = "2024" license = "Apache-2.0" diff --git a/core/connectors/sdk/README.md b/core/connectors/sdk/README.md index fb85d055b3..422bd697c1 100644 --- a/core/connectors/sdk/README.md +++ b/core/connectors/sdk/README.md @@ -4,6 +4,33 @@ SDK provides the commonly used structs and traits such as `Sink` and `Source`, a The macros automatically export the connector's version (from `CARGO_PKG_VERSION`) via FFI, allowing the runtime to report per-connector version information in the `/stats` endpoint. +## Source delivery acknowledgment + +Source connectors use a one-in-flight-batch contract between the plugin and the runtime: + +1. `Source::poll()` returns messages and candidate state without committing cursor changes or destructive operations. +2. The runtime sends the batch to Iggy and waits for the producer result. +3. After a successful send, the runtime persists the candidate state. +4. The runtime reports `SourceBatchResult::Ack` to the plugin. A send or state-save failure reports `SourceBatchResult::Nack` instead. +5. `Source::on_batch_result()` commits or discards the plugin's staged work before the next poll starts. + +An empty batch follows the same handshake. This prevents a successful no-op send from persisting state left over from an earlier failed delivery. Producer errors, including request timeouts, report a NACK. A successful send from the legacy Iggy server is still an ACK even though that server returns an empty confirmation list. + +The crash behavior is intentionally at-least-once: + +| Crash point | Recovery behavior | +| --- | --- | +| Before Iggy commits the batch | Persisted state is unchanged and the source can poll the batch again. | +| After Iggy commits but before the runtime observes success | Persisted state is unchanged, so the batch may be delivered again. | +| After send success but before state persistence | Persisted state is unchanged, so the batch may be delivered again. | +| After state persistence but before the plugin processes the ACK | The restored state records the delivered batch. Deferred source-side cleanup may still be pending. | +| After the plugin processes the ACK | The state and plugin cursor both record the delivered batch. | + +Source-side ACK work should be idempotent because process termination can interrupt it. NACK handling must discard staged cursor changes and staged delete or mark operations so polling can redeliver the batch. +The SDK stops polling if `Source::on_batch_result()` returns an error, preventing a failed rollback from advancing to another batch. + +This contract is a breaking FFI change. Source plugins must be rebuilt with the matching SDK. `iggy_source_handle` now supplies a batch ID to the runtime callback, and source plugins export `iggy_source_batch_result` for the corresponding ACK or NACK. + Moreover, it contains both, the `decoders` and `encoders` modules, implementing either `StreamDecoder` or `StreamEncoder` traits, which are used when consuming or producing data from/to Iggy streams. SDK is WiP, and it'd certainly benefit from having the support of multiple format schemas, such as Protobuf, Avro, Flatbuffers etc. including decoding/encoding the data between the different formats (when applicable) and supporting the data transformations whenever possible (easy for JSON, but complex for Bincode for example). diff --git a/core/connectors/sdk/src/lib.rs b/core/connectors/sdk/src/lib.rs index c8ed2ff94d..94eeb28289 100644 --- a/core/connectors/sdk/src/lib.rs +++ b/core/connectors/sdk/src/lib.rs @@ -111,6 +111,17 @@ pub trait Source: Send + Sync { /// Invoked every time a batch of messages is produced to the configured stream and topic. async fn poll(&self) -> Result; + /// Invoked after the runtime has finished processing the most recently polled batch. + /// + /// Sources that track cursors or perform destructive operations should stage those changes + /// in [`Source::poll`] and apply them only after receiving [`source::SourceBatchResult::Ack`]. + /// A [`source::SourceBatchResult::Nack`] means the staged changes must be discarded so the + /// batch can be polled again. The SDK allows only one batch to be in flight at a time and + /// stops polling if this method returns an error. + async fn on_batch_result(&self, _result: source::SourceBatchResult) -> Result<(), Error> { + Ok(()) + } + /// Invoked when the source is closed, allowing it to perform any necessary cleanup. async fn close(&mut self) -> Result<(), Error>; } diff --git a/core/connectors/sdk/src/source.rs b/core/connectors/sdk/src/source.rs index 713e00c6ae..028c5f7edb 100644 --- a/core/connectors/sdk/src/source.rs +++ b/core/connectors/sdk/src/source.rs @@ -16,10 +16,13 @@ // under the License. use crate::log::{CallbackLayer, LogCallback}; -use crate::{ConnectorState, Error, Source, get_runtime}; +use crate::{ConnectorState, Source, get_runtime}; use serde::de::DeserializeOwned; -use std::sync::Arc; -use tokio::{sync::watch, task::JoinHandle}; +use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; +use tokio::{ + sync::{oneshot, watch}, + task::JoinHandle, +}; use tracing::{error, info}; use tracing_subscriber::{EnvFilter, Registry, layer::SubscriberExt, util::SubscriberInitExt}; @@ -34,7 +37,50 @@ pub struct RawMessage { pub type HandleCallback = extern "C" fn(plugin_id: u32, callback: SendCallback) -> i32; -pub type SendCallback = extern "C" fn(plugin_id: u32, messages_ptr: *const u8, messages_len: usize); +pub type SendCallback = extern "C" fn( + plugin_id: u32, + batch_id: u64, + messages_ptr: *const u8, + messages_len: usize, +) -> i32; + +pub type BatchResultCallback = extern "C" fn(plugin_id: u32, batch_id: u64, result: u8) -> i32; + +/// Delivery result for the single batch currently in flight from a source plugin. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +pub enum SourceBatchResult { + /// The runtime sent the complete batch and persisted its candidate state. + Ack = 0, + /// The runtime could not send the batch or persist its candidate state. + Nack = 1, +} + +impl TryFrom for SourceBatchResult { + type Error = (); + + fn try_from(value: u8) -> Result { + match value { + 0 => Ok(Self::Ack), + 1 => Ok(Self::Nack), + _ => Err(()), + } + } +} + +struct PendingBatch { + id: u64, + result_sender: oneshot::Sender, +} + +impl std::fmt::Debug for PendingBatch { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("PendingBatch") + .field("id", &self.id) + .finish() + } +} #[derive(Debug)] pub struct SourceContainer { @@ -42,15 +88,17 @@ pub struct SourceContainer { source: Option>, shutdown: Option>, task: Option>, + pending_batch: Arc>>, } impl SourceContainer { - pub const fn new(id: u32) -> Self { + pub fn new(id: u32) -> Self { Self { id, source: None, shutdown: None, task: None, + pending_batch: Arc::new(Mutex::new(None)), } } @@ -157,22 +205,50 @@ impl SourceContainer { let (shutdown_tx, shutdown_rx) = watch::channel(()); let plugin_id = self.id; let source = Arc::clone(source); + let pending_batch = Arc::clone(&self.pending_batch); let handle = runtime.spawn(async move { - let _ = handle_messages(plugin_id, source, callback, shutdown_rx).await; + handle_messages( + plugin_id, + source, + move |plugin_id, batch_id, messages_ptr, messages_len| { + callback(plugin_id, batch_id, messages_ptr, messages_len) + }, + shutdown_rx, + pending_batch, + ) + .await; }); self.shutdown = Some(shutdown_tx); self.task = Some(handle); 0 } + + #[doc(hidden)] + pub fn complete_batch(&self, batch_id: u64, result: u8) -> i32 { + let Ok(result) = SourceBatchResult::try_from(result) else { + error!( + "Invalid batch result: {result} for source connector with ID: {}", + self.id + ); + return -1; + }; + + complete_pending_batch(&self.pending_batch, batch_id, result, self.id) + } } -async fn handle_messages( +async fn handle_messages( plugin_id: u32, source: Arc, - callback: SendCallback, + callback: F, mut shutdown: watch::Receiver<()>, -) -> Result<(), Error> { + pending_batch: Arc>>, +) where + T: Source, + F: Fn(u32, u64, *const u8, usize) -> i32, +{ + let mut batch_id = 1u64; loop { tokio::select! { _ = shutdown.changed() => { @@ -192,16 +268,111 @@ async fn handle_messages( Ok(messages) => messages, Err(err) => { error!("Failed to serialize messages for source connector with ID: {plugin_id}. {err}"); + if !notify_source(&source, SourceBatchResult::Nack, plugin_id).await { + break; + } continue; } }; - callback(plugin_id, messages.as_ptr(), messages.len()); + let (result_sender, result_receiver) = oneshot::channel(); + { + let mut pending = lock_pending_batch(&pending_batch); + *pending = Some(PendingBatch { + id: batch_id, + result_sender, + }); + } + + if callback(plugin_id, batch_id, messages.as_ptr(), messages.len()) != 0 { + _ = complete_pending_batch( + &pending_batch, + batch_id, + SourceBatchResult::Nack, + plugin_id, + ); + } + + let (result, shutting_down) = tokio::select! { + biased; + result = result_receiver => { + (result.unwrap_or(SourceBatchResult::Nack), false) + }, + _ = shutdown.changed() => { + _ = complete_pending_batch( + &pending_batch, + batch_id, + SourceBatchResult::Nack, + plugin_id, + ); + (SourceBatchResult::Nack, true) + } + }; + if !notify_source(&source, result, plugin_id).await { + break; + } + + if shutting_down { + info!("Shutting down source connector with ID: {plugin_id}"); + break; + } + + batch_id = batch_id.wrapping_add(1); + if batch_id == 0 { + batch_id = 1; + } } } } +} + +fn complete_pending_batch( + pending_batch: &Mutex>, + batch_id: u64, + result: SourceBatchResult, + plugin_id: u32, +) -> i32 { + let mut pending = lock_pending_batch(pending_batch); + let Some(current) = pending.as_ref() else { + error!("No batch is awaiting a result for source connector with ID: {plugin_id}"); + return -1; + }; + if current.id != batch_id { + error!( + "Batch result ID mismatch for source connector with ID: {plugin_id}. Expected: {}, received: {batch_id}", + current.id + ); + return -1; + } + + let Some(current) = pending.take() else { + return -1; + }; + if current.result_sender.send(result).is_err() { + error!( + "Failed to deliver batch result for source connector with ID: {plugin_id}, batch ID: {batch_id}" + ); + return -1; + } + 0 +} - Ok(()) +fn lock_pending_batch( + pending_batch: &Mutex>, +) -> MutexGuard<'_, Option> { + pending_batch.lock().unwrap_or_else(PoisonError::into_inner) +} + +async fn notify_source( + source: &Arc, + result: SourceBatchResult, + plugin_id: u32, +) -> bool { + if let Err(err) = source.on_batch_result(result).await { + error!("Failed to process {result:?} for source connector with ID: {plugin_id}. {err}"); + return false; + } + true } #[macro_export] @@ -264,6 +435,18 @@ macro_rules! source_connector { instance.handle(callback) } + #[cfg(not(test))] + #[unsafe(no_mangle)] + extern "C" fn iggy_source_batch_result(id: u32, batch_id: u64, result: u8) -> i32 { + let Some(instance) = INSTANCES.get(&id) else { + tracing::error!( + "Source connector with ID: {id} was not found and cannot complete batch {batch_id}." + ); + return -1; + }; + instance.complete_batch(batch_id, result) + } + #[cfg(not(test))] #[unsafe(no_mangle)] unsafe extern "C" fn iggy_source_close(id: u32) -> i32 { @@ -284,3 +467,243 @@ macro_rules! source_connector { } }; } + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ProducedMessages, Schema}; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use std::time::Duration; + use tokio::sync::mpsc; + + #[derive(Debug, Default)] + struct TestSource { + polls: AtomicUsize, + results: Mutex>, + fail_batch_result: AtomicBool, + } + + #[async_trait::async_trait] + impl Source for TestSource { + async fn open(&mut self) -> Result<(), crate::Error> { + Ok(()) + } + + async fn poll(&self) -> Result { + self.polls.fetch_add(1, Ordering::SeqCst); + Ok(ProducedMessages { + schema: Schema::Raw, + messages: Vec::new(), + state: None, + }) + } + + async fn on_batch_result(&self, result: SourceBatchResult) -> Result<(), crate::Error> { + self.results + .lock() + .unwrap_or_else(PoisonError::into_inner) + .push(result); + if self.fail_batch_result.load(Ordering::SeqCst) { + return Err(crate::Error::Storage( + "failed to apply batch result".to_string(), + )); + } + Ok(()) + } + + async fn close(&mut self) -> Result<(), crate::Error> { + Ok(()) + } + } + + #[test] + fn given_batch_without_result_should_not_poll_again() { + let runtime = tokio::runtime::Runtime::new().expect("failed to create test runtime"); + runtime.block_on(async { + let source = Arc::new(TestSource::default()); + let pending_batch = Arc::new(Mutex::new(None)); + let pending_for_task = Arc::clone(&pending_batch); + let (shutdown_sender, shutdown_receiver) = watch::channel(()); + let (batch_sender, mut batch_receiver) = mpsc::unbounded_channel(); + + let source_for_task = Arc::clone(&source); + let task = tokio::spawn(handle_messages( + 7, + source_for_task, + move |_, batch_id, _, _| { + batch_sender + .send(batch_id) + .expect("batch receiver should remain open"); + 0 + }, + shutdown_receiver, + pending_for_task, + )); + + let batch_id = tokio::time::timeout(Duration::from_secs(1), batch_receiver.recv()) + .await + .expect("first batch was not sent") + .expect("batch channel closed"); + assert_eq!(batch_id, 1); + assert_eq!(source.polls.load(Ordering::SeqCst), 1); + assert!( + tokio::time::timeout(Duration::from_millis(50), batch_receiver.recv()) + .await + .is_err(), + "source polled again before the first batch was completed" + ); + + assert_eq!( + complete_pending_batch(&pending_batch, batch_id, SourceBatchResult::Ack, 7), + 0 + ); + let next_batch_id = tokio::time::timeout(Duration::from_secs(1), batch_receiver.recv()) + .await + .expect("source did not poll after ACK") + .expect("batch channel closed"); + assert_eq!(next_batch_id, 2); + + shutdown_sender + .send(()) + .expect("source task should remain active"); + task.await.expect("source task failed"); + assert_eq!( + *source + .results + .lock() + .unwrap_or_else(PoisonError::into_inner), + vec![SourceBatchResult::Ack, SourceBatchResult::Nack] + ); + }); + } + + #[test] + fn given_nack_when_batch_is_pending_should_allow_redelivery() { + let runtime = tokio::runtime::Runtime::new().expect("failed to create test runtime"); + runtime.block_on(async { + let source = Arc::new(TestSource::default()); + let pending_batch = Arc::new(Mutex::new(None)); + let pending_for_task = Arc::clone(&pending_batch); + let (shutdown_sender, shutdown_receiver) = watch::channel(()); + let (batch_sender, mut batch_receiver) = mpsc::unbounded_channel(); + + let source_for_task = Arc::clone(&source); + let task = tokio::spawn(handle_messages( + 9, + source_for_task, + move |_, batch_id, _, _| { + batch_sender + .send(batch_id) + .expect("batch receiver should remain open"); + 0 + }, + shutdown_receiver, + pending_for_task, + )); + + let batch_id = tokio::time::timeout(Duration::from_secs(1), batch_receiver.recv()) + .await + .expect("first batch was not sent") + .expect("batch channel closed"); + assert_eq!( + complete_pending_batch(&pending_batch, batch_id, SourceBatchResult::Nack, 9), + 0 + ); + let next_batch_id = tokio::time::timeout(Duration::from_secs(1), batch_receiver.recv()) + .await + .expect("source did not poll after NACK") + .expect("batch channel closed"); + assert_eq!(next_batch_id, 2); + + shutdown_sender + .send(()) + .expect("source task should remain active"); + task.await.expect("source task failed"); + assert_eq!( + *source + .results + .lock() + .unwrap_or_else(PoisonError::into_inner), + vec![SourceBatchResult::Nack, SourceBatchResult::Nack] + ); + }); + } + + #[test] + fn given_mismatched_batch_id_should_reject_result() { + let pending_batch = Mutex::new(None); + let (result_sender, result_receiver) = oneshot::channel(); + *lock_pending_batch(&pending_batch) = Some(PendingBatch { + id: 41, + result_sender, + }); + + assert_eq!( + complete_pending_batch(&pending_batch, 42, SourceBatchResult::Ack, 11), + -1 + ); + assert_eq!( + complete_pending_batch(&pending_batch, 41, SourceBatchResult::Ack, 11), + 0 + ); + + let runtime = tokio::runtime::Runtime::new().expect("failed to create test runtime"); + assert_eq!( + runtime + .block_on(result_receiver) + .expect("batch result sender was dropped"), + SourceBatchResult::Ack + ); + } + + #[test] + fn given_batch_result_handler_failure_should_stop_polling() { + let runtime = tokio::runtime::Runtime::new().expect("failed to create test runtime"); + runtime.block_on(async { + let source = Arc::new(TestSource { + fail_batch_result: AtomicBool::new(true), + ..TestSource::default() + }); + let pending_batch = Arc::new(Mutex::new(None)); + let pending_for_task = Arc::clone(&pending_batch); + let (_shutdown_sender, shutdown_receiver) = watch::channel(()); + let (batch_sender, mut batch_receiver) = mpsc::unbounded_channel(); + + let source_for_task = Arc::clone(&source); + let task = tokio::spawn(handle_messages( + 13, + source_for_task, + move |_, batch_id, _, _| { + batch_sender + .send(batch_id) + .expect("batch receiver should remain open"); + 0 + }, + shutdown_receiver, + pending_for_task, + )); + + let batch_id = tokio::time::timeout(Duration::from_secs(1), batch_receiver.recv()) + .await + .expect("first batch was not sent") + .expect("batch channel closed"); + assert_eq!( + complete_pending_batch(&pending_batch, batch_id, SourceBatchResult::Nack, 13), + 0 + ); + tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("source task did not stop after batch result failure") + .expect("source task failed"); + + assert_eq!(source.polls.load(Ordering::SeqCst), 1); + assert_eq!( + *source + .results + .lock() + .unwrap_or_else(PoisonError::into_inner), + vec![SourceBatchResult::Nack] + ); + }); + } +}