feat(connectors): add source batch acknowledgments

Cargo.toml

@@ -203,7 +203,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"

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<String, Arc<Container<SourceApi>>> = 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<SourceConnectorPlugin>,
 }
 

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(),
         );
 

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<ProducedMessages>,
+    pub(crate) sender: Sender<ProducedBatch>,
     // Owned errors counter (Arc<AtomicU64> 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<DashMap<u32, SourceSenderEntry>> =
     LazyLock::new(DashMap::new);
 
@@ -364,7 +371,8 @@ pub(crate) async fn source_forwarding_loop(
     encoder: Arc<dyn StreamEncoder>,
     transforms: Vec<Arc<dyn Transform>>,
     state_storage: StateStorage,
-    receiver: Receiver<ProducedMessages>,
+    receiver: Receiver<ProducedBatch>,
+    batch_result_callback: BatchResultCallback,
     context: Arc<RuntimeContext>,
     labels: Arc<SourceLabels>,
 ) {
@@ -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<u64> = 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<dyn StreamEncoder>,
     transforms: Vec<Arc<dyn Transform>>,
     state_storage: StateStorage,
-    callback: HandleCallback,
+    handle_callback: HandleCallback,
+    batch_result_callback: BatchResultCallback,
     context: Arc<RuntimeContext>,
 ) -> Vec<JoinHandle<()>> {
     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::<ProducedMessages>(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
+        );
+    }
+}

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"

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).