Implement SnappyByteBuffDecompressor for zero-copy HFile block decomp…

hbase-compression/hbase-compression-snappy/src/main/java/org/apache/hadoop/hbase/io/compress/xerial/SnappyByteBuffDecompressor.java

@@ -0,0 +1,127 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.hadoop.hbase.io.compress.xerial;
+
+import edu.umd.cs.findbugs.annotations.Nullable;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import org.apache.hadoop.hbase.io.compress.BlockDecompressorHelper;
+import org.apache.hadoop.hbase.io.compress.ByteBuffDecompressor;
+import org.apache.hadoop.hbase.io.compress.Compression;
+import org.apache.hadoop.hbase.nio.ByteBuff;
+import org.apache.hadoop.hbase.nio.SingleByteBuff;
+import org.apache.yetus.audience.InterfaceAudience;
+import org.xerial.snappy.Snappy;
+
+@InterfaceAudience.Private
+public class SnappyByteBuffDecompressor implements ByteBuffDecompressor {
+
+  private boolean allowByteBuffDecompression;
+
+  SnappyByteBuffDecompressor() {
+    allowByteBuffDecompression = true;
+  }
+
+  @Override
+  public boolean canDecompress(ByteBuff output, ByteBuff input) {
+    if (!allowByteBuffDecompression) {
+      return false;
+    }
+    if (!(output instanceof SingleByteBuff) || !(input instanceof SingleByteBuff)) {
+      return false;
+    }
+    ByteBuffer nioInput = input.nioByteBuffers()[0];
+    ByteBuffer nioOutput = output.nioByteBuffers()[0];
+    // Xerial Snappy's ByteBuffer API only supports direct-to-direct; the byte array API handles
+    // heap-to-heap. Mixed direct/heap is not supported.
+    return (nioInput.isDirect() && nioOutput.isDirect())
+      || (nioInput.hasArray() && nioOutput.hasArray());
+  }
+
+  @Override
+  public int decompress(ByteBuff output, ByteBuff input, int inputLen) throws IOException {
+    return BlockDecompressorHelper.decompress(output, input, inputLen, this::decompressRaw);
+  }
+
+  private int decompressRaw(ByteBuff output, ByteBuff input, int inputLen) throws IOException {
+    if (!(output instanceof SingleByteBuff) || !(input instanceof SingleByteBuff)) {
+      throw new IllegalStateException(
+        "At least one buffer is not a SingleByteBuff, this is not supported");
+    }
+    ByteBuffer nioOutput = output.nioByteBuffers()[0];
+    ByteBuffer nioInput = input.nioByteBuffers()[0];
+    int origOutputPos = nioOutput.position();
+    int origInputPos = nioInput.position();
+    int bytesDecompressed;
+    if (nioInput.isDirect() && nioOutput.isDirect()) {
+      bytesDecompressed = decompressRawDirect(nioInput, nioOutput, inputLen);
+    } else if (nioInput.hasArray() && nioOutput.hasArray()) {
+      bytesDecompressed = decompressRawHeap(nioInput, nioOutput, inputLen);
+    } else {
+      throw new IllegalStateException(
+        "SnappyByteBuffDecompressor only supports direct-to-direct or heap-to-heap decompression,"
+          + " this should never happen since canDecompress() would have returned false");
+    }
+    nioOutput.position(origOutputPos + bytesDecompressed);
+    nioInput.position(origInputPos + inputLen);
+    return bytesDecompressed;
+  }
+
+
+  private int decompressRawDirect(ByteBuffer nioInput, ByteBuffer nioOutput, int inputLen)
+    throws IOException {
+    int savedInputLimit = nioInput.limit();
+    nioInput.limit(nioInput.position() + inputLen);
+    try {
+      return Snappy.uncompress(nioInput, nioOutput);
+    } catch (IOException e) {
+      throw new IOException("Snappy decompression failed: " + e.getMessage(), e);
+    } finally {
+      nioInput.limit(savedInputLimit);
+    }
+  }
+
+  private int decompressRawHeap(ByteBuffer nioInput, ByteBuffer nioOutput, int inputLen)
+    throws IOException {
+    try {
+      return Snappy.uncompress(nioInput.array(), nioInput.arrayOffset() + nioInput.position(),
+        inputLen, nioOutput.array(), nioOutput.arrayOffset() + nioOutput.position());
+    } catch (IOException e) {
+      throw new IOException("Snappy decompression failed: " + e.getMessage(), e);
+    }
+  }
+
+  @Override
+  public void reinit(@Nullable Compression.HFileDecompressionContext newHFileDecompressionContext) {
+    if (newHFileDecompressionContext == null) {
+      return;
+    }
+    if (!(newHFileDecompressionContext instanceof SnappyHFileDecompressionContext)) {
+      throw new IllegalArgumentException(
+        "SnappyByteBuffDecompressor#reinit() was given an HFileDecompressionContext that was not "
+          + "a SnappyHFileDecompressionContext, this should never happen");
+    }
+    SnappyHFileDecompressionContext ctx =
+      (SnappyHFileDecompressionContext) newHFileDecompressionContext;
+    allowByteBuffDecompression = ctx.isAllowByteBuffDecompression();
+  }
+
+  @Override
+  public void close() {
+  }
+}

