refactor: write and resolve should return the same type and Orchestrator to orchestrate retries (#3190)

Signed-off-by: Vigith Maurice <vigith@gmail.com>

Author: vigith

rust/numaflow-core/src/pipeline/isb.rs

@@ -6,7 +6,6 @@
 use std::time::Duration;
 
 use async_trait::async_trait;
-use tokio_util::sync::CancellationToken;
 
 use crate::error::Result;
 use crate::message::{Message, Offset};
@@ -92,28 +91,28 @@ impl std::fmt::Display for WriteError {
 
 impl std::error::Error for WriteError {}
 
-/// Result of resolving a pending write operation.
+/// Result of a write operation.
 ///
 /// Contains the offset of the written message along with additional metadata
 /// that may be useful for logging, metrics, or debugging.
 #[derive(Debug, Clone)]
-pub struct ResolveResult {
+pub struct WriteResult {
     /// The offset of the written message
     pub offset: Offset,
     /// Whether this was a duplicate message (already existed in the buffer)
     pub is_duplicate: bool,
 }
 
-impl ResolveResult {
-    /// Creates a new ResolveResult with the given offset and no duplicate flag.
+impl WriteResult {
+    /// Creates a new WriteResult with the given offset and no duplicate flag.
     pub fn new(offset: Offset) -> Self {
         Self {
             offset,
             is_duplicate: false,
         }
     }
 
-    /// Creates a new ResolveResult marked as a duplicate.
+    /// Creates a new WriteResult marked as a duplicate.
     pub fn duplicate(offset: Offset) -> Self {
         Self {
             offset,
@@ -128,7 +127,11 @@ impl ResolveResult {
 /// 1. **High-performance async pattern**: Use `async_write()` to get a `PendingWrite` handle
 ///    immediately, then resolve it later with `resolve()`. This allows batching acknowledgments
 ///    for higher throughput.
-/// 2. **Simple blocking pattern**: Use `blocking_write()` which blocks until the write is confirmed.
+/// 2. **Simple synchronous pattern**: Use `write()` which writes and waits for confirmation
+///    in a single call. This is useful as a fallback when async writes fail.
+///
+/// Both `resolve()` and `write()` return `Result<WriteResult, WriteError>` for consistency.
+/// The orchestrator is responsible for retry logic and handling cancellation.
 ///
 /// Implementations must be cheaply cloneable (e.g., using Arc internally).
 #[async_trait]
@@ -156,7 +159,7 @@ pub(crate) trait ISBWriter: Send + Sync + Clone {
 
     /// Resolves a pending write to get the result.
     ///
-    /// This waits for the write acknowledgment and returns a `ResolveResult` containing
+    /// This waits for the write acknowledgment and returns a `WriteResult` containing
     /// the offset of the written message along with additional metadata (e.g., whether
     /// the message was a duplicate).
     ///
@@ -165,21 +168,19 @@ pub(crate) trait ISBWriter: Send + Sync + Clone {
     async fn resolve(
         &self,
         pending: Self::PendingWrite,
-    ) -> std::result::Result<ResolveResult, WriteError>;
+    ) -> std::result::Result<WriteResult, WriteError>;
 
-    /// Writes a message and blocks until confirmed, returning the offset.
+    /// Writes a message and waits for confirmation, returning the result.
+    ///
+    /// This is a single-attempt write that returns immediately after the write
+    /// completes or fails. The orchestrator is responsible for retry logic.
     ///
-    /// This is the simple blocking write method with infinite retries until success
-    /// or cancellation. Use this when you don't need the high-performance async pattern.
+    /// Returns `Err(WriteError::BufferFull)` if the buffer is full.
+    /// Returns `Err(WriteError::WriteFailed)` if the write operation fails.
     ///
     /// # Arguments
     /// * `message` - The message to write
-    /// * `cln_token` - Cancellation token for graceful shutdown
-    async fn blocking_write(
-        &self,
-        message: Message,
-        cln_token: CancellationToken,
-    ) -> std::result::Result<Offset, WriteError>;
+    async fn write(&self, message: Message) -> std::result::Result<WriteResult, WriteError>;
 
     /// Returns the name/identifier of this writer (e.g., stream name).
     #[allow(dead_code)]

rust/numaflow-core/src/pipeline/isb/jetstream/js_writer.rs

@@ -6,10 +6,9 @@ use async_nats::jetstream::Context;
 use async_nats::jetstream::consumer::PullConsumer;
 use async_nats::jetstream::context::{PublishAckFuture, PublishErrorKind};
 use async_nats::jetstream::message::PublishMessage;
-use async_nats::jetstream::publish::PublishAck;
 use async_nats::jetstream::stream::RetentionPolicy::Limits;
 use bytes::{Bytes, BytesMut};
-use tokio::time::{Instant, sleep};
+use tokio::time::Instant;
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error};
 
@@ -44,7 +43,6 @@ struct BufferInfo {
     num_ack_pending: usize,
 }
 
-const DEFAULT_RETRY_INTERVAL_MILLIS: u64 = 10;
 const DEFAULT_REFRESH_INTERVAL_SECS: u64 = 1;
 
 /// Lightweight JetStream Writer for a single stream.
@@ -212,64 +210,82 @@ impl JetStreamWriter {
         }
     }
 
-    /// Writes the message to the JetStream ISB and returns the PublishAck. It will do
-    /// infinite retries until the message gets published successfully. If it returns
-    /// an error it means it is fatal non-retryable error.
-    pub(crate) async fn blocking_write(
+    /// Writes a message to the JetStream ISB and waits for confirmation.
+    ///
+    /// This is a single-attempt write that returns immediately after the write
+    /// completes or fails. The orchestrator is responsible for retry logic.
+    ///
+    /// Returns `Err(WriteError::BufferFull)` if the buffer is full.
+    /// Returns `Err(WriteError::WriteFailed)` if the write operation fails.
+    pub(crate) async fn write(
         &self,
         message: Message,
-        cln_token: CancellationToken,
-    ) -> Result<PublishAck> {
+    ) -> std::result::Result<crate::pipeline::isb::WriteResult, WriteError> {
         let start_time = Instant::now();
 
+        // Check if buffer is full
+        if self.is_full() {
+            pipeline_metrics()
+                .jetstream_isb
+                .isfull_total
+                .get_or_create(&self.buffer_labels)
+                .inc();
+            return Err(WriteError::BufferFull);
+        }
+
         // Compress the message value if compression is enabled
         let mut message = message;
-        message.value = match self.compression_type {
-            Some(compression_type) => bytes::Bytes::from(
-                compression::compress(compression_type, &message.value).map_err(|e| {
-                    Error::ISB(ISBError::Encode(format!("Compression failed: {}", e)))
-                })?,
-            ),
-            None => message.value,
-        };
+        if let Some(compression_type) = self.compression_type {
+            message.value = bytes::Bytes::from(
+                compression::compress(compression_type, &message.value)
+                    .map_err(|e| WriteError::WriteFailed(format!("Compression failed: {}", e)))?,
+            );
+        }
 
         let payload: Bytes = message
             .try_into()
             .expect("message serialization should not fail");
 
-        loop {
-            match self.js_ctx.publish(self.stream.name, payload.clone()).await {
-                Ok(paf) => match paf.await {
-                    Ok(ack) => {
-                        debug!(
-                            elapsed_ms = start_time.elapsed().as_millis(),
-                            "Blocking write successful in",
-                        );
-                        pipeline_metrics()
-                            .jetstream_isb
-                            .write_time_total
-                            .get_or_create(&self.buffer_labels)
-                            .observe(start_time.elapsed().as_micros() as f64);
-                        return Ok(ack);
-                    }
-                    Err(e) => {
-                        error!(?e, "awaiting publish ack failed, retrying");
-                        sleep(Duration::from_millis(10)).await;
-                    }
-                },
-                Err(e) => {
+        // Publish and await acknowledgment
+        match self.js_ctx.publish(self.stream.name, payload).await {
+            Ok(paf) => match paf.await {
+                Ok(ack) => {
+                    debug!(
+                        elapsed_ms = start_time.elapsed().as_millis(),
+                        "Write successful",
+                    );
                     pipeline_metrics()
                         .jetstream_isb
-                        .write_error_total
+                        .write_time_total
                         .get_or_create(&self.buffer_labels)
-                        .inc();
-                    error!(?e, "publishing failed, retrying");
-                    sleep(Duration::from_millis(DEFAULT_RETRY_INTERVAL_MILLIS)).await;
+                        .observe(start_time.elapsed().as_micros() as f64);
+
+                    // Convert sequence number to Offset using the stream's partition
+                    let offset = crate::message::Offset::Int(crate::message::IntOffset::new(
+                        ack.sequence as i64,
+                        self.stream.partition,
+                    ));
+
+                    // Check if this was a duplicate message
+                    if ack.duplicate {
+                        Ok(crate::pipeline::isb::WriteResult::duplicate(offset))
+                    } else {
+                        Ok(crate::pipeline::isb::WriteResult::new(offset))
+                    }
                 }
-            }
-
-            if cln_token.is_cancelled() {
-                return Err(Error::Cancelled());
+                Err(e) => {
+                    error!(?e, "awaiting publish ack failed");
+                    Err(WriteError::WriteFailed(e.to_string()))
+                }
+            },
+            Err(e) => {
+                pipeline_metrics()
+                    .jetstream_isb
+                    .write_error_total
+                    .get_or_create(&self.buffer_labels)
+                    .inc();
+                error!(?e, "publishing failed");
+                Err(WriteError::WriteFailed(e.to_string()))
             }
         }
     }
@@ -395,7 +411,7 @@ impl crate::pipeline::isb::ISBWriter for JetStreamWriter {
     async fn resolve(
         &self,
         pending: Self::PendingWrite,
-    ) -> std::result::Result<crate::pipeline::isb::ResolveResult, crate::pipeline::isb::WriteError>
+    ) -> std::result::Result<crate::pipeline::isb::WriteResult, crate::pipeline::isb::WriteError>
     {
         // Await the PAF to get the PublishAck
         let ack = pending
@@ -410,28 +426,22 @@ impl crate::pipeline::isb::ISBWriter for JetStreamWriter {
 
         // Check if this was a duplicate message
         if ack.duplicate {
-            Ok(crate::pipeline::isb::ResolveResult::duplicate(offset))
+            Ok(crate::pipeline::isb::WriteResult::duplicate(offset))
         } else {
-            Ok(crate::pipeline::isb::ResolveResult::new(offset))
+            Ok(crate::pipeline::isb::WriteResult::new(offset))
         }
     }
 
-    async fn blocking_write(
+    async fn write(
         &self,
         message: Message,
-        cln_token: CancellationToken,
-    ) -> std::result::Result<crate::message::Offset, crate::pipeline::isb::WriteError> {
-        // Use the existing blocking_write which handles retries internally
-        let ack = self
-            .blocking_write(message, cln_token)
-            .await
-            .map_err(|e| crate::pipeline::isb::WriteError::WriteFailed(e.to_string()))?;
-
-        // Convert sequence number to Offset using the stream's partition
-        Ok(crate::message::Offset::Int(crate::message::IntOffset::new(
-            ack.sequence as i64,
-            self.stream.partition,
-        )))
+    ) -> std::result::Result<crate::pipeline::isb::WriteResult, crate::pipeline::isb::WriteError>
+    {
+        // Delegate to the inherent write method
+        self.write(message).await.map_err(|e| match e {
+            WriteError::BufferFull => crate::pipeline::isb::WriteError::BufferFull,
+            WriteError::WriteFailed(msg) => crate::pipeline::isb::WriteError::WriteFailed(msg),
+        })
     }
 
     fn name(&self) -> &'static str {
@@ -688,13 +698,13 @@ mod tests {
 
     #[cfg(feature = "nats-tests")]
     #[tokio::test]
-    async fn test_blocking_write_with_compression() {
+    async fn test_write_with_compression() {
         let js_url = "localhost:4222";
         let client = async_nats::connect(js_url).await.unwrap();
         let context = jetstream::new(client);
         let cln_token = CancellationToken::new();
 
-        let stream = Stream::new("test-blocking-compress", "temp", 0);
+        let stream = Stream::new("test-write-compress", "temp", 0);
         let _ = context.delete_stream(stream.name).await;
         let _stream = context
             .get_or_create_stream(stream::Config {
@@ -752,11 +762,12 @@ mod tests {
             ..Default::default()
         };
 
-        let result = writer.blocking_write(message, cln_token).await;
-        assert!(
-            result.is_ok(),
-            "blocking_write with compression should succeed"
-        );
+        let result = writer.write(message).await;
+        assert!(result.is_ok(), "write with compression should succeed");
+
+        // Verify the WriteResult has the expected structure
+        let write_result = result.unwrap();
+        assert!(!write_result.is_duplicate, "should not be a duplicate");
 
         context.delete_stream(stream.name).await.unwrap();
     }