[FLINK-35280] PR comments

Author: mateczagany

flink-connector-hbase-2.6/src/main/java/org/apache/flink/connector/hbase2/sink/HBaseDynamicTableSink.java

@@ -25,9 +25,9 @@
 import org.apache.flink.connector.base.table.sink.AsyncDynamicTableSinkBuilder;
 import org.apache.flink.connector.hbase.sink.HBaseSinkBuilder;
 import org.apache.flink.connector.hbase.sink.RowDataToMutationElementConverter;
-import org.apache.flink.connector.hbase.sink.SerializableMutation;
 import org.apache.flink.connector.hbase.sink.WritableMetadata;
 import org.apache.flink.connector.hbase.util.HBaseTableSchema;
+import org.apache.flink.connector.hbase.util.SerializableMutation;
 import org.apache.flink.table.connector.ChangelogMode;
 import org.apache.flink.table.connector.sink.DynamicTableSink;
 import org.apache.flink.table.connector.sink.SinkV2Provider;

flink-connector-hbase-base/src/main/java/org/apache/flink/connector/hbase/sink/HBaseSink.java

@@ -7,6 +7,7 @@
 import org.apache.flink.connector.base.sink.writer.BufferedRequestState;
 import org.apache.flink.connector.base.sink.writer.ElementConverter;
 import org.apache.flink.connector.hbase.util.HBaseConfigurationUtil;
+import org.apache.flink.connector.hbase.util.SerializableMutation;
 import org.apache.flink.core.io.SimpleVersionedSerializer;
 
 import org.apache.hadoop.conf.Configuration;

flink-connector-hbase-base/src/main/java/org/apache/flink/connector/hbase/sink/HBaseSinkBuilder.java

@@ -3,6 +3,8 @@
 import org.apache.flink.annotation.PublicEvolving;
 import org.apache.flink.connector.base.sink.AsyncSinkBaseBuilder;
 import org.apache.flink.connector.base.sink.writer.ElementConverter;
+import org.apache.flink.connector.hbase.util.SerializableMutation;
+import org.apache.flink.connector.hbase.util.WrappedElementConverter;
 
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.client.Mutation;

flink-connector-hbase-base/src/main/java/org/apache/flink/connector/hbase/sink/HBaseStateSerializer.java

@@ -2,7 +2,13 @@
 
 import org.apache.flink.annotation.Internal;
 import org.apache.flink.connector.base.sink.writer.AsyncSinkWriterStateSerializer;
-import org.apache.flink.connector.hbase.util.HBaseMutationSerialization;
+import org.apache.flink.connector.hbase.util.SerializableMutation;
+
+import org.apache.hadoop.hbase.client.Delete;
+import org.apache.hadoop.hbase.client.Mutation;
+import org.apache.hadoop.hbase.client.Put;
+import org.apache.hadoop.hbase.protobuf.ProtobufUtil;
+import org.apache.hadoop.hbase.protobuf.generated.ClientProtos;
 
 import java.io.DataInputStream;
 import java.io.DataOutputStream;
@@ -14,13 +20,30 @@ public class HBaseStateSerializer extends AsyncSinkWriterStateSerializer<Seriali
     @Override
     protected void serializeRequestToStream(SerializableMutation request, DataOutputStream out)
             throws IOException {
-        HBaseMutationSerialization.serialize(request.get(), out);
+        Mutation mutation = request.get();
+        ClientProtos.MutationProto.MutationType type;
+
+        if (mutation instanceof Put) {
+            type = ClientProtos.MutationProto.MutationType.PUT;
+        } else if (mutation instanceof Delete) {
+            type = ClientProtos.MutationProto.MutationType.DELETE;
+        } else {
+            throw new IllegalArgumentException(
+                    String.format(
+                            "Unknown HBase mutation type, cannot serialize: %s",
+                            mutation.getClass()));
+        }
+
+        ClientProtos.MutationProto proto = ProtobufUtil.toMutation(type, mutation);
+        proto.writeDelimitedTo(out);
     }
 
     @Override
     protected SerializableMutation deserializeRequestFromStream(
             long requestSize, DataInputStream in) throws IOException {
-        return new SerializableMutation(HBaseMutationSerialization.deserialize(in));
+        ClientProtos.MutationProto proto = ClientProtos.MutationProto.parseDelimitedFrom(in);
+        Mutation mutation = ProtobufUtil.toMutation(proto);
+        return new SerializableMutation(mutation);
     }
 
     @Override

