chore: Batch ack requests and sink responses for better performance (#149)

yhl25 · web-flow · commit 134414ab1179 · 2024-10-31T20:43:23.000+05:30
Signed-off-by: Yashash H L &lt;yashashhl25@gmail.com&gt;
diff --git a/examples/src/main/java/io/numaproj/numaflow/examples/source/simple/SimpleSource.java b/examples/src/main/java/io/numaproj/numaflow/examples/source/simple/SimpleSource.java
@@ -72,9 +72,11 @@ public void read(ReadRequest request, OutputObserver observer) {
 
     @Override
     public void ack(AckRequest request) {
-        Long offset = Longs.fromByteArray(request.getOffset().getValue());
-        // remove the acknowledged messages from the map
-        messages.remove(offset);
+        for (Offset offset : request.getOffsets()) {
+            Long decoded_offset = Longs.fromByteArray(offset.getValue());
+            // remove the acknowledged messages from the map
+            messages.remove(decoded_offset);
+        }
     }
 
     @Override
diff --git a/examples/src/test/java/io/numaproj/numaflow/examples/source/simple/SimpleSourceTest.java b/examples/src/test/java/io/numaproj/numaflow/examples/source/simple/SimpleSourceTest.java
@@ -39,11 +39,9 @@ public void test_ReadAndAck() {
             offsets.add(message.getOffset());
         }
 
-        for (Offset offset : offsets) {
-            SourcerTestKit.TestAckRequest ackRequest = SourcerTestKit.TestAckRequest.builder()
-                    .offset(offset).build();
-            simpleSource.ack(ackRequest);
-        }
+        SourcerTestKit.TestAckRequest ackRequest = SourcerTestKit.TestAckRequest.builder()
+                .offsets(offsets).build();
+        simpleSource.ack(ackRequest);
 
         // Try reading 6 more messages
         // Since the previous batch got acked, the data source should allow us to read more messages
diff --git a/src/main/java/io/numaproj/numaflow/sinker/Service.java b/src/main/java/io/numaproj/numaflow/sinker/Service.java
@@ -73,10 +73,11 @@ public void onNext(SinkOuterClass.SinkRequest request) {
                         datumStream.writeMessage(HandlerDatum.EOF_DATUM);
 
                         ResponseList responses = result.join();
-                        responses.getResponses().forEach(response -> {
-                            SinkOuterClass.SinkResponse sinkResponse = buildResponse(response);
-                            responseObserver.onNext(sinkResponse);
-                        });
+                        SinkOuterClass.SinkResponse.Builder responseBuilder = SinkOuterClass.SinkResponse.newBuilder();
+                        for (Response response : responses.getResponses()) {
+                            responseBuilder.addResults(buildResult(response));
+                        }
+                        responseObserver.onNext(responseBuilder.build());
 
                         // send eot response to indicate end of transmission for the batch
                         SinkOuterClass.SinkResponse eotResponse = SinkOuterClass.SinkResponse
@@ -113,15 +114,13 @@ public void onCompleted() {
         };
     }
 
-    private SinkOuterClass.SinkResponse buildResponse(Response response) {
+    private SinkOuterClass.SinkResponse.Result buildResult(Response response) {
         SinkOuterClass.Status status = response.getFallback() ? SinkOuterClass.Status.FALLBACK :
                 response.getSuccess() ? SinkOuterClass.Status.SUCCESS : SinkOuterClass.Status.FAILURE;
-        return SinkOuterClass.SinkResponse.newBuilder()
-                .setResult(SinkOuterClass.SinkResponse.Result.newBuilder()
-                        .setId(response.getId() == null ? "" : response.getId())
-                        .setErrMsg(response.getErr() == null ? "" : response.getErr())
-                        .setStatus(status)
-                        .build())
+        return SinkOuterClass.SinkResponse.Result.newBuilder()
+                .setId(response.getId() == null ? "" : response.getId())
+                .setErrMsg(response.getErr() == null ? "" : response.getErr())
+                .setStatus(status)
                 .build();
     }
 
diff --git a/src/main/java/io/numaproj/numaflow/sinker/SinkerTestKit.java b/src/main/java/io/numaproj/numaflow/sinker/SinkerTestKit.java
@@ -197,19 +197,22 @@ 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()
-                            .getId()));
-                } else if (result.getResult().getStatus() == SinkOuterClass.Status.FALLBACK) {
-                    responseListBuilder.addResponse(Response.responseFallback(
-                            result.getResult().getId()));
-                } else {
-                    responseListBuilder.addResponse(Response.responseFailure(
-                            result.getResult().getId(), result.getResult().getErrMsg()));
+
+                for (SinkOuterClass.SinkResponse.Result response : result.getResultsList()) {
+                    if (response.getStatus() == SinkOuterClass.Status.SUCCESS) {
+                        responseListBuilder.addResponse(Response.responseOK(response
+                                .getId()));
+                    } else if (response.getStatus() == SinkOuterClass.Status.FALLBACK) {
+                        responseListBuilder.addResponse(Response.responseFallback(
+                                response.getId()));
+                    } else {
+                        responseListBuilder.addResponse(Response.responseFailure(
+                                response.getId(), response.getErrMsg()));
+                    }
                 }
             }
 
