chore: group serve-store and builtin serve-sink (#2582)

Signed-off-by: Vigith Maurice <[email protected]>

Author: vigith

rust/numaflow-core/src/lib.rs

@@ -61,9 +61,6 @@ mod tracker;
 /// [Map]: https://numaflow.numaproj.io/user-guide/user-defined-functions/map/map/
 mod mapper;
 
-/// Serving store to store the result of the serving pipeline.
-mod serving_store;
-
 /// [Watermark] _is a monotonically increasing timestamp of the oldest work/event not yet completed_
 ///
 ///

rust/numaflow-core/src/pipeline.rs

@@ -18,11 +18,11 @@ use crate::pipeline::forwarder::source_forwarder;
 use crate::pipeline::isb::jetstream::reader::JetStreamReader;
 use crate::pipeline::isb::jetstream::writer::JetstreamWriter;
 use crate::pipeline::pipeline::isb::BufferReaderConfig;
-use crate::serving_store::ServingStore;
-use crate::serving_store::nats::NatsServingStore;
-use crate::serving_store::user_defined::UserDefinedStore;
 use crate::shared::create_components;
 use crate::shared::metrics::start_metrics_server;
+use crate::sink::serve::ServingStore;
+use crate::sink::serve::nats::NatsServingStore;
+use crate::sink::serve::user_defined::UserDefinedStore;
 use crate::tracker::TrackerHandle;
 use crate::watermark::WatermarkHandle;
 use crate::watermark::isb::ISBWatermarkHandle;

rust/numaflow-core/src/serving_store.rs

@@ -14,9 +14,9 @@ use crate::config::pipeline::map::{MapMode, MapType, MapVtxConfig};
 use crate::config::pipeline::{DEFAULT_BATCH_MAP_SOCKET, DEFAULT_STREAM_MAP_SOCKET};
 use crate::error::Error;
 use crate::mapper::map::MapHandle;
-use crate::serving_store::ServingStore;
 use crate::shared::grpc;
 use crate::shared::server_info::{ContainerType, sdk_server_info};
+use crate::sink::serve::ServingStore;
 use crate::sink::{SinkClientType, SinkWriter, SinkWriterBuilder};
 use crate::source::Source;
 use crate::source::generator::new_generator;

rust/numaflow-core/src/shared/create_components.rs

@@ -31,7 +31,7 @@ use crate::metrics::{
     monovertex_metrics, mvtx_forward_metric_labels, pipeline_forward_metric_labels,
     pipeline_metrics,
 };
-use crate::serving_store::{ServingStore, StoreEntry};
+use crate::sink::serve::{ServingStore, StoreEntry};
 use crate::tracker::TrackerHandle;
 
 /// A [Blackhole] sink which reads but never writes to anywhere, semantic equivalent of `/dev/null`.
@@ -44,7 +44,9 @@ mod blackhole;
 /// [Log]: https://numaflow.numaproj.io/user-guide/sinks/log/
 mod log;
 
-mod serve;
+/// Serving [ServingStore] to store the result of the serving pipeline. It also contains the builtin [serve::ServeSink]
+/// to write to the serving store.
+pub mod serve;
 
 /// [User-Defined Sink] extends Numaflow to add custom sources supported outside the builtins.
 ///
@@ -958,12 +960,11 @@ mod tests {
     use super::*;
     use crate::config::pipeline::NatsStoreConfig;
     use crate::message::{IntOffset, Message, MessageID, Offset, ReadAck};
-    use crate::serving_store::nats::NatsServingStore;
     use crate::shared::grpc::create_rpc_channel;
+    use crate::sink::serve::nats::NatsServingStore;
     use async_nats::jetstream;
     use async_nats::jetstream::kv::Config;
     use chrono::{TimeZone, Utc};
-    use futures::StreamExt;
     use numaflow::sink;
     use numaflow_pb::clients::sink::{SinkRequest, SinkResponse};
     use std::sync::Arc;

rust/numaflow-core/src/sink.rs

@@ -1,6 +1,36 @@
+use crate::error;
+use crate::message::Message;
 use crate::sink::{ResponseFromSink, ResponseStatusFromSink, Sink};
-use crate::{error, message::Message};
+use bytes::Bytes;
 
+/// User defined serving store to store the serving responses.
+pub(crate) mod user_defined;
+
+/// Nats serving store to store the serving responses.
+pub(crate) mod nats;
+
+/// Enum to represent different types of serving stores.
+#[derive(Clone)]
+pub(crate) enum ServingStore {
+    UserDefined(user_defined::UserDefinedStore),
+    Nats(nats::NatsServingStore),
+}
+
+/// Entry in the serving store.
+#[derive(Clone, Debug)]
+pub(crate) struct StoreEntry {
+    /// Pod Hash is for filtering the stream by each request originating pod while listening for
+    /// "put" onto the KV store. This is used only on ISB KV store and this enables SSE (converse is
+    /// that SSE is not supported on user defined store).
+    pub(crate) pod_hash: String,
+    /// Unique ID Of the request to which the result belongs. There could be multiple results for
+    /// the same request.
+    pub(crate) id: String,
+    /// The result of the computation.
+    pub(crate) value: Bytes,
+}
+
+/// Builtin Sink to write to the serving store.
 #[derive(Debug, Default, Clone)]
 pub(crate) struct ServeSink;
 
@@ -24,9 +54,9 @@ mod tests {
 
     use chrono::Utc;
 
-    use super::ServeSink;
     use crate::message::IntOffset;
     use crate::message::{Message, MessageID, Offset};
+    use crate::sink::serve::ServeSink;
     use crate::sink::{ResponseFromSink, ResponseStatusFromSink, Sink};
 
     #[tokio::test]

rust/numaflow-core/src/sink/serve.rs

@@ -1,5 +1,5 @@
 use crate::config::pipeline::NatsStoreConfig;
-use crate::serving_store::StoreEntry;
+use crate::sink::serve::StoreEntry;
 use async_nats::jetstream::Context;
 use async_nats::jetstream::kv::Store;
 use chrono::Utc;

rust/numaflow-core/src/serving_store/nats.rs renamed to rust/numaflow-core/src/sink/serve/nats.rs

@@ -4,8 +4,8 @@ use tokio::task::JoinSet;
 use tonic::transport::Channel;
 
 use crate::config::pipeline::UserDefinedStoreConfig;
-use crate::serving_store::StoreEntry;
 use crate::shared;
+use crate::sink::serve::StoreEntry;
 
 /// User defined serving store to store the serving responses.
 #[derive(Clone)]