feat: introduce API to fail a message in udf/transformer

src/main/java/io/numaproj/numaflow/batchmapper/Message.java

@@ -8,6 +8,7 @@
 public class Message {
   private static final String[] DROP_TAGS = {"U+005C__DROP__"};
   private static final String[] NACK_TAGS = {"U+005C__NACK__"};
+  private static final String[] FAIL_TAGS = {"U+005C__FAIL__"};
   private final String[] keys;
   private final byte[] value;
   private final String[] tags;
@@ -69,4 +70,13 @@ public static Message toDrop() {
   public static Message toNack(NackOptions nackOptions) {
     return new Message(new byte[0], null, NACK_TAGS, nackOptions);
   }
+
+  /**
+   * creates a Message that marks the input message as failed, triggering a retry by numaflow-core.
+   *
+   * @return the Message which will be failed
+   */
+  public static Message toFail() {
+    return new Message(new byte[0], null, FAIL_TAGS);
+  }
 }

src/main/java/io/numaproj/numaflow/mapper/Message.java

@@ -13,6 +13,7 @@
 public class Message {
   private static final String[] DROP_TAGS = {"U+005C__DROP__"};
   private static final String[] NACK_TAGS = {"U+005C__NACK__"};
+  private static final String[] FAIL_TAGS = {"U+005C__FAIL__"};
   private final String[] keys;
   private final byte[] value;
   private final String[] tags;
@@ -89,4 +90,13 @@ public static Message toDrop() {
   public static Message toNack(NackOptions nackOptions) {
     return new Message(new byte[0], null, NACK_TAGS, null, nackOptions);
   }
+
+  /**
+   * creates a Message that marks the input message as failed, triggering a retry by numaflow-core.
+   *
+   * @return the Message which will be failed
+   */
+  public static Message toFail() {
+    return new Message(new byte[0], null, FAIL_TAGS, null);
+  }
 }

src/main/java/io/numaproj/numaflow/mapstreamer/Message.java

@@ -8,6 +8,7 @@
 public class Message {
   private static final String[] DROP_TAGS = {"U+005C__DROP__"};
   private static final String[] NACK_TAGS = {"U+005C__NACK__"};
+  private static final String[] FAIL_TAGS = {"U+005C__FAIL__"};
   private final String[] keys;
   private final byte[] value;
   private final String[] tags;
@@ -69,4 +70,13 @@ public static Message toDrop() {
   public static Message toNack(NackOptions nackOptions) {
     return new Message(new byte[0], null, NACK_TAGS, nackOptions);
   }
+
+  /**
+   * creates a Message that marks the input message as failed, triggering a retry by numaflow-core.
+   *
+   * @return the Message which will be failed
+   */
+  public static Message toFail() {
+    return new Message(new byte[0], null, FAIL_TAGS);
+  }
 }

src/main/java/io/numaproj/numaflow/sourcetransformer/Message.java

@@ -11,6 +11,7 @@
 public class Message {
   private static final String[] DROP_TAGS = {"U+005C__DROP__"};
   private static final String[] NACK_TAGS = {"U+005C__NACK__"};
+  private static final String[] FAIL_TAGS = {"U+005C__FAIL__"};
   private final String[] keys;
   private final byte[] value;
   private final Instant eventTime;
@@ -99,4 +100,16 @@ public static Message toDrop(Instant eventTime) {
   public static Message toNack(Instant eventTime, NackOptions nackOptions) {
     return new Message(new byte[0], eventTime, null, NACK_TAGS, null, nackOptions);
   }
+
+  /**
+   * creates a Message that marks the input message as failed, triggering a retry by numaflow-core.
+   * eventTime is required: even though the message is failed, it is considered processed for
+   * watermark purposes.
+   *
+   * @param eventTime message eventTime
+   * @return the Message which will be failed
+   */
+  public static Message toFail(Instant eventTime) {
+    return new Message(new byte[0], eventTime, null, FAIL_TAGS, null);
+  }
 }

src/test/java/io/numaproj/numaflow/batchmapper/ServerFailTest.java

@@ -0,0 +1,97 @@
+package io.numaproj.numaflow.batchmapper;
+
+import com.google.protobuf.ByteString;
+import io.grpc.ManagedChannel;
+import io.grpc.inprocess.InProcessChannelBuilder;
+import io.grpc.inprocess.InProcessServerBuilder;
+import io.grpc.stub.StreamObserver;
+import io.grpc.testing.GrpcCleanupRule;
+import io.numaproj.numaflow.map.v1.MapGrpc;
+import io.numaproj.numaflow.map.v1.MapOuterClass;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+public class ServerFailTest {
+    @Rule
+    public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+    private Server server;
+    private ManagedChannel inProcessChannel;
+
+    @Before
+    public void setUp() throws Exception {
+        String serverName = InProcessServerBuilder.generateName();
+        GRPCConfig grpcServerConfig = GRPCConfig.newBuilder()
+                .maxMessageSize(Constants.DEFAULT_MESSAGE_SIZE)
+                .socketPath(Constants.DEFAULT_SOCKET_PATH)
+                .infoFilePath("/tmp/numaflow-test-server-info)")
+                .build();
+        server = new Server(grpcServerConfig, new FailBatchMapFn(), null, serverName);
+        server.start();
+        inProcessChannel = grpcCleanup.register(
+                InProcessChannelBuilder.forName(serverName).directExecutor().build());
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        server.stop();
+    }
+
+    @Test
+    public void batchMapFail() {
+        // expect: handshake resp + 1 per-id response + 1 EOT = 3
+        BatchMapOutputStreamObserver outputStreamObserver = new BatchMapOutputStreamObserver(3);
+        StreamObserver<MapOuterClass.MapRequest> in = MapGrpc.newStub(inProcessChannel)
+                .mapFn(outputStreamObserver);
+        in.onNext(MapOuterClass.MapRequest.newBuilder()
+                .setHandshake(MapOuterClass.Handshake.newBuilder().setSot(true)).build());
+        in.onNext(MapOuterClass.MapRequest.newBuilder()
+                .setRequest(MapOuterClass.MapRequest.Request.newBuilder()
+                        .setValue(ByteString.copyFromUtf8("x")).addKeys("k").build())
+                .setId("id-1").build());
+        in.onNext(MapOuterClass.MapRequest.newBuilder()
+                .setStatus(MapOuterClass.TransmissionStatus.newBuilder().setEot(true)).build());
+        in.onCompleted();
+        try {
+            outputStreamObserver.done.get();
+        } catch (InterruptedException | ExecutionException e) {
+            fail("Error in getting done signal " + e.getMessage());
+        }
+        List<MapOuterClass.MapResponse> result = outputStreamObserver.getMapResponses();
+        MapOuterClass.MapResponse.Result r = result.stream()
+                .filter(resp -> resp.getResultsCount() > 0)
+                .findFirst().orElseThrow(() -> new AssertionError("no result")).getResults(0);
+        assertEquals(Arrays.asList("U+005C__FAIL__"), r.getTagsList());
+    }
+
+    private static class FailBatchMapFn extends BatchMapper {
+        @Override
+        public BatchResponses processMessage(DatumIterator datumStream) {
+            BatchResponses batchResponses = new BatchResponses();
+            while (true) {
+                Datum datum;
+                try {
+                    datum = datumStream.next();
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    continue;
+                }
+                if (datum == null) {
+                    break;
+                }
+                BatchResponse batchResponse = new BatchResponse(datum.getId());
+                batchResponse.append(Message.toFail());
+                batchResponses.append(batchResponse);
+            }
+            return batchResponses;
+        }
+    }
+}

src/test/java/io/numaproj/numaflow/mapper/ServerFailTest.java

@@ -0,0 +1,79 @@
+package io.numaproj.numaflow.mapper;
+
+import com.google.protobuf.ByteString;
+import io.grpc.ManagedChannel;
+import io.grpc.inprocess.InProcessChannelBuilder;
+import io.grpc.inprocess.InProcessServerBuilder;
+import io.grpc.testing.GrpcCleanupRule;
+import io.numaproj.numaflow.map.v1.MapGrpc;
+import io.numaproj.numaflow.map.v1.MapOuterClass;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+public class ServerFailTest {
+    @Rule
+    public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+    private Server server;
+    private ManagedChannel inProcessChannel;
+
+    @Before
+    public void setUp() throws Exception {
+        String serverName = InProcessServerBuilder.generateName();
+        GRPCConfig grpcServerConfig = GRPCConfig.newBuilder()
+                .maxMessageSize(Constants.DEFAULT_MESSAGE_SIZE)
+                .socketPath(Constants.DEFAULT_SOCKET_PATH)
+                .infoFilePath("/tmp/numaflow-test-server-info)")
+                .build();
+        server = new Server(grpcServerConfig, new FailMapFn(), null, serverName);
+        server.start();
+        inProcessChannel = grpcCleanup.register(
+                InProcessChannelBuilder.forName(serverName).directExecutor().build());
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        server.stop();
+    }
+
+    @Test
+    public void mapperFail() {
+        MapOuterClass.MapRequest handshake = MapOuterClass.MapRequest.newBuilder()
+                .setHandshake(MapOuterClass.Handshake.newBuilder().setSot(true)).build();
+        MapOuterClass.MapRequest inDatum = MapOuterClass.MapRequest.newBuilder()
+                .setRequest(MapOuterClass.MapRequest.Request.newBuilder()
+                        .setValue(ByteString.copyFromUtf8("x")).addKeys("k").build()).build();
+
+        MapOutputStreamObserver responseObserver = new MapOutputStreamObserver(2);
+        var stub = MapGrpc.newStub(inProcessChannel);
+        var requestStreamObserver = stub.mapFn(responseObserver);
+        requestStreamObserver.onNext(handshake);
+        requestStreamObserver.onNext(inDatum);
+        try {
+            responseObserver.done.get();
+        } catch (InterruptedException | ExecutionException e) {
+            fail("Error while waiting for response" + e.getMessage());
+        }
+        List<MapOuterClass.MapResponse> responses = responseObserver.getMapResponses().subList(1, 2);
+        MapOuterClass.MapResponse.Result r = responses.get(0).getResults(0);
+        assertEquals(Arrays.asList("U+005C__FAIL__"), r.getTagsList());
+        requestStreamObserver.onCompleted();
+    }
+
+    private static class FailMapFn extends Mapper {
+        @Override
+        public MessageList processMessage(String[] keys, Datum datum) {
+            return MessageList.newBuilder()
+                    .addMessage(Message.toFail())
+                    .build();
+        }
+    }
+}

src/test/java/io/numaproj/numaflow/mapstreamer/ServiceFailTest.java

@@ -0,0 +1,77 @@
+package io.numaproj.numaflow.mapstreamer;
+
+import com.google.protobuf.ByteString;
+import io.grpc.ManagedChannel;
+import io.grpc.inprocess.InProcessChannelBuilder;
+import io.grpc.inprocess.InProcessServerBuilder;
+import io.grpc.testing.GrpcCleanupRule;
+import io.numaproj.numaflow.map.v1.MapGrpc;
+import io.numaproj.numaflow.map.v1.MapOuterClass;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+public class ServiceFailTest {
+    @Rule
+    public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+    private Service service;
+    private ManagedChannel inProcessChannel;
+
+    @Before
+    public void setUp() throws Exception {
+        String serverName = InProcessServerBuilder.generateName();
+        CompletableFuture<Void> shutdownSignal = new CompletableFuture<>();
+        service = new Service(new FailMapStreamer(), shutdownSignal);
+        grpcCleanup.register(InProcessServerBuilder.forName(serverName).directExecutor()
+                .addService(service).build().start());
+        inProcessChannel = grpcCleanup.register(
+                InProcessChannelBuilder.forName(serverName).directExecutor().build());
+    }
+
+    @After
+    public void tearDown() {
+        inProcessChannel.shutdownNow();
+    }
+
+    @Test
+    public void mapStreamerFail() {
+        MapOuterClass.MapRequest handshake = MapOuterClass.MapRequest.newBuilder()
+                .setHandshake(MapOuterClass.Handshake.newBuilder().setSot(true)).build();
+        MapOuterClass.MapRequest inDatum = MapOuterClass.MapRequest.newBuilder()
+                .setRequest(MapOuterClass.MapRequest.Request.newBuilder()
+                        .setValue(ByteString.copyFromUtf8("x")).addKeys("k").build()).build();
+
+        // expect: handshake resp + 1 result + 1 EOT = 3 responses
+        MapStreamOutputStreamObserver responseObserver = new MapStreamOutputStreamObserver(3);
+        var stub = MapGrpc.newStub(inProcessChannel);
+        var requestStreamObserver = stub.mapFn(responseObserver);
+        requestStreamObserver.onNext(handshake);
+        requestStreamObserver.onNext(inDatum);
+        try {
+            responseObserver.done.get();
+        } catch (Exception e) {
+            fail("Error while waiting for response" + e.getMessage());
+        }
+        List<MapOuterClass.MapResponse> responses = responseObserver.getMapResponses();
+        MapOuterClass.MapResponse.Result r = responses.stream()
+                .filter(resp -> resp.getResultsCount() > 0)
+                .findFirst().orElseThrow(() -> new AssertionError("no result")).getResults(0);
+        assertEquals(Arrays.asList("U+005C__FAIL__"), r.getTagsList());
+        requestStreamObserver.onCompleted();
+    }
+
+    private static class FailMapStreamer extends MapStreamer {
+        @Override
+        public void processMessage(String[] keys, Datum datum, OutputObserver outputObserver) {
+            outputObserver.send(Message.toFail());
+        }
+    }
+}