feat: Custom Serve Store for integration with any user-defined store (#113)

Signed-off-by: Sreekanth <[email protected]>
Signed-off-by: Yashash H L <[email protected]>
Co-authored-by: Yashash H L <[email protected]>
Co-authored-by: Vigith Maurice <[email protected]>

Author: 3 people

Cargo.toml

@@ -5,4 +5,4 @@ default-members = ["numaflow"]
 resolver = "2"
 
 [workspace.package]
-rust-version = "1.82"
+rust-version = "1.84"

examples/simple-source/src/main.rs

@@ -1,4 +1,5 @@
-///! An example for simple User Defined Source. It generates a continuous increasing sequence of offsets and some data for each call to [`numaflow::source::sourcer::read`].
+//! An example for simple User Defined Source. It generates a continuous increasing sequence of offsets
+//! and some data for each call to [`numaflow::source::sourcer::read`].
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {

numaflow/Cargo.toml

@@ -2,7 +2,7 @@
 name = "numaflow"
 version = "0.2.1"
 edition = "2021"
-rust-version = "1.80"
+rust-version = "1.84"
 description = "Rust SDK for Numaflow"
 authors = ["Numaflow Developers"]
 license = "Apache-2.0"
@@ -34,10 +34,9 @@ tokio-stream = { version = "0.1.16", features = ["net"] }
 serde = { version = "1.0.210", features = ["derive"] }
 chrono = "0.4.38"
 serde_json = "1.0.128"
-futures-util = "0.3.30"
 tracing = "0.1.40"
 uuid = { version = "1.10.0", features = ["v4"] }
-thiserror = "1.0"
+thiserror = "2.0.12"
 hyper-util = "0.1.7"
 
 [build-dependencies]

numaflow/build.rs

@@ -15,6 +15,7 @@ fn main() {
                 "proto/reduce.proto",
                 "proto/sink.proto",
                 "proto/sideinput.proto",
+                "proto/store.proto",
             ],
             &["proto"],
         )

numaflow/proto/sink.proto

@@ -3,6 +3,7 @@ syntax = "proto3";
 import "google/protobuf/empty.proto";
 import "google/protobuf/timestamp.proto";
 
+
 package sink.v1;
 
 service Sink {
@@ -63,6 +64,7 @@ enum Status {
   SUCCESS = 0;
   FAILURE = 1;
   FALLBACK = 2;
+  SERVE = 3;
 }
 
 /**
@@ -76,8 +78,9 @@ message SinkResponse {
     Status status = 2;
     // err_msg is the error message, set it if success is set to false.
     string err_msg = 3;
+    optional bytes serve_response = 4;
   }
   repeated Result results = 1;
   optional Handshake handshake = 2;
   optional TransmissionStatus status = 3;
-}
+}

numaflow/proto/store.proto

@@ -0,0 +1,77 @@
+/*
+Copyright 2022 The Numaproj Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+syntax = "proto3";
+
+import "google/protobuf/empty.proto";
+import "google/protobuf/timestamp.proto";
+
+package serving.v1;
+
+// ServingStore defines a set of methods to interface with a user-defined Store.
+service ServingStore {
+  // Put is to put the PutRequest into the Store.
+  rpc Put(PutRequest) returns (PutResponse);
+
+  // Get gets the GetRequest from the Store.
+  rpc Get(GetRequest) returns (GetResponse);
+
+  // IsReady checks the health of the container interfacing the Store.
+  rpc IsReady(google.protobuf.Empty) returns (ReadyResponse);
+}
+
+// Payload that represent the output that is to be written into to the store.
+message Payload {
+  // Origin is the Vertex that generated this result.
+  string origin = 1;
+  // Value is the result of the computation.
+  bytes value = 2;
+}
+
+// PutRequest is the request sent to the Store.
+message PutRequest {
+  // ID is the unique id as provided by the user in the original request. If not provided, it will be a system generated
+  // uuid.
+  string id = 1;
+  // Payloads are one or more results generated (could be more than one due to flat-map).
+  repeated Payload payloads = 2;
+}
+
+// PutResponse is the result of the Put call.
+message PutResponse {
+  bool success = 1;
+}
+
+// GetRequest is the call to get the result stored in the Store.
+message GetRequest {
+  // ID is the unique id as provided by the user in the original request. If not provided, it will be a system generated
+  // uuid.
+  string id = 1;
+}
+
+// GetResponse is the result stored in the Store.
+message GetResponse {
+  string id = 1;
+  // Payloads are one or more results generated (could be more than one due to flat-map).
+  repeated Payload payloads = 2;
+}
+
+/**
+ * ReadyResponse is the health check result.
+ */
+message ReadyResponse {
+  bool ready = 1;
+}

numaflow/src/batchmap.rs

@@ -354,7 +354,7 @@ where
             };
 
             // we are done with this batch because eot=true
-            if message.status.map_or(false, |status| status.eot) {
+            if message.status.is_some_and(|status| status.eot) {
                 debug!("Batch Ended, received an EOT message");
                 break;
             }

numaflow/src/error.rs

@@ -33,4 +33,7 @@ pub enum Error {
 
     #[error("SideInput Error: {0}")]
     SideInputError(ErrorKind),
+
+    #[error("ServingStore Error: {0}")]
+    ServingStoreError(ErrorKind),
 }

numaflow/src/lib.rs

@@ -40,6 +40,8 @@ pub mod batchmap;
 /// mapstream is for writing the map in [stream mode](https://numaflow.numaproj.io/user-guide/user-defined-functions/map/map/#streaming-mode) handlers.
 pub mod mapstream;
 
+pub mod serving_store;
+
 mod servers;
 
 // Error handling on Numaflow SDKs!

numaflow/src/servers.rs

@@ -21,3 +21,7 @@ pub mod source;
 #[path = "servers/sourcetransformer.v1.rs"]
 #[rustfmt::skip]
 pub mod sourcetransformer;
+
+#[path = "servers/serving.v1.rs"]
+#[rustfmt::skip]
+pub mod serving;