chore: code review

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

Author: vigith

rust/Cargo.lock

@@ -18,7 +18,7 @@ all-tests:
 	cargo test --features nats-tests,pulsar-tests,sqs-tests,kafka-tests,redis-tests --workspace --all
 	cargo test --features global-state-tests test_map_with_panic
 	cargo test --features global-state-tests test_batch_map_with_panic
-	RUST_LOG=debug cargo test --features global-state-tests test_map_stream_with_panic -- --nocapture
+	cargo test --features global-state-tests test_map_stream_with_panic
 	cargo test --features global-state-tests test_transform_stream_with_panic
 
 .PHONY: clean

rust/Makefile

@@ -75,6 +75,5 @@ k8s-openapi = { version = "0.24.0", features = ["v1_32"] }
 numaflow-kafka = { workspace = true, features = ["kafka-tests-utils"] }
 numaflow-pulsar = { workspace = true, features = ["pulsar-tests-utils"] }
 serial_test = "3.3.1"
-tracing-subscriber = "0.3.20"
 
 [build-dependencies]

rust/numaflow-core/Cargo.toml

@@ -44,11 +44,11 @@ enum ConcurrentMapper {
 }
 
 /// MapperType is an enum to store the different types of mappers.
-/// Concurrent mappers (Unary/Stream) process messages individually with concurrency control.
-/// Batch mapper processes messages in batches sequentially.
 #[derive(Clone)]
 enum MapperType {
+    /// Concurrent mappers (Unary/Stream) process messages individually with concurrency control.
     Concurrent(ConcurrentMapper),
+    /// Batch mapper processes messages in batches sequentially.
     Batch(UserDefinedBatchMap),
 }
 
@@ -205,7 +205,7 @@ impl MapHandle {
                 .acquire_many_owned(self.concurrency as u32)
                 .await
                 .map_err(|e| Error::Mapper(format!("failed to acquire semaphore: {e}")))?;
-            info!(status=?self.final_result, "Map component is completed with status");
+            info!(status=?self.final_result, "Map component is completed");
 
             // abort the shutdown handle since we are done processing, no need to wait for the
             // hard shutdown
@@ -409,13 +409,6 @@ fn update_udf_process_time_metric(is_mono_vertex: bool) {
     }
 }
 
