feat: Add Retryable Configuration for GRPC Sink (using CEL)

build.gradle

@@ -103,6 +103,8 @@ dependencies {
     implementation 'org.apache.logging.log4j:log4j-core:2.20.0'
     implementation group: 'com.gotocompany', name: 'depot', version: '0.9.1'
     implementation group: 'com.networknt', name: 'json-schema-validator', version: '1.0.59' exclude group: 'org.slf4j'
+    implementation(enforcedPlatform("org.projectnessie.cel:cel-bom:0.4.4"))
+    implementation("org.projectnessie.cel:cel-tools")
 
     testImplementation group: 'junit', name: 'junit', version: '4.11'
     testImplementation 'org.hamcrest:hamcrest-all:1.3'
@@ -117,11 +119,11 @@ dependencies {
 protobuf {
     generatedFilesBaseDir = "$projectDir/src/generated"
     protoc {
-        artifact = "com.google.protobuf:protoc:3.1.0"
+        artifact = "com.google.protobuf:protoc:3.17.3"
     }
     plugins {
         grpc {
-            artifact = "io.grpc:protoc-gen-grpc-java:1.0.3"
+            artifact = "io.grpc:protoc-gen-grpc-java:1.60.1"
         }
     }
     generateProtoTasks {

src/main/java/com/gotocompany/firehose/config/GrpcSinkConfig.java

@@ -30,6 +30,10 @@ public interface GrpcSinkConfig extends AppConfig {
     @Config.Key("SINK_GRPC_ARG_DEADLINE_MS")
     Long getSinkGrpcArgDeadlineMS();
 
+    @Config.Key("SINK_GRPC_RESPONSE_RETRY_CEL_EXPRESSION")
+    @DefaultValue("")
+    String getSinkGrpcResponseRetryCELExpression();
+
     @Key("SINK_GRPC_METADATA")
     @DefaultValue("")
     @ConverterClass(GrpcMetadataConverter.class)

src/main/java/com/gotocompany/firehose/evaluator/GrpcResponseCELPayloadEvaluator.java

@@ -0,0 +1,66 @@
+package com.gotocompany.firehose.evaluator;
+
+import com.google.protobuf.Descriptors;
+import com.google.protobuf.DynamicMessage;
+import com.google.protobuf.Message;
+import lombok.extern.slf4j.Slf4j;
+import org.projectnessie.cel.checker.Decls;
+import org.projectnessie.cel.tools.Script;
+import org.projectnessie.cel.tools.ScriptCreateException;
+import org.projectnessie.cel.tools.ScriptException;
+import org.projectnessie.cel.tools.ScriptHost;
+
+import java.util.HashMap;
+import java.util.Map;
+
+@Slf4j
+public class GrpcResponseCELPayloadEvaluator implements PayloadEvaluator<Message> {
+
+    private final String celExpression;
+    private Script script;
+    private Descriptors.Descriptor descriptor;
+
+    public GrpcResponseCELPayloadEvaluator(Descriptors.Descriptor descriptor, String celExpression) {
+        this.celExpression = celExpression;
+        this.descriptor = descriptor;
+        this.script = buildScript(descriptor);
+    }
+
+    @Override
+    public boolean evaluate(Message payload) {
+        try {
+            Map<String, Object> arguments = new HashMap<>();
+            arguments.put(payload.getDescriptorForType().getFullName(), payload);
+            return getScript(payload.getDescriptorForType()).execute(Boolean.class, arguments);
+        } catch (ScriptException e) {
+            throw new IllegalArgumentException(
+                    "Failed to evaluate payload with CEL Expression with reason: " + e.getMessage(), e);
+        }
+    }
+
+    private Script getScript(Descriptors.Descriptor payloadDescriptor) throws ScriptCreateException {
+        if (!payloadDescriptor.equals(this.descriptor)) {
+            synchronized (this) {
+                if (!payloadDescriptor.equals(this.descriptor)) {
+                    this.script = buildScript(payloadDescriptor);
+                    this.descriptor = payloadDescriptor;
+                }
+            }
+        }
+        return this.script;
+    }
+
+    private Script buildScript(Descriptors.Descriptor payloadDescriptor) {
+        try {
+            log.info("Building new CEL Script");
+            return ScriptHost.newBuilder()
+                    .build()
+                    .buildScript(this.celExpression)
+                    .withDeclarations(Decls.newVar(payloadDescriptor.getFullName(), Decls.newObjectType(payloadDescriptor.getFullName())))
+                    .withTypes(DynamicMessage.newBuilder(payloadDescriptor).getDefaultInstanceForType())
+                    .build();
+        } catch (ScriptCreateException e) {
+            throw new IllegalArgumentException("Failed to build CEL Script due to : " + e.getMessage(), e);
+        }
+    }
+}

src/main/java/com/gotocompany/firehose/evaluator/PayloadEvaluator.java

@@ -0,0 +1,5 @@
+package com.gotocompany.firehose.evaluator;
+
+public interface PayloadEvaluator<T> {
+    boolean evaluate(T payload);
+}

src/main/java/com/gotocompany/firehose/sink/grpc/GrpcSink.java

@@ -1,14 +1,19 @@
 package com.gotocompany.firehose.sink.grpc;
 
 
-
+import com.gotocompany.depot.error.ErrorInfo;
+import com.gotocompany.depot.error.ErrorType;
+import com.gotocompany.firehose.config.GrpcSinkConfig;
+import com.gotocompany.firehose.evaluator.GrpcResponseCELPayloadEvaluator;
+import com.gotocompany.firehose.exception.DefaultException;
 import com.gotocompany.firehose.exception.DeserializerException;
 import com.gotocompany.firehose.message.Message;
 import com.gotocompany.firehose.metrics.FirehoseInstrumentation;
 import com.gotocompany.firehose.sink.AbstractSink;
 import com.gotocompany.firehose.sink.grpc.client.GrpcClient;
 import com.google.protobuf.DynamicMessage;
 import com.gotocompany.stencil.client.StencilClient;
+import org.apache.commons.lang3.StringUtils;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -21,13 +26,24 @@
 public class GrpcSink extends AbstractSink {
 
     private final GrpcClient grpcClient;
+    private final StencilClient stencilClient;
+    private final GrpcSinkConfig grpcSinkConfig;
     private List<Message> messages;
-    private StencilClient stencilClient;
+    private GrpcResponseCELPayloadEvaluator retryEvaluator;
 
-    public GrpcSink(FirehoseInstrumentation firehoseInstrumentation, GrpcClient grpcClient, StencilClient stencilClient) {
+    public GrpcSink(FirehoseInstrumentation firehoseInstrumentation,
+                    GrpcClient grpcClient,
+                    StencilClient stencilClient,
+                    GrpcSinkConfig grpcSinkConfig) {
         super(firehoseInstrumentation, "grpc");
         this.grpcClient = grpcClient;
         this.stencilClient = stencilClient;
+        this.grpcSinkConfig = grpcSinkConfig;
+        if (StringUtils.isNotBlank(grpcSinkConfig.getSinkGrpcResponseRetryCELExpression())) {
+            this.retryEvaluator = new GrpcResponseCELPayloadEvaluator(
+                    stencilClient.get(grpcSinkConfig.getSinkGrpcResponseSchemaProtoClass()),
+                    grpcSinkConfig.getSinkGrpcResponseRetryCELExpression());
+        }
     }
 
     @Override
@@ -44,6 +60,9 @@ protected List<Message> execute() throws Exception {
                 getFirehoseInstrumentation().logWarn("Grpc Service returned error");
                 failedMessages.add(message);
             }
+            if (StringUtils.isNotBlank(grpcSinkConfig.getSinkGrpcResponseRetryCELExpression())) {
+                setRetryEvaluatorErrorInfo(message, response);
+            }
         }
         getFirehoseInstrumentation().logDebug("Failed messages count: {}", failedMessages.size());
         return failedMessages;
@@ -60,4 +79,13 @@ public void close() throws IOException {
         this.messages = new ArrayList<>();
         stencilClient.close();
     }
+
+    private void setRetryEvaluatorErrorInfo(Message message, DynamicMessage dynamicMessage) {
+        boolean eligibleToRetry = retryEvaluator.evaluate(dynamicMessage);
+        if (eligibleToRetry) {
+            message.setErrorInfo(new ErrorInfo(new DefaultException("Retryable gRPC Error"), ErrorType.SINK_RETRYABLE_ERROR));
+            return;
+        }
+        message.setErrorInfo(new ErrorInfo(new DefaultException("Non Retryable gRPC Error"), ErrorType.SINK_NON_RETRYABLE_ERROR));
+    }
 }

src/main/java/com/gotocompany/firehose/sink/grpc/GrpcSinkFactory.java

@@ -38,7 +38,7 @@ public static AbstractSink create(Map<String, String> configuration, StatsDRepor
         GrpcClient grpcClient = new GrpcClient(new FirehoseInstrumentation(statsDReporter, GrpcClient.class), grpcConfig, managedChannel, stencilClient);
         firehoseInstrumentation.logInfo("GRPC connection established");
 
-        return new GrpcSink(new FirehoseInstrumentation(statsDReporter, GrpcSink.class), grpcClient, stencilClient);
+        return new GrpcSink(new FirehoseInstrumentation(statsDReporter, GrpcSink.class), grpcClient, stencilClient, grpcConfig);
     }
 
 }

src/test/java/com/gotocompany/firehose/evaluator/GrpcResponseCELPayloadEvaluatorTest.java

@@ -0,0 +1,94 @@
+package com.gotocompany.firehose.evaluator;
+
+import com.google.protobuf.DescriptorProtos;
+import com.google.protobuf.Descriptors;
+import com.google.protobuf.DynamicMessage;
+import com.google.protobuf.Message;
+import com.gotocompany.firehose.consumer.GenericError;
+import com.gotocompany.firehose.consumer.GenericResponse;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.jupiter.api.Assertions;
+
+public class GrpcResponseCELPayloadEvaluatorTest {
+
+    private static final String CEL_EXPRESSION = "GenericResponse.success == false && GenericResponse.errors.exists(e, e.code == \"400\")";
+    private PayloadEvaluator<Message> grpcPayloadEvaluator;
+
+    @Before
+    public void setup() {
+        this.grpcPayloadEvaluator = new GrpcResponseCELPayloadEvaluator(GenericResponse.getDescriptor(), CEL_EXPRESSION);
+    }
+
+    @Test
+    public void shouldEvaluateResponseToTrueWhenCELExpressionMatchesPayload() {
+        GenericResponse genericResponse = GenericResponse.newBuilder()
+                .setSuccess(false)
+                .setDetail("Detail Message")
+                .addErrors(GenericError.newBuilder()
+                        .setCode("400")
+                        .setEntity("GoFin")
+                        .setCause("Unknown")
+                        .build())
+                .build();
+
+        boolean result = grpcPayloadEvaluator.evaluate(genericResponse);
+
+        Assertions.assertTrue(result);
+    }
+
+    @Test
+    public void shouldEvaluateResponseToFalseWhenCELExpressionDoesntMatchPayload() {
+        GenericResponse genericResponse = GenericResponse.newBuilder()
+                .setSuccess(false)
+                .setDetail("Detail Message")
+                .addErrors(GenericError.newBuilder()
+                        .setCode("50000")
+                        .setEntity("GoFin")
+                        .setCause("Unknown")
+                        .build())
+                .build();
+
+        boolean result = grpcPayloadEvaluator.evaluate(genericResponse);
+
+        Assertions.assertFalse(result);
+    }
+
+    @Test
+    public void shouldEvaluateResponseWhenDescriptorUpdated() throws Descriptors.DescriptorValidationException {
+        Descriptors.Descriptor baseDescriptor = GenericResponse.getDescriptor();
+        DescriptorProtos.FieldDescriptorProto newFieldProto = DescriptorProtos.FieldDescriptorProto.newBuilder()
+                .setName("new_field")
+                .setNumber(4)
+                .setType(DescriptorProtos.FieldDescriptorProto.Type.TYPE_STRING)
+                .build();
+        DescriptorProtos.DescriptorProto newDescriptorProto = baseDescriptor.toProto().toBuilder()
+                .addField(newFieldProto)
+                .build();
+        Descriptors.FileDescriptor newFileDescriptor = Descriptors.FileDescriptor.buildFrom(DescriptorProtos.FileDescriptorProto
+                .newBuilder()
+                .setName("new.proto")
+                .addMessageType(newDescriptorProto)
+                .addMessageType(GenericError.getDescriptor().toProto())
+                .build(), new Descriptors.FileDescriptor[]{});
+        Descriptors.Descriptor genericResponseDescriptor = newFileDescriptor.findMessageTypeByName("GenericResponse");
+        DynamicMessage dynamicMessage = DynamicMessage.newBuilder(genericResponseDescriptor)
+                .setField(genericResponseDescriptor.findFieldByName("success"), false)
+                .setField(genericResponseDescriptor.findFieldByName("new_field"), "new_field")
+                .build();
+
+        boolean result = grpcPayloadEvaluator.evaluate(dynamicMessage);
+
+        Assertions.assertFalse(result);
+    }
+
+    @Test
+    public void shouldThrowIllegalArgumentExceptionWhenPayloadNotMatchingDescriptor() {
+        Assertions.assertThrows(IllegalArgumentException.class,
+                () -> grpcPayloadEvaluator.evaluate(GenericError.newBuilder()
+                        .setCause("Unknown")
+                        .setCode("500")
+                        .setEntity("GoFin")
+                        .build()));
+    }
+}

src/test/java/com/gotocompany/firehose/sink/grpc/GrpcSinkTest.java

@@ -1,5 +1,9 @@
 package com.gotocompany.firehose.sink.grpc;
 
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.gotocompany.firehose.config.GrpcSinkConfig;
+import com.gotocompany.firehose.consumer.GenericError;
+import com.gotocompany.firehose.consumer.GenericResponse;
 import com.gotocompany.firehose.exception.DeserializerException;
 import com.gotocompany.firehose.message.Message;
 import com.gotocompany.firehose.metrics.FirehoseInstrumentation;
@@ -12,6 +16,7 @@
 import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.jupiter.api.Assertions;
 import org.mockito.Mock;
 
 import java.io.IOException;
@@ -40,10 +45,13 @@ public class GrpcSinkTest {
     @Mock
     private FirehoseInstrumentation firehoseInstrumentation;
 
+    @Mock
+    private GrpcSinkConfig grpcSinkConfig;
+
     @Before
     public void setUp() {
         initMocks(this);
-        sink = new GrpcSink(firehoseInstrumentation, grpcClient, stencilClient);
+        sink = new GrpcSink(firehoseInstrumentation, grpcClient, stencilClient, grpcSinkConfig);
     }
 
     @Test
@@ -85,17 +93,77 @@ public void shouldReturnBackListOfFailedMessages() throws IOException, Deseriali
 
     @Test
     public void shouldCloseStencilClient() throws IOException {
-        sink = new GrpcSink(firehoseInstrumentation, grpcClient, stencilClient);
+        sink = new GrpcSink(firehoseInstrumentation, grpcClient, stencilClient, grpcSinkConfig);
 
         sink.close();
         verify(stencilClient, times(1)).close();
     }
 
     @Test
     public void shouldLogWhenClosingConnection() throws IOException {
-        sink = new GrpcSink(firehoseInstrumentation, grpcClient, stencilClient);
+        sink = new GrpcSink(firehoseInstrumentation, grpcClient, stencilClient, grpcSinkConfig);
 
         sink.close();
         verify(firehoseInstrumentation, times(1)).logInfo("GRPC connection closing");
     }
+
+    @Test
+    public void shouldReturnFailedMessagesWithRetryableErrorsWhenCELExpressionMatches() throws InvalidProtocolBufferException {
+        Message payload = new Message(new byte[]{}, new byte[]{}, "topic", 0, 1);
+        GenericResponse response = GenericResponse.newBuilder()
+                .setSuccess(false)
+                .setDetail("detail")
+                .addErrors(GenericError.newBuilder()
+                        .setCode("4000")
+                        .setCause("cause")
+                        .setEntity("gtf")
+                        .build())
+                .build();
+        DynamicMessage dynamicMessage = DynamicMessage.parseFrom(
+                response.getDescriptorForType(),
+                response.toByteArray()
+        );
+        when(grpcSinkConfig.getSinkGrpcResponseRetryCELExpression())
+                .thenReturn("GenericResponse.success == false && GenericResponse.errors.exists(e, e.code == \"4000\")");
+        when(grpcClient.execute(any(), any()))
+                .thenReturn(dynamicMessage);
+        when(stencilClient.get(any()))
+                .thenReturn(GenericResponse.getDescriptor());
+        sink = new GrpcSink(firehoseInstrumentation, grpcClient, stencilClient, grpcSinkConfig);
+
+        List<Message> result = sink.pushMessage(Collections.singletonList(payload));
+
+        Assertions.assertEquals(1, result.size());
+        Assertions.assertEquals(result.get(0).getErrorInfo().getErrorType(), ErrorType.SINK_RETRYABLE_ERROR);
+    }
+
+    @Test
+    public void shouldReturnFailedMessagesWithNonRetryableErrorsWhenCELExpressionDoesntMatch() throws InvalidProtocolBufferException {
+        Message payload = new Message(new byte[]{}, new byte[]{}, "topic", 0, 1);
+        GenericResponse response = GenericResponse.newBuilder()
+                .setSuccess(false)
+                .setDetail("detail")
+                .addErrors(GenericError.newBuilder()
+                        .setCode("4000")
+                        .setCause("cause")
+                        .setEntity("gtf")
+                        .build())
+                .build();
+        DynamicMessage dynamicMessage = DynamicMessage.parseFrom(
+                response.getDescriptorForType(),
+                response.toByteArray()
+        );
+        when(grpcSinkConfig.getSinkGrpcResponseRetryCELExpression())
+                .thenReturn("GenericResponse.success == false && GenericResponse.errors.exists(e, e.code == \"5000\")");
+        when(grpcClient.execute(any(), any()))
+                .thenReturn(dynamicMessage);
+        when(stencilClient.get(any()))
+                .thenReturn(GenericResponse.getDescriptor());
+        sink = new GrpcSink(firehoseInstrumentation, grpcClient, stencilClient, grpcSinkConfig);
+
+        List<Message> result = sink.pushMessage(Collections.singletonList(payload));
+
+        Assertions.assertEquals(1, result.size());
+        Assertions.assertEquals(result.get(0).getErrorInfo().getErrorType(), ErrorType.SINK_NON_RETRYABLE_ERROR);
+    }
 }