HDDS-9370. Avoid copying ByteString in ByteStringCodec.

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/utils/db/BooleanCodec.java

@@ -18,7 +18,6 @@
 package org.apache.hadoop.hdds.utils.db;
 
 import javax.annotation.Nonnull;
-import java.util.function.IntFunction;
 
 /**
  * Codec to serialize/deserialize {@link Boolean}.
@@ -44,7 +43,7 @@ public boolean supportCodecBuffer() {
 
   @Override
   public CodecBuffer toCodecBuffer(Boolean object,
-      IntFunction<CodecBuffer> allocator) {
+      CodecBuffer.Allocator allocator) {
     return allocator.apply(1).put(TRUE);
   }
 

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/utils/db/Codec.java

@@ -19,7 +19,6 @@
 
 import javax.annotation.Nonnull;
 import java.io.IOException;
-import java.util.function.IntFunction;
 
 /**
  * Codec interface to serialize/deserialize objects to/from bytes.
@@ -34,7 +33,7 @@ public interface Codec<T> {
   /**
    * Does this {@link Codec} support the {@link CodecBuffer} methods?
    * If this method returns true, this class must implement both
-   * {@link #toCodecBuffer(Object, IntFunction)} and
+   * {@link #toCodecBuffer(Object, CodecBuffer.Allocator)} and
    * {@link #fromCodecBuffer(CodecBuffer)}.
    *
    * @return ture iff this class supports the {@link CodecBuffer} methods.
@@ -59,7 +58,7 @@ default int getSerializedSizeUpperBound(T object) {
    * @return a buffer storing the serialized bytes.
    */
   default CodecBuffer toCodecBuffer(@Nonnull T object,
-      IntFunction<CodecBuffer> allocator) throws IOException {
+      CodecBuffer.Allocator allocator) throws IOException {
     throw new UnsupportedOperationException();
   }
 