diff --git a/src/main/java/io/numaproj/numaflow/sourcer/AckRequest.java b/src/main/java/io/numaproj/numaflow/sourcer/AckRequest.java
@@ -1,12 +1,14 @@
 package io.numaproj.numaflow.sourcer;
 
 
+import java.util.List;
+
 /**
  * AckRequest request for acknowledging messages.
  */
 public interface AckRequest {
     /**
-     * @return the offset to be acknowledged
+     * @return the list of offsets to be acknowledged
      */
-    Offset getOffset();
+    List<Offset> getOffsets();
 }
diff --git a/src/main/java/io/numaproj/numaflow/sourcer/AckRequestImpl.java b/src/main/java/io/numaproj/numaflow/sourcer/AckRequestImpl.java
@@ -2,15 +2,17 @@
 
 import lombok.AllArgsConstructor;
 
+import java.util.List;
+
 /**
  * AckRequestImpl is the implementation of AckRequest.
  */
 @AllArgsConstructor
 class AckRequestImpl implements AckRequest {
-    private final Offset offset;
+    private final List<Offset> offsets;
 
     @Override
-    public Offset getOffset() {
-        return this.offset;
+    public List<Offset> getOffsets() {
+        return this.offsets;
     }
 }
diff --git a/src/main/java/io/numaproj/numaflow/sourcer/Service.java b/src/main/java/io/numaproj/numaflow/sourcer/Service.java
@@ -7,6 +7,7 @@
 import io.numaproj.numaflow.source.v1.SourceOuterClass;
 
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.List;
 
 import static io.numaproj.numaflow.source.v1.SourceGrpc.getPendingFnMethod;
@@ -115,14 +116,16 @@ public void onNext(SourceOuterClass.AckRequest request) {
                     return;
                 }
 
-                SourceOuterClass.Offset offset = request.getRequest().getOffset();
+                List<Offset> offsets = new ArrayList<>(request.getRequest().getOffsetsCount());
+                for (SourceOuterClass.Offset offset : request.getRequest().getOffsetsList()) {
+                    offsets.add(new Offset(
+                            offset.getOffset().toByteArray(),
+                            offset.getPartitionId()));
+                }
 
-                AckRequestImpl ackRequest = new AckRequestImpl(new Offset(
-                        offset.getOffset().toByteArray(),
-                        offset.getPartitionId()));
+                AckRequestImpl ackRequest = new AckRequestImpl(offsets);
 
                 // invoke the sourcer's ack method
-
                 sourcer.ack(ackRequest);
 
                 // send an ack response to the client after acking the message
diff --git a/src/main/java/io/numaproj/numaflow/sourcer/SourcerTestKit.java b/src/main/java/io/numaproj/numaflow/sourcer/SourcerTestKit.java
@@ -16,6 +16,7 @@
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 
 /**
  * SourcerTestKit is a test kit for testing Sourcer implementations.
@@ -165,14 +166,15 @@ public void onCompleted() {
          */
         public void sendAckRequest(AckRequest request) throws Exception {
             CompletableFuture<SourceOuterClass.AckResponse> future = new CompletableFuture<>();
+
             SourceOuterClass.AckRequest.Request.Builder builder = SourceOuterClass.AckRequest.Request
                     .newBuilder()
-                    .setOffset(SourceOuterClass.Offset.newBuilder()
-                            .setOffset(com.google.protobuf.ByteString.copyFrom(request
-                                    .getOffset()
-                                    .getValue()))
-                            .setPartitionId(request.getOffset().getPartitionId())
-                            .build());
+                    .addAllOffsets(request.getOffsets().stream().map(
+                            offset -> SourceOuterClass.Offset.newBuilder()
+                                    .setOffset(com.google.protobuf.ByteString.copyFrom(offset.getValue()))
+                                    .setPartitionId(offset.getPartitionId())
+                                    .build()
+                    ).collect(Collectors.toList()));
 
             SourceOuterClass.AckRequest grpcRequest = SourceOuterClass.AckRequest.newBuilder()
                     .setRequest(builder.build())
@@ -288,7 +290,7 @@ public static class TestReadRequest implements ReadRequest {
     @Setter
     @Builder
     public static class TestAckRequest implements AckRequest {
-        Offset offset;
+        List<Offset> offsets;
     }
 
     /**
diff --git a/src/main/proto/sink/v1/sink.proto b/src/main/proto/sink/v1/sink.proto
@@ -79,7 +79,7 @@ message SinkResponse {
     // err_msg is the error message, set it if success is set to false.
     string err_msg = 3;
   }
-  Result result = 1;
+  repeated Result results = 1;
   optional Handshake handshake = 2;
   optional TransmissionStatus status = 3;
 }
diff --git a/src/main/proto/source/v1/source.proto b/src/main/proto/source/v1/source.proto
@@ -112,8 +112,8 @@ message ReadResponse {
  */
 message AckRequest {
   message Request {
-    // Required field holding the offset to be acked
-    Offset offset = 1;
+    // Required field holding the offsets to be acked
+    repeated Offset offsets = 1;
   }
   // Required field holding the request. The list will be ordered and will have the same order as the original Read response.
   Request request = 1;
diff --git a/src/test/java/io/numaproj/numaflow/sinker/ServerTest.java b/src/test/java/io/numaproj/numaflow/sinker/ServerTest.java
@@ -114,19 +114,19 @@ public void sinkerSuccess() {
 
         while (!outputStreamObserver.completed.get()) ;
         List<SinkOuterClass.SinkResponse> responseList = outputStreamObserver.getSinkResponse();
-        assertEquals(111, responseList.size());
+        assertEquals(21, 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()) {
+        var response = responseList.get(0);
+        response.getResultsList().forEach(result -> {
+            if (result.getStatus() == SinkOuterClass.Status.FAILURE) {
+                assertEquals(result.getErrMsg(), "error message");
                 return;
             }
-            assertEquals(response.getResult().getId(), expectedId);
-            if (response.getResult().getStatus() == SinkOuterClass.Status.FAILURE) {
-                assertEquals(response.getResult().getErrMsg(), "error message");
-            }
+            assertEquals(result.getId(), expectedId);
+            assertEquals(result.getStatus(), SinkOuterClass.Status.SUCCESS);
         });
     }
 
diff --git a/src/test/java/io/numaproj/numaflow/sourcer/ServerTest.java b/src/test/java/io/numaproj/numaflow/sourcer/ServerTest.java
@@ -104,7 +104,7 @@ public void onNext(SourceOuterClass.ReadResponse readResponse) {
                         .newBuilder()
                         .getRequest()
                         .toBuilder()
-                        .setOffset(offset)
+                        .addOffsets(offset)
                         .build();
                 ackRequests.add(SourceOuterClass.AckRequest
                         .newBuilder()
@@ -243,8 +243,10 @@ public List<Integer> getPartitions() {
 
         @Override
         public void ack(AckRequest request) {
-            Integer offset = ByteBuffer.wrap(request.getOffset().getValue()).getInt();
-            yetToBeAcked.remove(offset);
+            for (Offset offset : request.getOffsets()) {
+                Integer decoded_offset = ByteBuffer.wrap(offset.getValue()).getInt();
+                yetToBeAcked.remove(decoded_offset);
+            }
         }
 
         @Override

Original file line number	Diff line number	Diff line change
`@@ -79,7 +79,7 @@ message SinkResponse {`
`79`	`79`	`// err_msg is the error message, set it if success is set to false.`
`80`	`80`	`string err_msg = 3;`
`81`	`81`	`}`
`82`		`- Result result = 1;`
	`82`	`+ repeated Result results = 1;`
`83`	`83`	`optional Handshake handshake = 2;`
`84`	`84`	`optional TransmissionStatus status = 3;`
`85`	`85`	`}`