Fix leader redirect for append v2

src/main/java/io/kurrent/dbclient/GrpcUtils.java

@@ -2,31 +2,57 @@
 
 import io.kurrent.dbclient.proto.shared.Shared;
 import io.kurrent.dbclient.proto.streams.StreamsOuterClass;
+import io.kurrentdb.protocol.v2.common.errors.NotLeaderNodeErrorDetails;
+import com.google.protobuf.Any;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.InvalidProtocolBufferException;
 import io.grpc.Metadata;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
+import io.grpc.protobuf.StatusProto;
 import io.grpc.stub.*;
 
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 
 final class GrpcUtils {
     static boolean handleNotLeaderError(Throwable t, CompletableFuture<?> dest) {
-        if (t instanceof StatusRuntimeException) {
-            StatusRuntimeException e = (StatusRuntimeException) t;
+        if (!(t instanceof StatusRuntimeException))
+            return false;
+
+        StatusRuntimeException e = (StatusRuntimeException) t;
+
+        try {
             Metadata trailers = e.getTrailers();
 
             if (trailers != null) {
                 String leaderHost = trailers.get(Metadata.Key.of("leader-endpoint-host", Metadata.ASCII_STRING_MARSHALLER));
                 String leaderPort = trailers.get(Metadata.Key.of("leader-endpoint-port", Metadata.ASCII_STRING_MARSHALLER));
 
                 if (leaderHost != null && leaderPort != null) {
-                    NotLeaderException reason = new NotLeaderException(leaderHost, Integer.parseInt(leaderPort));
-                    dest.completeExceptionally(reason);
+                    dest.completeExceptionally(new NotLeaderException(leaderHost, Integer.parseInt(leaderPort)));
+                    return true;
+                }
+            }
+
+            com.google.rpc.Status status = StatusProto.fromThrowable(e);
+
+            if (status != null) {
+                for (Any detail : status.getDetailsList()) {
+                    if (!detail.is(NotLeaderNodeErrorDetails.class))
+                        continue;
+
+                    NotLeaderNodeErrorDetails.NodeInfo leader =
+                            detail.unpack(NotLeaderNodeErrorDetails.class).getCurrentLeader();
+                    dest.completeExceptionally(new NotLeaderException(leader.getHost(), leader.getPort()));
                     return true;
                 }
             }
+        } catch (IllegalArgumentException ignored) {
+            dest.completeExceptionally(t);
+            return true;
+        } catch (InvalidProtocolBufferException ignored) {
+            return false;
         }
 
         return false;

src/test/java/io/kurrent/dbclient/LeaderRedirectUnitTest.java

@@ -1,13 +1,20 @@
 package io.kurrent.dbclient;
 
+import com.google.protobuf.Any;
+import com.google.rpc.ErrorInfo;
 import io.grpc.Metadata;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
+import io.grpc.protobuf.ProtoUtils;
+import io.grpc.protobuf.StatusProto;
+import io.kurrentdb.protocol.v2.common.errors.NotLeaderNodeErrorDetails;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
+import java.net.InetSocketAddress;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
@@ -85,4 +92,64 @@ public void onCancelled(Subscription subscription, Throwable exception) {
         Assertions.assertTrue(cancelErr.get() instanceof NotLeaderException,
             "listener should receive a NotLeaderException, got: " + cancelErr.get());
     }
+
+    @Test
+    public void handleNotLeaderErrorReadsNotLeaderNodeErrorDetails() {
+        com.google.rpc.Status status = com.google.rpc.Status.newBuilder()
+            .setCode(Status.Code.FAILED_PRECONDITION.value())
+            .setMessage("The server is not the leader node and cannot handle the request.")
+            .addDetails(Any.pack(ErrorInfo.newBuilder().setReason("SERVER_ERROR_NOT_LEADER_NODE").build()))
+            .addDetails(Any.pack(NotLeaderNodeErrorDetails.newBuilder()
+                .setCurrentLeader(NotLeaderNodeErrorDetails.NodeInfo.newBuilder().setHost("127.0.0.1").setPort(2113))
+                .build()))
+            .build();
+        CompletableFuture<Object> dest = new CompletableFuture<>();
+
+        Assertions.assertTrue(GrpcUtils.handleNotLeaderError(StatusProto.toStatusRuntimeException(status), dest));
+
+        ExecutionException error = Assertions.assertThrows(ExecutionException.class, dest::get);
+        Assertions.assertTrue(error.getCause() instanceof NotLeaderException,
+            "should complete with a NotLeaderException, got: " + error.getCause());
+        Assertions.assertEquals(new InetSocketAddress("127.0.0.1", 2113),
+            ((NotLeaderException) error.getCause()).getLeaderEndpoint());
+    }
+
+    @Test
+    public void handleNotLeaderErrorFailsWithOriginalErrorOnInvalidTrailerPort() {
+        Metadata trailers = new Metadata();
+        trailers.put(Metadata.Key.of("leader-endpoint-host", Metadata.ASCII_STRING_MARSHALLER), "127.0.0.1");
+        trailers.put(Metadata.Key.of("leader-endpoint-port", Metadata.ASCII_STRING_MARSHALLER), "abc");
+
+        assertFailsWithOriginalError(new StatusRuntimeException(Status.NOT_FOUND, trailers));
+    }
+
+    @Test
+    public void handleNotLeaderErrorFailsWithOriginalErrorOnOutOfRangeDetailsPort() {
+        com.google.rpc.Status status = com.google.rpc.Status.newBuilder()
+            .setCode(Status.Code.FAILED_PRECONDITION.value())
+            .addDetails(Any.pack(NotLeaderNodeErrorDetails.newBuilder()
+                .setCurrentLeader(NotLeaderNodeErrorDetails.NodeInfo.newBuilder().setHost("127.0.0.1").setPort(99999))
+                .build()))
+            .build();
+
+        assertFailsWithOriginalError(StatusProto.toStatusRuntimeException(status));
+    }
+
+    @Test
+    public void handleNotLeaderErrorFailsWithOriginalErrorOnStatusCodeMismatch() {
+        Metadata trailers = new Metadata();
+        trailers.put(Metadata.Key.of("grpc-status-details-bin", ProtoUtils.metadataMarshaller(com.google.rpc.Status.getDefaultInstance())),
+            com.google.rpc.Status.newBuilder().setCode(Status.Code.FAILED_PRECONDITION.value()).build());
+
+        assertFailsWithOriginalError(new StatusRuntimeException(Status.INTERNAL, trailers));
+    }
+
+    private static void assertFailsWithOriginalError(StatusRuntimeException original) {
+        CompletableFuture<Object> dest = new CompletableFuture<>();
+
+        Assertions.assertTrue(GrpcUtils.handleNotLeaderError(original, dest));
+
+        ExecutionException error = Assertions.assertThrows(ExecutionException.class, dest::get);
+        Assertions.assertSame(original, error.getCause());
+    }
 }