-/*
- t1 -> receiver dropped
- t2 -> map drained
- t3 -> written to the tx
- t4 -> hung
-*/
-
 fn update_udf_write_only_metric(is_mono_vertex: bool) {
     if !is_mono_vertex {
         pipeline_metrics()
@@ -1055,10 +1048,9 @@ mod tests {
         }
     }
 
-    // #[cfg(feature = "global-state-tests")]
+    #[cfg(feature = "global-state-tests")]
     #[tokio::test]
     async fn test_batch_map_with_panic() -> Result<()> {
-        tracing_subscriber::fmt::init();
         let cln_token = CancellationToken::new();
         let (_shutdown_tx, shutdown_rx) = oneshot::channel();
         let tmp_dir = TempDir::new().unwrap();

rust/numaflow-core/src/mapper/map.rs

@@ -23,12 +23,19 @@ use tracing::error;
 /// Type alias for the batch response - raw results from the UDF
 pub(in crate::mapper) type BatchMapResponse = Vec<map::map_response::Result>;
 
-/// Type aliases
+/// Type aliases for HashMap used to track the oneshot response sender for each request keyed by
+/// message id.
 type ResponseSenderMap = HashMap<String, oneshot::Sender<Result<BatchMapResponse>>>;
 
+/// Shared state for tracking batch map senders between the sender and the receiver tasks.
+/// We have BiDi gRPC stream so we have 2 different set of tasks for sending and receiving.
 #[derive(Default)]
 pub(in crate::mapper) struct BatchSenderMapState {
+    /// Map of oneshot response senders keyed by message id.
     map: ResponseSenderMap,
+    /// Flag to indicate whether the rx task has closed the stream and cleared the `map`.
+    /// This is because `tx.send()` could return `Ok()` even after the receiver task has closed the
+    /// stream.
     closed: bool,
 }
 

rust/numaflow-core/src/mapper/map/batch.rs

@@ -13,7 +13,7 @@ use tokio_util::sync::CancellationToken;
 use tokio_util::task::AbortOnDropHandle;
 use tonic::Streaming;
 use tonic::transport::Channel;
-use tracing::error;
+use tracing::{error, warn};
 
 use super::{
     ParentMessageInfo, STREAMING_MAP_RESP_CHANNEL_SIZE, SharedMapTaskContext, UserDefinedMessage,
@@ -24,11 +24,19 @@ use super::{
 /// Type alias for the stream response - raw results from the UDF
 pub(in crate::mapper) type StreamMapResponse = Vec<map::map_response::Result>;
 
+/// Type aliases for HashMap used to track the oneshot response sender for each request keyed by
+/// message id.
 type StreamResponseSenderMap = HashMap<String, mpsc::Sender<Result<StreamMapResponse>>>;
 
+/// Shared state for tracking batch map senders between the sender and the receiver tasks.
+/// We have BiDi gRPC stream so we have 2 different set of tasks for sending and receiving.
 #[derive(Default)]
 pub(in crate::mapper) struct StreamSenderMapState {
+    /// Map of oneshot response senders keyed by message id.
     map: StreamResponseSenderMap,
+    /// Flag to indicate whether the rx task has closed the stream and cleared the `map`.
+    /// This is because `tx.send()` could return `Ok()` even after the receiver task has closed the
+    /// stream.
     closed: bool,
 }
 
@@ -261,7 +269,8 @@ impl UserDefinedStreamMap {
                 .send(Err(Error::Mapper(format!(
                     "failed to send message to stream map server: {e}"
                 ))))
-                .await;
+                .await
+                .inspect(|_| warn!("failed to send error to oneshot receiver"));
             return rx;
         }
 

rust/numaflow-core/src/mapper/map/stream.rs

@@ -13,7 +13,7 @@ use tokio_util::sync::CancellationToken;
 use tokio_util::task::AbortOnDropHandle;
 use tonic::Streaming;
 use tonic::transport::Channel;
-use tracing::error;
+use tracing::{error, warn};
 
 use super::{
     ParentMessageInfo, SharedMapTaskContext, UserDefinedMessage, create_response_stream,
@@ -23,24 +23,34 @@ use super::{
 /// Type alias for the response - raw results from the UDF
 pub(in crate::mapper) type UnaryMapResponse = Vec<map::map_response::Result>;
 
+/// Type aliases for HashMap used to track the oneshot response sender for each request keyed by
+/// message id.
 type ResponseSenderMap = HashMap<String, oneshot::Sender<Result<UnaryMapResponse>>>;
 
+/// Shared state for tracking batch map senders between the sender and the receiver tasks.
+/// We have BiDi gRPC stream so we have 2 different set of tasks for sending and receiving.
 #[derive(Default)]
 pub(in crate::mapper) struct UnarySenderMapState {
+    /// Map of oneshot response senders keyed by message id.
     map: ResponseSenderMap,
+    /// Flag to indicate whether the rx task has closed the stream and cleared the `map`.
+    /// This is because `tx.send()` could return `Ok()` even after the receiver task has closed the
+    /// stream.
     closed: bool,
 }
 
-/// MapUnaryTask encapsulates all the context needed to execute a unary map operation.
+/// MapUnaryTask encapsulates all the context needed to execute a unary map operation per message.
 pub(in crate::mapper) struct MapUnaryTask {
     pub mapper: UserDefinedUnaryMap,
+    /// Permit to achieve structured concurrency by ensuring we do not exceed the concurrency limit
+    /// and all the tasks are cleaned up when the component is shutting down.
     pub permit: OwnedSemaphorePermit,
     pub message: Message,
     pub shared_ctx: Arc<SharedMapTaskContext>,
 }
 
 impl MapUnaryTask {
-    /// Spawns the unary map task as a tokio task.
+    /// Spawns the unary map task as a tokio task per [MapUnaryTask].
     /// The task will process the message through the UDF and send results downstream.
     pub fn spawn(self) {
         tokio::spawn(async move {
@@ -229,7 +239,9 @@ impl UserDefinedUnaryMap {
             if !senders_guard.closed {
                 senders_guard.map.insert(key.clone(), tx);
             } else {
-                let _ = tx.send(Err(Error::Mapper("mapper closed".to_string())));
+                let _ = tx
+                    .send(Err(Error::Mapper("mapper closed".to_string())))
+                    .inspect(|_| warn!("failed to send error to oneshot receiver"));
             }
         };