HubSpot Backport: HBASE-28256 Enhance ByteBufferUtils.readVLong to re…

…ad more bytes at a time (apache#5576) Signed-off-by: Bryan Beaudreault <bbeaudreault@apache.org> Signed-off-by: Duo Zhang <zhangduo@apache.org>
HubSpot · Jan 16, 2024 · 87db97e · 87db97e
1 parent e48abda
commit 87db97e
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 23 deletions.
diff --git a/hbase-common/src/main/java/org/apache/hadoop/hbase/util/ByteBufferUtils.java b/hbase-common/src/main/java/org/apache/hadoop/hbase/util/ByteBufferUtils.java
@@ -471,38 +471,75 @@ public static void writeVLong(ByteBuffer out, long i) {
     }
   }
 
-  private interface ByteVisitor {
-    byte get();
-  }
-
-  private static long readVLong(ByteVisitor visitor) {
-    byte firstByte = visitor.get();
+  /**
+   * Similar to {@link WritableUtils#readVLong(java.io.DataInput)} but reads from a
+   * {@link ByteBuff}.
+   */
+  public static long readVLong(ByteBuff buf) {
+    byte firstByte = buf.get();
     int len = WritableUtils.decodeVIntSize(firstByte);
     if (len == 1) {
       return firstByte;
+    } else {
+      int remaining = len - 1;
+      long i = 0;
+      int offsetFromPos = 0;
+      if (remaining >= Bytes.SIZEOF_INT) {
+        // The int read has to be converted to unsigned long so the & op
+        i = (buf.getIntAfterPosition(offsetFromPos) & 0x00000000ffffffffL);
+        remaining -= Bytes.SIZEOF_INT;
+        offsetFromPos += Bytes.SIZEOF_INT;
+      }
+      if (remaining >= Bytes.SIZEOF_SHORT) {
+        short s = buf.getShortAfterPosition(offsetFromPos);
+        i = i << 16;
+        i = i | (s & 0xFFFF);
+        remaining -= Bytes.SIZEOF_SHORT;
+        offsetFromPos += Bytes.SIZEOF_SHORT;
+      }
+      for (int idx = 0; idx < remaining; idx++) {
+        byte b = buf.getByteAfterPosition(offsetFromPos + idx);
+        i = i << 8;
+        i = i | (b & 0xFF);
+      }
+      buf.skip(len - 1);
+      return WritableUtils.isNegativeVInt(firstByte) ? ~i : i;
     }
-    long i = 0;
-    for (int idx = 0; idx < len - 1; idx++) {
-      byte b = visitor.get();
-      i = i << 8;
-      i = i | (b & 0xFF);
-    }
-    return (WritableUtils.isNegativeVInt(firstByte) ? (i ^ -1L) : i);
   }
 
   /**
    * Similar to {@link WritableUtils#readVLong(DataInput)} but reads from a {@link ByteBuffer}.
    */
-  public static long readVLong(ByteBuffer in) {
-    return readVLong(in::get);
-  }
-
-  /**
-   * Similar to {@link WritableUtils#readVLong(java.io.DataInput)} but reads from a
-   * {@link ByteBuff}.
-   */
-  public static long readVLong(ByteBuff in) {
-    return readVLong(in::get);
+  public static long readVLong(ByteBuffer buf) {
+    byte firstByte = buf.get();
+    int len = WritableUtils.decodeVIntSize(firstByte);
+    if (len == 1) {
+      return firstByte;
+    } else {
+      int remaining = len - 1;
+      long i = 0;
+      int offsetFromPos = 0;
+      if (remaining >= Bytes.SIZEOF_INT) {
+        // The int read has to be converted to unsigned long so the & op
+        i = (buf.getInt(buf.position() + offsetFromPos) & 0x00000000ffffffffL);
+        remaining -= Bytes.SIZEOF_INT;
+        offsetFromPos += Bytes.SIZEOF_INT;
+      }
+      if (remaining >= Bytes.SIZEOF_SHORT) {
+        short s = buf.getShort(buf.position() + offsetFromPos);
+        i = i << 16;
+        i = i | (s & 0xFFFF);
+        remaining -= Bytes.SIZEOF_SHORT;
+        offsetFromPos += Bytes.SIZEOF_SHORT;
+      }
+      for (int idx = 0; idx < remaining; idx++) {
+        byte b = buf.get(buf.position() + offsetFromPos + idx);
+        i = i << 8;
+        i = i | (b & 0xFF);
+      }
+      buf.position(buf.position() + len - 1);
+      return WritableUtils.isNegativeVInt(firstByte) ? ~i : i;
+    }
   }
 
   /**

diff --git a/hbase-common/src/test/java/org/apache/hadoop/hbase/util/TestByteBufferUtils.java b/hbase-common/src/test/java/org/apache/hadoop/hbase/util/TestByteBufferUtils.java
@@ -46,6 +46,7 @@
 import org.apache.hadoop.hbase.HBaseClassTestRule;
 import org.apache.hadoop.hbase.HBaseCommonTestingUtility;
 import org.apache.hadoop.hbase.HConstants;
+import org.apache.hadoop.hbase.nio.ByteBuff;
 import org.apache.hadoop.hbase.testclassification.MediumTests;
 import org.apache.hadoop.hbase.testclassification.MiscTests;
 import org.apache.hadoop.hbase.unsafe.HBasePlatformDependent;
@@ -179,6 +180,23 @@ public void testReadWriteVLong() {
       ByteBufferUtils.writeVLong(b, l);
       b.flip();
       assertEquals(l, ByteBufferUtils.readVLong(b));
+      b.flip();
+      assertEquals(l, ByteBufferUtils.readVLong(ByteBuff.wrap(b)));
+    }
+  }
+
+  @Test
+  public void testReadWriteConsecutiveVLong() {
+    for (long l : testNumbers) {
+      ByteBuffer b = ByteBuffer.allocate(2 * MAX_VLONG_LENGTH);
+      ByteBufferUtils.writeVLong(b, l);
+      ByteBufferUtils.writeVLong(b, l - 4);
+      b.flip();
+      assertEquals(l, ByteBufferUtils.readVLong(b));
+      assertEquals(l - 4, ByteBufferUtils.readVLong(b));
+      b.flip();
+      assertEquals(l, ByteBufferUtils.readVLong(ByteBuff.wrap(b)));
+      assertEquals(l - 4, ByteBufferUtils.readVLong(ByteBuff.wrap(b)));
     }
   }