rename and add eot to sink

Signed-off-by: Yashash H L <[email protected]>
numaproj · Oct 14, 2024 · 4b35fa4 · 4b35fa4
1 parent 0f0ec35
commit 4b35fa4
Show file tree

Hide file tree

Showing 10 changed files with 40 additions and 15 deletions.
diff --git a/src/main/java/io/numaproj/numaflow/batchmapper/Service.java b/src/main/java/io/numaproj/numaflow/batchmapper/Service.java
@@ -149,7 +149,8 @@ private void buildAndStreamResponse(
         // Send an EOT message to indicate the end of the transmission for the batch.
         MapOuterClass.MapResponse eotResponse = MapOuterClass.MapResponse
                 .newBuilder()
-                .setStatus(MapOuterClass.Status.newBuilder().setEot(true).build()).build();
+                .setStatus(MapOuterClass.TransmissionStatus.newBuilder().setEot(true).build())
+                .build();
         responseObserver.onNext(eotResponse);
         responseObserver.onCompleted();
     }

diff --git a/src/main/java/io/numaproj/numaflow/mapstreamer/Service.java b/src/main/java/io/numaproj/numaflow/mapstreamer/Service.java
@@ -66,7 +66,10 @@ public void onNext(MapOuterClass.MapRequest request) {
                 // Send an EOT message to indicate the end of the transmission for the batch.
                 MapOuterClass.MapResponse eotResponse = MapOuterClass.MapResponse
                         .newBuilder()
-                        .setStatus(MapOuterClass.Status.newBuilder().setEot(true).build()).build();
+                        .setStatus(MapOuterClass.TransmissionStatus
+                                .newBuilder()
+                                .setEot(true)
+                                .build()).build();
                 responseObserver.onNext(eotResponse);
             }
 

diff --git a/src/main/java/io/numaproj/numaflow/sinker/Service.java b/src/main/java/io/numaproj/numaflow/sinker/Service.java
@@ -78,6 +78,16 @@ public void onNext(SinkOuterClass.SinkRequest request) {
                             responseObserver.onNext(sinkResponse);
                         });
 
+                        // send eot response to indicate end of transmission for the batch
+                        SinkOuterClass.SinkResponse eotResponse = SinkOuterClass.SinkResponse
+                                .newBuilder()
+                                .setStatus(SinkOuterClass.TransmissionStatus
+                                        .newBuilder()
+                                        .setEot(true)
+                                        .build())
+                                .build();
+                        responseObserver.onNext(eotResponse);
+
                         // reset the startOfStream flag, since the stream has ended
                         // so that the next request will be treated as the start of the stream
                         startOfStream = true;

diff --git a/src/main/java/io/numaproj/numaflow/sinker/SinkerTestKit.java b/src/main/java/io/numaproj/numaflow/sinker/SinkerTestKit.java
@@ -181,7 +181,7 @@ public void onCompleted() {
             }
             // send end of transmission message
             requestObserver.onNext(SinkOuterClass.SinkRequest.newBuilder().setStatus(
-                    SinkOuterClass.SinkRequest.Status.newBuilder().setEot(true)).build());
+                    SinkOuterClass.TransmissionStatus.newBuilder().setEot(true)).build());
 
             requestObserver.onCompleted();
 