flink-connector-hbase-base/src/main/java/org/apache/flink/connector/hbase/sink/HBaseWriter.java

@@ -12,6 +12,7 @@
 import org.apache.flink.connector.base.sink.writer.strategy.CongestionControlRateLimitingStrategy;
 import org.apache.flink.connector.base.sink.writer.strategy.RateLimitingStrategy;
 import org.apache.flink.connector.hbase.util.HBaseConfigurationUtil;
+import org.apache.flink.connector.hbase.util.SerializableMutation;
 import org.apache.flink.metrics.Counter;
 import org.apache.flink.metrics.groups.SinkWriterMetricGroup;
 import org.apache.flink.util.StringUtils;

flink-connector-hbase-base/src/main/java/org/apache/flink/connector/hbase/sink/HBaseWriterAsyncHandler.java

@@ -2,8 +2,8 @@
 
 import org.apache.flink.annotation.Internal;
 import org.apache.flink.connector.base.sink.writer.ResultHandler;
+import org.apache.flink.connector.hbase.util.SerializableMutation;
 import org.apache.flink.metrics.Counter;
-import org.apache.flink.util.Preconditions;
 
 import org.apache.hadoop.hbase.DoNotRetryIOException;
 import org.apache.hadoop.hbase.client.Mutation;
@@ -13,10 +13,14 @@
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.apache.flink.util.Preconditions.checkArgument;
 
 /**
  * This class is responsible for managing the async calls to HBase and managing the {@link
@@ -47,34 +51,47 @@ public void handleWriteFutures(
             List<CompletableFuture<Mutation>> futures,
             List<SerializableMutation> processedMutationsInOrder,
             ResultHandler<SerializableMutation> resultHandler) {
-        Preconditions.checkArgument(
+        checkArgument(
                 futures.size() == processedMutationsInOrder.size(),
                 "Different number of HBase futures was supplied than mutations.");
 
-        ConcurrentLinkedQueue<FailedMutation> failedMutations = new ConcurrentLinkedQueue<>();
+        Queue<FailedMutation> failedMutations = new ConcurrentLinkedQueue<>();
 
         // Handle each future separately and store failures.
-        CompletableFuture<?>[] handledFutures = new CompletableFuture[futures.size()];
-        for (int i = 0; i < futures.size(); i++) {
-            final int index = i;
-            handledFutures[index] =
-                    futures.get(index)
-                            .exceptionally(
-                                    throwable -> {
-                                        failedMutations.add(
-                                                new FailedMutation(
-                                                        processedMutationsInOrder.get(index),
-                                                        throwable));
-                                        return null;
-                                    });
-        }
+        CompletableFuture<?>[] handledFutures =
+                IntStream.range(0, futures.size())
+                        .mapToObj(
+                                i ->
+                                        addFailedCallback(
+                                                futures.get(i),
+                                                failedMutations,
+                                                processedMutationsInOrder.get(i)))
+                        .toArray(CompletableFuture[]::new);
 
         // Exceptions are already handled here, so it's safe to use `thenRun()`.
         CompletableFuture.allOf(handledFutures)
-                .thenRun(
-                        () -> {
-                            handleFailedRequests(failedMutations, resultHandler);
-                        });
+                .thenRun(() -> handleFailedRequests(failedMutations, resultHandler));
+    }
+
+    /**
+     * For a {@link CompletableFuture} HBase write future, return a new future that adds a {@link
+     * FailedMutation} to the failedMutations queue to mark that write as failed.
+     *
+     * @param cf future of a mutation write
+     * @param failedMutations queue of failed mutations to add a new entry to upon failure
+     * @param mutation the actual mutation that's being written by the future
+     * @return new future with the exception callback
+     */
+    private CompletableFuture<Mutation> addFailedCallback(
+            CompletableFuture<Mutation> cf,
+            Queue<FailedMutation> failedMutations,
+            SerializableMutation mutation) {
+
+        return cf.exceptionally(
+                throwable -> {
+                    failedMutations.add(new FailedMutation(mutation, throwable));
+                    return null;
+                });
     }
 
     /**

flink-connector-hbase-base/src/main/java/org/apache/flink/connector/hbase/util/HBaseMutationSerialization.java

@@ -1,6 +1,8 @@
-package org.apache.flink.connector.hbase.sink;
+package org.apache.flink.connector.hbase.util;
 
 import org.apache.flink.annotation.Internal;
+import org.apache.flink.connector.hbase.sink.HBaseSink;
+import org.apache.flink.connector.hbase.sink.HBaseWriter;
 
 import org.apache.hadoop.hbase.client.Mutation;
 
@@ -14,7 +16,7 @@
 public class SerializableMutation implements Serializable {
     private static final long serialVersionUID = 1L;
 
-    private transient Mutation mutation;
+    private final transient Mutation mutation;
 
     public SerializableMutation(Mutation mutation) {
         this.mutation = mutation;

flink-connector-hbase-base/src/main/java/org/apache/flink/connector/hbase/sink/SerializableMutation.java renamed to flink-connector-hbase-base/src/main/java/org/apache/flink/connector/hbase/util/SerializableMutation.java

@@ -1,4 +1,4 @@
-package org.apache.flink.connector.hbase.sink;
+package org.apache.flink.connector.hbase.util;
 
 import org.apache.flink.annotation.Internal;
 import org.apache.flink.api.connector.sink2.Sink;

flink-connector-hbase-base/src/main/java/org/apache/flink/connector/hbase/sink/WrappedElementConverter.java renamed to flink-connector-hbase-base/src/main/java/org/apache/flink/connector/hbase/util/WrappedElementConverter.java

@@ -0,0 +1,69 @@
+package org.apache.flink.connector.hbase.sink;
+
+import org.apache.flink.connector.hbase.util.SerializableMutation;
+
+import org.apache.hadoop.hbase.client.Append;
+import org.apache.hadoop.hbase.client.Delete;
+import org.apache.hadoop.hbase.client.Mutation;
+import org.apache.hadoop.hbase.client.Put;
+import org.junit.jupiter.api.Test;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+class HBaseStateSerializerTest {
+
+    private final HBaseStateSerializer serializer = new HBaseStateSerializer();
+
+    @Test
+    public void testSerdePut() throws IOException {
+        SerializableMutation mut = new SerializableMutation(new Put(new byte[] {0x00, 0x01}, 1001));
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        DataOutputStream out = new DataOutputStream(baos);
+
+        serializer.serializeRequestToStream(mut, out);
+
+        ByteArrayInputStream bais = new ByteArrayInputStream(baos.toByteArray());
+        DataInputStream in = new DataInputStream(bais);
+
+        Mutation mutationResult = serializer.deserializeRequestFromStream(0, in).get();
+
+        assertThat(mutationResult.getRow()).isEqualTo(new byte[] {0x00, 0x01});
+        assertThat(mutationResult.getTimestamp()).isEqualTo(1001);
+    }
+
+    @Test
+    public void testSerdeDelete() throws IOException {
+        SerializableMutation mut =
+                new SerializableMutation(new Delete(new byte[] {0x00, 0x02}, 1002));
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        DataOutputStream out = new DataOutputStream(baos);
+
+        serializer.serializeRequestToStream(mut, out);
+
+        ByteArrayInputStream bais = new ByteArrayInputStream(baos.toByteArray());
+        DataInputStream in = new DataInputStream(bais);
+
+        Mutation mutationResult = serializer.deserializeRequestFromStream(0, in).get();
+
+        assertThat(mutationResult.getRow()).isEqualTo(new byte[] {0x00, 0x02});
+        assertThat(mutationResult.getTimestamp()).isEqualTo(1002);
+    }
+
+    @Test
+    public void testSerializationUnsupportedMutationType() {
+        SerializableMutation mut = new SerializableMutation(new Append(new byte[] {0x00, 0x02}));
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        DataOutputStream out = new DataOutputStream(baos);
+
+        assertThrows(
+                IllegalArgumentException.class,
+                () -> serializer.serializeRequestToStream(mut, out));
+    }
+}