hbase-compression/hbase-compression-snappy/src/main/java/org/apache/hadoop/hbase/io/compress/xerial/SnappyCodec.java

@@ -23,6 +23,9 @@
 import org.apache.hadoop.conf.Configurable;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.CommonConfigurationKeys;
+import org.apache.hadoop.hbase.io.compress.ByteBuffDecompressionCodec;
+import org.apache.hadoop.hbase.io.compress.ByteBuffDecompressor;
+import org.apache.hadoop.hbase.io.compress.Compression;
 import org.apache.hadoop.io.compress.BlockCompressorStream;
 import org.apache.hadoop.io.compress.BlockDecompressorStream;
 import org.apache.hadoop.io.compress.CompressionCodec;
@@ -41,7 +44,7 @@
  * This is data format compatible with Hadoop's native snappy codec.
  */
 @InterfaceAudience.Private
-public class SnappyCodec implements Configurable, CompressionCodec {
+public class SnappyCodec implements Configurable, CompressionCodec, ByteBuffDecompressionCodec {
 
   public static final String SNAPPY_BUFFER_SIZE_KEY = "hbase.io.compress.snappy.buffersize";
 
@@ -133,6 +136,22 @@ public String getDefaultExtension() {
     return ".snappy";
   }
 
+  @Override
+  public ByteBuffDecompressor createByteBuffDecompressor() {
+    return new SnappyByteBuffDecompressor();
+  }
+
+  @Override
+  public Class<? extends ByteBuffDecompressor> getByteBuffDecompressorType() {
+    return SnappyByteBuffDecompressor.class;
+  }
+
+  @Override
+  public Compression.HFileDecompressionContext
+    getDecompressionContextFromConfiguration(Configuration conf) {
+    return SnappyHFileDecompressionContext.fromConfiguration(conf);
+  }
+
   // Package private
 
   static int getBufferSize(Configuration conf) {

hbase-compression/hbase-compression-snappy/src/main/java/org/apache/hadoop/hbase/io/compress/xerial/SnappyHFileDecompressionContext.java

@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.hadoop.hbase.io.compress.xerial;
+
+import java.io.IOException;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.hbase.io.compress.Compression;
+import org.apache.hadoop.hbase.util.ClassSize;
+import org.apache.yetus.audience.InterfaceAudience;
+
+@InterfaceAudience.Private
+public final class SnappyHFileDecompressionContext extends Compression.HFileDecompressionContext {
+
+  public static final long FIXED_OVERHEAD =
+    ClassSize.estimateBase(SnappyHFileDecompressionContext.class, false);
+
+  public static final String ALLOW_BYTE_BUFF_DECOMPRESSION_KEY =
+    "hbase.io.compress.snappy.allowByteBuffDecompression";
+
+  private final boolean allowByteBuffDecompression;
+
+  private SnappyHFileDecompressionContext(boolean allowByteBuffDecompression) {
+    this.allowByteBuffDecompression = allowByteBuffDecompression;
+  }
+
+  public boolean isAllowByteBuffDecompression() {
+    return allowByteBuffDecompression;
+  }
+
+  public static SnappyHFileDecompressionContext fromConfiguration(Configuration conf) {
+    return new SnappyHFileDecompressionContext(
+      conf.getBoolean(ALLOW_BYTE_BUFF_DECOMPRESSION_KEY, true));
+  }
+
+  @Override
+  public void close() throws IOException {
+  }
+
+  @Override
+  public long heapSize() {
+    return FIXED_OVERHEAD;
+  }
+
+  @Override
+  public String toString() {
+    return "SnappyHFileDecompressionContext{allowByteBuffDecompression=" + allowByteBuffDecompression
+      + '}';
+  }
+}