@@ -197,6 +197,9 @@ public void onCompleted() {
                 if (result.getHandshake().getSot()) {
                     continue;
                 }
+                if (result.hasStatus() && result.getStatus().getEot()) {
+                    continue;
+                }
                 if (result.getResult().getStatus() == SinkOuterClass.Status.SUCCESS) {
                     responseListBuilder.addResponse(Response.responseOK(result
                             .getResult()

diff --git a/src/main/proto/map/v1/map.proto b/src/main/proto/map/v1/map.proto
@@ -30,7 +30,7 @@ message MapRequest {
   // This ID is used to uniquely identify a map request
   string id = 2;
   optional Handshake handshake = 3;
-  optional Status status = 4;
+  optional TransmissionStatus status = 4;
 }
 
 /*
@@ -44,7 +44,7 @@ message Handshake {
 /*
  * Status message to indicate the status of the message.
  */
-message Status {
+message TransmissionStatus {
   bool eot = 1;
 }
 
@@ -61,7 +61,7 @@ message MapResponse {
   // This ID is used to refer the responses to the request it corresponds to.
   string id = 2;
   optional Handshake handshake = 3;
-  optional Status status = 4;
+  optional TransmissionStatus status = 4;
 }
 
 /**

diff --git a/src/main/proto/sink/v1/sink.proto b/src/main/proto/sink/v1/sink.proto
@@ -27,14 +27,11 @@ message SinkRequest {
     string id = 5;
     map<string, string> headers = 6;
   }
-  message Status {
-    bool eot = 1;
-  }
   // Required field indicating the request.
   Request request = 1;
   // Required field indicating the status of the request.
   // If eot is set to true, it indicates the end of transmission.
-  Status status = 2;
+  TransmissionStatus status = 2;
   // optional field indicating the handshake message.
   optional Handshake handshake = 3;
 }
@@ -54,6 +51,13 @@ message ReadyResponse {
   bool ready = 1;
 }
 
+/**
+ * TransmissionStatus is the status of the transmission.
+ */
+message TransmissionStatus {
+  bool eot = 1;
+}
+
 /*
   * Status is the status of the response.
  */
@@ -77,4 +81,5 @@ message SinkResponse {
   }
   Result result = 1;
   optional Handshake handshake = 2;
+  optional TransmissionStatus status = 3;
 }
diff --git a/src/test/java/io/numaproj/numaflow/batchmapper/ServerErrTest.java b/src/test/java/io/numaproj/numaflow/batchmapper/ServerErrTest.java
@@ -127,7 +127,7 @@ public void testErrorFromUDF() {
         inputStreamObserver.onNext(request);
         inputStreamObserver.onNext(MapOuterClass.MapRequest
                 .newBuilder()
-                .setStatus(MapOuterClass.Status.newBuilder().setEot(true))
+                .setStatus(MapOuterClass.TransmissionStatus.newBuilder().setEot(true))
                 .build());
         inputStreamObserver.onCompleted();
         try {

diff --git a/src/test/java/io/numaproj/numaflow/batchmapper/ServerTest.java b/src/test/java/io/numaproj/numaflow/batchmapper/ServerTest.java
@@ -90,7 +90,7 @@ public void testBatchMapHappyPath() {
 
         inputStreamObserver.onNext(MapOuterClass.MapRequest
                 .newBuilder()
-                .setStatus(MapOuterClass.Status.newBuilder().setEot(true))
+                .setStatus(MapOuterClass.TransmissionStatus.newBuilder().setEot(true))
                 .build());
         inputStreamObserver.onCompleted();
 

diff --git a/src/test/java/io/numaproj/numaflow/sinker/ServerErrTest.java b/src/test/java/io/numaproj/numaflow/sinker/ServerErrTest.java
@@ -110,7 +110,7 @@ public void sinkerException() {
 
         // send eot message
         inputStreamObserver.onNext(SinkOuterClass.SinkRequest.newBuilder()
-                .setStatus(SinkOuterClass.SinkRequest.Status.newBuilder().setEot(true)).build());
+                .setStatus(SinkOuterClass.TransmissionStatus.newBuilder().setEot(true)).build());
 
         inputStreamObserver.onCompleted();
 

diff --git a/src/test/java/io/numaproj/numaflow/sinker/ServerTest.java b/src/test/java/io/numaproj/numaflow/sinker/ServerTest.java
@@ -101,7 +101,7 @@ public void sinkerSuccess() {
             // If it's the end of the batch, send an EOT message
             if (i % 10 == 0) {
                 SinkOuterClass.SinkRequest eotRequest = SinkOuterClass.SinkRequest.newBuilder()
-                        .setStatus(SinkOuterClass.SinkRequest.Status
+                        .setStatus(SinkOuterClass.TransmissionStatus
                                 .newBuilder()
                                 .setEot(true)
                                 .build())
@@ -114,12 +114,15 @@ public void sinkerSuccess() {
 
         while (!outputStreamObserver.completed.get()) ;
         List<SinkOuterClass.SinkResponse> responseList = outputStreamObserver.getSinkResponse();
-        assertEquals(101, responseList.size());
+        assertEquals(111, responseList.size());
         // first response is the handshake response
         assertTrue(responseList.get(0).getHandshake().getSot());
 
         responseList = responseList.subList(1, responseList.size());
         responseList.forEach(response -> {
+            if (response.hasStatus() && response.getStatus().getEot()) {
+                return;
+            }
             assertEquals(response.getResult().getId(), expectedId);
             if (response.getResult().getStatus() == SinkOuterClass.Status.FAILURE) {
                 assertEquals(response.getResult().getErrMsg(), "error message");