@@ -71,7 +70,7 @@ default CodecBuffer toCodecBuffer(@Nonnull T object,
    */
   default CodecBuffer toDirectCodecBuffer(@Nonnull T object)
       throws IOException {
-    return toCodecBuffer(object, CodecBuffer::allocateDirect);
+    return toCodecBuffer(object, CodecBuffer.Allocator.getDirect());
   }
 
   /**
@@ -82,7 +81,7 @@ default CodecBuffer toDirectCodecBuffer(@Nonnull T object)
    */
   default CodecBuffer toHeapCodecBuffer(@Nonnull T object)
       throws IOException {
-    return toCodecBuffer(object, CodecBuffer::allocateHeap);
+    return toCodecBuffer(object, CodecBuffer.Allocator.getHeap());
   }
 
   /**

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/utils/db/CodecBuffer.java

@@ -17,11 +17,13 @@
  */
 package org.apache.hadoop.hdds.utils.db;
 
+import com.google.protobuf.ByteString;
 import org.apache.hadoop.hdds.StringUtils;
 import org.apache.ratis.thirdparty.io.netty.buffer.ByteBuf;
 import org.apache.ratis.thirdparty.io.netty.buffer.ByteBufAllocator;
 import org.apache.ratis.thirdparty.io.netty.buffer.ByteBufInputStream;
 import org.apache.ratis.thirdparty.io.netty.buffer.ByteBufOutputStream;
+import org.apache.ratis.thirdparty.io.netty.buffer.EmptyByteBuf;
 import org.apache.ratis.thirdparty.io.netty.buffer.PooledByteBufAllocator;
 import org.apache.ratis.thirdparty.io.netty.buffer.Unpooled;
 import org.apache.ratis.util.MemoizedSupplier;
@@ -37,7 +39,7 @@
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.function.Function;
+import java.util.function.BiFunction;
 import java.util.function.IntFunction;
 import java.util.function.ToIntFunction;
 
@@ -53,22 +55,26 @@ public class CodecBuffer implements AutoCloseable {
 
   /** To create {@link CodecBuffer} instances. */
   private static class Factory {
-    private static volatile Function<ByteBuf, CodecBuffer> constructor
+    private static volatile BiFunction<ByteBuf, Object, CodecBuffer> constructor
         = CodecBuffer::new;
-    static void set(Function<ByteBuf, CodecBuffer> f) {
+    static void set(BiFunction<ByteBuf, Object, CodecBuffer> f) {
       constructor = f;
       LOG.info("Successfully set constructor to " + f);
     }
 
     static CodecBuffer newCodecBuffer(ByteBuf buf) {
-      return constructor.apply(buf);
+      return newCodecBuffer(buf, null);
+    }
+
+    static CodecBuffer newCodecBuffer(ByteBuf buf, Object wrapped) {
+      return constructor.apply(buf, wrapped);
     }
   }
 
   /** To detect buffer leak. */
   private static class LeakDetector {
-    static CodecBuffer newCodecBuffer(ByteBuf buf) {
-      return new CodecBuffer(buf) {
+    static CodecBuffer newCodecBuffer(ByteBuf buf, Object wrapped) {
+      return new CodecBuffer(buf, wrapped) {
         @Override
         protected void finalize() {
           detectLeaks();
@@ -126,6 +132,51 @@ public void increase(int required) {
       ? POOL.heapBuffer(c, c)   // allocate exact size
       : POOL.heapBuffer(-c);    // allocate a resizable buffer
 
+  private static final CodecBuffer EMPTY_BUFFER = new CodecBuffer(
+      new EmptyByteBuf(POOL), null);
+
+  public static CodecBuffer getEmptyBuffer() {
+    return EMPTY_BUFFER;
+  }
+
+  /** To allocate {@link CodecBuffer} objects. */
+  public interface Allocator extends IntFunction<CodecBuffer> {
+    Allocator DIRECT = new Allocator() {
+      @Override
+      public CodecBuffer apply(int capacity) {
+        return allocate(capacity, POOL_DIRECT);
+      }
+
+      @Override
+      public boolean isDirect() {
+        return true;
+      }
+    };
+
+    static Allocator getDirect() {
+      return DIRECT;
+    }
+
+    Allocator HEAP = new Allocator() {
+      @Override
+      public CodecBuffer apply(int capacity) {
+        return allocate(capacity, POOL_HEAP);
+      }
+
+      @Override
+      public boolean isDirect() {
+        return false;
+      }
+    };
+
+    static Allocator getHeap() {
+      return HEAP;
+    }
+
+    /** Does this object allocate direct buffers? */
+    boolean isDirect();
+  }
+
   private final StackTraceElement[] elements;
 
   /**
@@ -162,7 +213,13 @@ public static CodecBuffer allocateHeap(int capacity) {
 
   /** Wrap the given array. */
   public static CodecBuffer wrap(byte[] array) {
-    return Factory.newCodecBuffer(Unpooled.wrappedBuffer(array));
+    return Factory.newCodecBuffer(Unpooled.wrappedBuffer(array), array);
+  }
+
+  /** Wrap the given {@link ByteString}. */
+  public static CodecBuffer wrap(ByteString bytes) {
+    return Factory.newCodecBuffer(
+        Unpooled.wrappedBuffer(bytes.asReadOnlyByteBuffer()), bytes);
   }
 
   private static final AtomicInteger LEAK_COUNT = new AtomicInteger();
@@ -176,14 +233,28 @@ public static void assertNoLeaks() {
   }
 
   private final ByteBuf buf;
+  private final Object wrapped;
   private final CompletableFuture<Void> released = new CompletableFuture<>();
 
-  private CodecBuffer(ByteBuf buf) {
+  private CodecBuffer(ByteBuf buf, Object wrapped) {
     this.buf = buf;
+    this.wrapped = wrapped;
     this.elements = getStackTrace(LOG);
     assertRefCnt(1);
   }
 
+  public boolean isDirect() {
+    return buf.isDirect();
+  }
+
+  /**
+   * @return the wrapped object if this buffer is created by wrapping it;
+   *         otherwise, return null.
+   */
+  public Object getWrapped() {
+    return wrapped;
+  }
+
   private void assertRefCnt(int expected) {
     Preconditions.assertSame(expected, buf.refCnt(), "refCnt");
   }
@@ -222,7 +293,10 @@ public void close() {
   /** Release this buffer and return it back to the pool. */
   public void release() {
     final boolean set = released.complete(null);
-    Preconditions.assertTrue(set, () -> "Already released: " + this);
+    if (!set) {
+      // Allow a zero capacity buffer to be released multiple times.
+      Preconditions.assertSame(0, buf.capacity(), "capacity");
+    }
     if (buf.release()) {
       assertRefCnt(0);
     } else {

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/utils/db/DelegatedCodec.java

@@ -21,7 +21,6 @@
 
 import javax.annotation.Nonnull;
 import java.io.IOException;
-import java.util.function.IntFunction;
 
 /**
  * A {@link Codec} to serialize/deserialize objects by delegation.
@@ -82,7 +81,7 @@ public final boolean supportCodecBuffer() {
 
   @Override
   public final CodecBuffer toCodecBuffer(@Nonnull T message,
-      IntFunction<CodecBuffer> allocator) throws IOException {
+      CodecBuffer.Allocator allocator) throws IOException {
     return delegate.toCodecBuffer(backward.apply(message), allocator);
   }
 

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/utils/db/IntegerCodec.java

@@ -20,7 +20,6 @@
 
 import javax.annotation.Nonnull;
 import java.nio.ByteBuffer;
-import java.util.function.IntFunction;
 
 /**
  * Codec to serialize/deserialize {@link Integer}.
@@ -44,7 +43,7 @@ public boolean supportCodecBuffer() {
 
   @Override
   public CodecBuffer toCodecBuffer(@Nonnull Integer object,
-      IntFunction<CodecBuffer> allocator) {
+      CodecBuffer.Allocator allocator) {
     return allocator.apply(Integer.BYTES).putInt(object);
   }
 

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/utils/db/LongCodec.java

@@ -20,7 +20,6 @@
 
 import javax.annotation.Nonnull;
 import java.nio.ByteBuffer;
-import java.util.function.IntFunction;
 
 /**
  * Codec to serialize/deserialize {@link Long}.
@@ -41,7 +40,7 @@ public boolean supportCodecBuffer() {
 
   @Override
   public CodecBuffer toCodecBuffer(@Nonnull Long object,
-      IntFunction<CodecBuffer> allocator) {
+      CodecBuffer.Allocator allocator) {
     return allocator.apply(Long.BYTES).putLong(object);
   }
 

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/utils/db/Proto2Codec.java

@@ -28,7 +28,6 @@
 import java.io.OutputStream;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
-import java.util.function.IntFunction;
 
 /**
  * Codecs to serialize/deserialize Protobuf v2 messages.
@@ -70,7 +69,7 @@ public boolean supportCodecBuffer() {
 
   @Override
   public CodecBuffer toCodecBuffer(@Nonnull M message,
-      IntFunction<CodecBuffer> allocator) throws IOException {
+      CodecBuffer.Allocator allocator) throws IOException {
     final int size = message.getSerializedSize();
     return allocator.apply(size).put(writeTo(message, size));
   }

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/utils/db/Proto3Codec.java

@@ -27,7 +27,6 @@
 import java.nio.ByteBuffer;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
-import java.util.function.IntFunction;
 import java.util.function.ToIntFunction;
 
 /**
@@ -82,7 +81,7 @@ private ToIntFunction<ByteBuffer> writeTo(M message, int size) {
 
   @Override
   public CodecBuffer toCodecBuffer(@Nonnull M message,
-      IntFunction<CodecBuffer> allocator) {
+      CodecBuffer.Allocator allocator) {
     final int size = message.getSerializedSize();
     return allocator.apply(size).put(writeTo(message, size));
   }

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/utils/db/ShortCodec.java

@@ -19,7 +19,6 @@
 package org.apache.hadoop.hdds.utils.db;
 
 import java.nio.ByteBuffer;
-import java.util.function.IntFunction;
 
 import javax.annotation.Nonnull;
 
@@ -45,7 +44,7 @@ public boolean supportCodecBuffer() {
 
   @Override
   public CodecBuffer toCodecBuffer(@Nonnull Short object,
-      IntFunction<CodecBuffer> allocator) {
+      CodecBuffer.Allocator allocator) {
     return allocator.apply(Short.BYTES).putShort(object);
   }
 

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/utils/db/StringCodecBase.java

@@ -34,7 +34,6 @@
 import java.nio.charset.CodingErrorAction;
 import java.util.Objects;
 import java.util.function.Function;
-import java.util.function.IntFunction;
 
 /**
  * An abstract {@link Codec} to serialize/deserialize {@link String}
@@ -169,7 +168,7 @@ public boolean supportCodecBuffer() {
 
   @Override
   public CodecBuffer toCodecBuffer(@Nonnull String object,
-      IntFunction<CodecBuffer> allocator) throws IOException {
+      CodecBuffer.Allocator allocator) throws IOException {
     // allocate a larger buffer to avoid encoding twice.
     final int upperBound = getSerializedSizeUpperBound(object);
     final CodecBuffer buffer = allocator.apply(upperBound);

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/utils/db/UuidCodec.java

@@ -21,7 +21,6 @@
 import javax.annotation.Nonnull;
 import java.nio.ByteBuffer;
 import java.util.UUID;
-import java.util.function.IntFunction;
 
 /**
  * Codec to serialize/deserialize {@link UUID}.
@@ -48,7 +47,7 @@ public boolean supportCodecBuffer() {
 
   @Override
   public CodecBuffer toCodecBuffer(@Nonnull UUID id,
-      IntFunction<CodecBuffer> allocator) {
+      CodecBuffer.Allocator allocator) {
     return allocator.apply(SERIALIZED_SIZE)
         .putLong(id.getMostSignificantBits())
         .putLong(id.getLeastSignificantBits());

hadoop-hdds/common/src/test/java/org/apache/hadoop/hdds/utils/db/CodecTestUtil.java

@@ -71,7 +71,7 @@ public static <T> void runTest(Codec<T> codec, T original,
 
     // serialize to CodecBuffer
     final CodecBuffer codecBuffer = codec.toCodecBuffer(
-        original, CodecBuffer::allocateHeap);
+        original, CodecBuffer.Allocator.getHeap());
     Assertions.assertEquals(array.length, codecBuffer.readableBytes());
     final ByteBuffer byteBuffer = codecBuffer.asReadOnlyByteBuffer();
     Assertions.assertEquals(array.length, byteBuffer.remaining());