Add support for on success sink

Signed-off-by: vtiwari5 <[email protected]>

Author: vtiwari5

src/main/java/io/numaproj/numaflow/sinker/Constants.java

@@ -3,14 +3,18 @@
 class Constants {
   public static final String DEFAULT_SOCKET_PATH = "/var/run/numaflow/sink.sock";
   public static final String DEFAULT_FB_SINK_SOCKET_PATH = "/var/run/numaflow/fb-sink.sock";
+  public static final String DEFAULT_ON_SUCCESS_SINK_SOCKET_PATH = "/var/run/numaflow/ons-sink.sock";
   public static final String DEFAULT_SERVER_INFO_FILE_PATH = "/var/run/numaflow/sinker-server-info";
   public static final String DEFAULT_FB_SERVER_INFO_FILE_PATH =
       "/var/run/numaflow/fb-sinker-server-info";
+  public static final String DEFAULT_ON_SUCCESS_SERVER_INFO_FILE_PATH =
+      "/var/run/numaflow/ons-sinker-server-info";
   public static final int DEFAULT_MESSAGE_SIZE = 1024 * 1024 * 64;
   public static final int DEFAULT_PORT = 50051;
   public static final String DEFAULT_HOST = "localhost";
   public static final String ENV_UD_CONTAINER_TYPE = "NUMAFLOW_UD_CONTAINER_TYPE";
   public static final String UD_CONTAINER_FALLBACK_SINK = "fb-udsink";
+  public static final String UD_CONTAINER_ON_SUCCESS_SINK = "ons-udsink";
 
   // Private constructor to prevent instantiation
   private Constants() {

src/main/java/io/numaproj/numaflow/sinker/GRPCConfig.java

@@ -36,6 +36,9 @@ static GRPCConfig defaultGrpcConfig() {
         if (Constants.UD_CONTAINER_FALLBACK_SINK.equals(containerType)) {
             socketPath = Constants.DEFAULT_FB_SINK_SOCKET_PATH;
             infoFilePath = Constants.DEFAULT_FB_SERVER_INFO_FILE_PATH;
+        } else if (Constants.UD_CONTAINER_ON_SUCCESS_SINK.equals(containerType)) {
+            socketPath = Constants.DEFAULT_ON_SUCCESS_SINK_SOCKET_PATH;
+            infoFilePath = Constants.DEFAULT_ON_SUCCESS_SERVER_INFO_FILE_PATH;
         }
         return GRPCConfig.newBuilder()
                 .infoFilePath(infoFilePath)

src/main/java/io/numaproj/numaflow/sinker/KeyValueGroup.java

@@ -0,0 +1,12 @@
+package io.numaproj.numaflow.sinker;
+
+import lombok.Builder;
+import lombok.Getter;
+
+import java.util.HashMap;
+
+@Getter
+@Builder
+public class KeyValueGroup {
+    private final HashMap<String, byte[]> keyValue;
+}

src/main/java/io/numaproj/numaflow/sinker/OnSuccessMessage.java

@@ -0,0 +1,16 @@
+package io.numaproj.numaflow.sinker;
+
+import lombok.Builder;
+import lombok.Getter;
+
+import java.util.HashMap;
+
+@Getter
+@Builder
+public class OnSuccessMessage {
+    private final byte[] value;
+    private final String key;
+    private final HashMap<String, KeyValueGroup> userMetadata;
+}
+
+

src/main/java/io/numaproj/numaflow/sinker/Response.java

@@ -1,9 +1,15 @@
 package io.numaproj.numaflow.sinker;
 
+import com.google.protobuf.ByteString;
+import common.MetadataOuterClass;
+import io.numaproj.numaflow.sink.v1.SinkOuterClass.SinkResponse.Result.Message;
 import lombok.AccessLevel;
 import lombok.AllArgsConstructor;
 import lombok.Getter;
 
+import java.util.Map;
+import java.util.stream.Collectors;
+
 /**
  * Response is used to send response from the user defined sinker. It contains the id of the
  * message, success status, an optional error message and a fallback status. Various static factory
@@ -16,6 +22,12 @@ public class Response {
   private final Boolean success;
   private final String err;
   private final Boolean fallback;
+  private final Boolean serve;
+  private final byte[] serveResponse;
+  private final Boolean onSuccess;
+  // FIXME: Should this be OnSuccessMessage object from package? That would allow parity with other SDKs (specially Go)
+  // Currently done this way to prevent conversion in buildResult method.
+  private final Message onSuccessMessage;
 
   /**
    * Static method to create response for successful message processing.
@@ -24,7 +36,7 @@ public class Response {
    * @return Response object with success status
    */
   public static Response responseOK(String id) {
-    return new Response(id, true, null, false);
+    return new Response(id, true, null, false, false, null, false, null);
   }
 
   /**
@@ -35,7 +47,7 @@ public static Response responseOK(String id) {
    * @return Response object with failure status and error message
    */
   public static Response responseFailure(String id, String errMsg) {
-    return new Response(id, false, errMsg, false);
+    return new Response(id, false, errMsg, false, false, null, false, null);
   }
 
   /**
@@ -46,6 +58,76 @@ public static Response responseFailure(String id, String errMsg) {
    * @return Response object with fallback status
    */
   public static Response responseFallback(String id) {
-    return new Response(id, false, null, true);
+    return new Response(id, false, null, true, false, null, false, null);
+  }
+
+  /**
+   * Static method to create response for serve message which is raw bytes.
+   * This indicates that the message should be sent to the serving store.
+   * Allows creation of serve message from raw bytes.
+   *
+   * @param id id of the message
+   * @param serveResponse Response object to be sent to the serving store
+   * @return Response object with serve status and serve response
+   */
+  public static Response responseServe(String id, byte[] serveResponse) {
+    return new Response(id, false, null, false, true, serveResponse, false, null);
+  }
+
+  /**
+   * Static method to create response for onSuccess message. Allows creation of onSuccess message
+   * from protobuf Message object.
+   *
+   * @param id id of the message
+   * @param onSuccessMessage OnSuccessMessage object to be sent to the onSuccess sink
+   * @return Response object with onSuccess status and onSuccess message
+   */
+  public static Response responseOnSuccess(String id, Message onSuccessMessage) {
+    return new Response(id, false, null, false, false, null, true, onSuccessMessage);
+  }
+
+  /**
+   * Overloaded static method to create response for onSuccess message. Allows creation of onSuccess message
+   * from OnSuccessMessage object.
+   *
+   * @param id id of the message
+   * @param onSuccessMessage OnSuccessMessage object to be sent to the onSuccess sink. Can be null
+   * if original message needs to be written to onSuccess sink
+   * @return Response object with onSuccess status and onSuccess message
+   */
+  public static Response responseOnSuccess(String id, OnSuccessMessage onSuccessMessage) {
+      if (onSuccessMessage == null) {
+          return new Response(id, false, null, false, false, null, true, null);
+      } else {
+          Map<String, MetadataOuterClass.KeyValueGroup> pbUserMetadata = onSuccessMessage.getUserMetadata()
+                  .entrySet()
+                  .stream()
+                  .collect(Collectors.toMap(
+                          Map.Entry::getKey,
+                          e -> MetadataOuterClass.KeyValueGroup.newBuilder()
+                                  .putAllKeyValue(e.getValue()
+                                          .getKeyValue()
+                                          .entrySet()
+                                          .stream()
+                                          .collect(Collectors.toMap(
+                                                  Map.Entry::getKey,
+                                                  kv -> ByteString.copyFrom(kv.getValue())
+                                          ))
+                                  )
+                                  .build()
+          ));
+
+          MetadataOuterClass.Metadata pbMetadata = MetadataOuterClass.Metadata.newBuilder()
+                  .putAllUserMetadata(pbUserMetadata)
+                  .build();
+
+          Message pbOnSuccessMessage = Message.newBuilder()
+                  .addKeys(onSuccessMessage.getKey())
+                  .setValue(ByteString.copyFrom(onSuccessMessage.getValue()))
+                  .setMetadata(pbMetadata)
+                  .build();
+
+          return new Response(id, false, null, false, false, null, true, pbOnSuccessMessage);
+      }
   }
 }

src/main/java/io/numaproj/numaflow/sinker/Service.java

@@ -1,6 +1,7 @@
 package io.numaproj.numaflow.sinker;
 
 import com.google.protobuf.Any;
+import com.google.protobuf.ByteString;
 import com.google.protobuf.Empty;
 import com.google.rpc.Code;
 import com.google.rpc.DebugInfo;
@@ -131,13 +132,38 @@ public void onCompleted() {
     }
 
     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.Result.newBuilder()
-                .setId(response.getId() == null ? "" : response.getId())
-                .setErrMsg(response.getErr() == null ? "" : response.getErr())
-                .setStatus(status)
-                .build();
+        if (response.getFallback()) {
+            return SinkOuterClass.SinkResponse.Result.newBuilder()
+                    .setId(response.getId() == null ? "" : response.getId())
+                    .setStatus(SinkOuterClass.Status.FALLBACK)
+                    .build();
+        } else if (response.getSuccess()) {
+            return SinkOuterClass.SinkResponse.Result.newBuilder()
+                    .setId(response.getId() == null ? "" : response.getId())
+                    .setStatus(SinkOuterClass.Status.SUCCESS)
+                    .build();
+        } else if (response.getServe()) {
+            // FIXME: Return error when serve response is not set?
+            return SinkOuterClass.SinkResponse.Result.newBuilder()
+                    .setId(response.getId() == null ? "" : response.getId())
+                    .setStatus(SinkOuterClass.Status.SERVE)
+                    .setServeResponse(response.getServeResponse() == null ? null : ByteString.copyFrom(
+                            response.getServeResponse()))
+                    .build();
+        } else if (response.getOnSuccess()) {
+            return SinkOuterClass.SinkResponse.Result.newBuilder()
+                    .setId(response.getId() == null ? "" : response.getId())
+                    .setStatus(SinkOuterClass.Status.ON_SUCCESS)
+                    .setOnSuccessMsg(response.getOnSuccessMessage() == null ? null : response.getOnSuccessMessage())
+                    .build();
+        } else {
+            // FIXME: Return error when error message is not set?
+            return SinkOuterClass.SinkResponse.Result.newBuilder()
+                    .setId(response.getId() == null ? "" : response.getId())
+                    .setStatus(SinkOuterClass.Status.FAILURE)
+                    .setErrMsg(response.getErr() == null ? "" : response.getErr())
+                    .build();
+        }
     }
 
     /**

src/main/proto/common/metadata.proto

@@ -0,0 +1,37 @@
+/*
+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";
+option go_package = "github.com/numaproj/numaflow-go/pkg/apis/proto/common";
+
+package common;
+
+// Metadata is the metadata of the message
+message Metadata {
+  // PreviousVertex is the name of the previous vertex
+  string previous_vertex = 1;
+  // SystemMetadata is the system metadata of the message
+  // Key of the map is the group name
+  map<string, KeyValueGroup> sys_metadata = 2;
+  // UserMetadata is the user metadata of the message
+  // Key of the map is the group name
+  map<string, KeyValueGroup> user_metadata = 3;
+}
+
+// KeyValueGroup is a group of key-value pairs for a given group.
+message KeyValueGroup {
+  map<string, bytes> key_value = 1;
+}

src/main/proto/sink/v1/sink.proto

@@ -4,6 +4,7 @@ option java_package = "io.numaproj.numaflow.sink.v1";
 
 import "google/protobuf/timestamp.proto";
 import "google/protobuf/empty.proto";
+import "common/metadata.proto";
 
 package sink.v1;
 
@@ -65,19 +66,29 @@ enum Status {
   SUCCESS = 0;
   FAILURE = 1;
   FALLBACK = 2;
+  SERVE = 3;
+  ON_SUCCESS = 4;
 }
 
 /**
  * SinkResponse is the individual response of each message written to the sink.
  */
 message SinkResponse {
   message Result {
+    message Message {
+      bytes value = 1;
+      repeated string keys = 2;
+      common.Metadata metadata = 3;
+    }
     // id is the ID of the message, can be used to uniquely identify the message.
     string id = 1;
     // status denotes the status of persisting to sink. It can be SUCCESS, FAILURE, or FALLBACK.
     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;
+    // on_success_msg is the message to be sent to on_success sink.
+    optional Message on_success_msg = 5;
   }
   repeated Result results = 1;
   optional Handshake handshake = 2;