feat(libhdfs): add direct ByteBuffer write path

Author: Asmoday

hadoop-hdds/client/src/main/java/org/apache/hadoop/hdds/scm/storage/BlockOutputStream.java

@@ -394,6 +394,45 @@ public void write(byte[] b, int off, int len) throws IOException {
     }
   }
 
+  /**
+   * Writes {@code len} bytes starting at absolute position {@code off} of the
+   * given {@link ByteBuffer}. This mirrors {@link #write(byte[], int, int)} but
+   * feeds the buffer straight into the pooled {@link ChunkBuffer} through
+   * {@link ChunkBuffer#put(ByteBuffer)}, avoiding an intermediate byte[] copy.
+   * The caller's buffer position and limit are left unchanged.
+   */
+  public void write(ByteBuffer b, int off, int len) throws IOException {
+    checkOpen();
+    if (b == null) {
+      throw new NullPointerException();
+    }
+    if ((off < 0) || (off > b.capacity()) || (len < 0)
+        || ((off + len) > b.capacity()) || ((off + len) < 0)) {
+      throw new IndexOutOfBoundsException("Offset=" + off + " and len=" + len
+          + " don't match the buffer capacity of " + b.capacity());
+    }
+    if (len == 0) {
+      return;
+    }
+    synchronized (this) {
+      final ByteBuffer src = b.duplicate();
+      src.clear();
+      src.position(off).limit(off + len);
+      while (src.hasRemaining()) {
+        allocateNewBufferIfNeeded();
+        final int writeLen = Math.min(currentBufferRemaining, src.remaining());
+        final int srcLimit = src.limit();
+        src.limit(src.position() + writeLen);
+        currentBuffer.put(src);
+        src.limit(srcLimit);
+        currentBufferRemaining -= writeLen;
+        updateWrittenDataLength(writeLen);
+        writeChunkIfNeeded();
+        doFlushOrWatchIfNeeded();
+      }
+    }
+  }
+
   protected synchronized void updateWrittenDataLength(int writeLen) {
     writtenDataLength += writeLen;
   }

hadoop-hdds/client/src/main/java/org/apache/hadoop/ozone/client/io/ByteArrayStreamOutput.java

@@ -40,7 +40,9 @@ public void write(ByteBuffer buffer, int off, int len) throws IOException {
     }
 
     if (buffer.hasArray()) {
-      write(buffer.array(), off, len);
+      // off is a logical position within the buffer; translate it to an index
+      // into the backing array (arrayOffset is non-zero e.g. for sliced buffers).
+      write(buffer.array(), buffer.arrayOffset() + off, len);
     } else {
       final byte[] array = new byte[Math.min(ARRAY_SIZE_LIMIT, len)];
       for (; len > 0;) {

hadoop-hdds/client/src/test/java/org/apache/hadoop/ozone/client/io/TestByteArrayStreamOutput.java

@@ -0,0 +1,104 @@
+/*
+ * 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.ozone.client.io;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import java.io.ByteArrayOutputStream;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Tests {@link ByteArrayStreamOutput#write(ByteBuffer, int, int)}, in
+ * particular that a logical {@code off} is translated correctly for heap
+ * buffers with a non-zero {@link ByteBuffer#arrayOffset()} (e.g. slices).
+ */
+public class TestByteArrayStreamOutput {
+
+  private static final byte[] DATA = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
+
+  /** Collects everything written through the byte[] sink. */
+  private static final class Collector extends ByteArrayStreamOutput {
+    private final ByteArrayOutputStream sink = new ByteArrayOutputStream();
+
+    @Override
+    public void write(byte[] b, int off, int len) {
+      sink.write(b, off, len);
+    }
+
+    @Override
+    public void flush() {
+    }
+
+    @Override
+    public void hflush() {
+    }
+
+    @Override
+    public void hsync() {
+    }
+
+    byte[] written() {
+      return sink.toByteArray();
+    }
+  }
+
+  private static void writeRemaining(Collector c, ByteBuffer buf)
+      throws java.io.IOException {
+    c.write(buf, buf.position(), buf.remaining());
+  }
+
+  @Test
+  public void writesSlicedHeapBufferFromCorrectOffset() throws Exception {
+    ByteBuffer base = ByteBuffer.allocate(DATA.length);
+    base.put(DATA);
+    base.position(3);
+    ByteBuffer slice = base.slice(); // arrayOffset() == 3, content = DATA[3..]
+    assertEquals(3, slice.arrayOffset());
+
+    Collector c = new Collector();
+    writeRemaining(c, slice);
+
+    assertArrayEquals(Arrays.copyOfRange(DATA, 3, DATA.length), c.written());
+  }
+
+  @Test
+  public void writesHeapBufferFromLogicalPosition() throws Exception {
+    ByteBuffer buf = ByteBuffer.wrap(DATA.clone());
+    buf.position(4).limit(9);
+
+    Collector c = new Collector();
+    writeRemaining(c, buf);
+
+    assertArrayEquals(Arrays.copyOfRange(DATA, 4, 9), c.written());
+  }
+
+  @Test
+  public void writesDirectBuffer() throws Exception {
+    ByteBuffer buf = ByteBuffer.allocateDirect(DATA.length);
+    buf.put(DATA);
+    buf.position(2).limit(8);
+
+    Collector c = new Collector();
+    writeRemaining(c, buf);
+
+    assertArrayEquals(Arrays.copyOfRange(DATA, 2, 8), c.written());
+  }
+}

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/BlockOutputStreamEntry.java

@@ -20,6 +20,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.nio.ByteBuffer;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.concurrent.ExecutorService;
@@ -184,6 +185,12 @@ public void write(byte[] b, int off, int len) throws IOException {
     incCurrentPosition(len);
   }
 
+  void write(ByteBuffer b, int off, int len) throws IOException {
+    checkStream();
+    outputStream.write(b, off, len);
+    incCurrentPosition(len);
+  }
+
   void writeOnRetry(long len) throws IOException {
     checkStream();
     BlockOutputStream out = (BlockOutputStream) getOutputStream();

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/ECKeyOutputStream.java

@@ -35,6 +35,7 @@
 import java.util.stream.IntStream;
 import org.apache.commons.lang3.NotImplementedException;
 import org.apache.hadoop.fs.FSExceptionMessages;
+import org.apache.hadoop.fs.StreamCapabilities;
 import org.apache.hadoop.hdds.client.ECReplicationConfig;
 import org.apache.hadoop.hdds.scm.OzoneClientConfig;
 import org.apache.hadoop.hdds.scm.client.HddsClientUtils;
@@ -158,6 +159,28 @@ public void write(byte[] b, int off, int len) throws IOException {
     writeOffset += len;
   }
 
+  /**
+   * EC striping does not support the direct ByteBuffer write path; drain the
+   * buffer through the byte[] path. Also see {@link #hasCapability(String)},
+   * which refuses {@code WRITEBYTEBUFFER} so a wrapping CryptoOutputStream never
+   * forwards a ciphertext buffer here.
+   */
+  @Override
+  public void write(ByteBuffer buf) throws IOException {
+    final int len = buf.remaining();
+    final byte[] tmp = new byte[len];
+    buf.get(tmp);
+    write(tmp, 0, len);
+  }
+
+  @Override
+  public boolean hasCapability(String capability) {
+    if (StreamCapabilities.WRITEBYTEBUFFER.equalsIgnoreCase(capability)) {
+      return false;
+    }
+    return super.hasCapability(capability);
+  }
+
   private void rollbackAndReset(ECChunkBuffers stripe) throws IOException {
     // Rollback the length/offset updated as part of this failed stripe write.
     final ByteBuffer[] dataBuffers = stripe.getDataBuffers();

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/KeyOutputStream.java

@@ -22,6 +22,7 @@
 import java.io.IOException;
 import java.io.InterruptedIOException;
 import java.io.OutputStream;
+import java.nio.ByteBuffer;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -33,7 +34,9 @@
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
+import org.apache.hadoop.fs.ByteBufferWritable;
 import org.apache.hadoop.fs.FSExceptionMessages;
+import org.apache.hadoop.fs.StreamCapabilities;
 import org.apache.hadoop.fs.Syncable;
 import org.apache.hadoop.hdds.client.ReplicationConfig;
 import org.apache.hadoop.hdds.protocol.DatanodeDetails;
@@ -73,7 +76,7 @@
  * TODO : currently not support multi-thread access.
  */
 public class KeyOutputStream extends OutputStream
-    implements Syncable, KeyMetadataAware {
+    implements Syncable, KeyMetadataAware, ByteBufferWritable, StreamCapabilities {
 
   private static final Logger LOG =
       LoggerFactory.getLogger(KeyOutputStream.class);
@@ -245,6 +248,58 @@ public void write(byte[] b, int off, int len)
     }
   }
 
+  /**
+   * Writes {@code len} bytes starting at absolute position {@code off} of the
+   * given {@link ByteBuffer}, mirroring {@link #write(byte[], int, int)} but
+   * routing the buffer down to the block layer without an intermediate heap
+   * copy. The caller's buffer position and limit are left unchanged.
+   */
+  public void write(ByteBuffer b, int off, int len) throws IOException {
+    try {
+      getRequestSemaphore().acquire();
+      checkNotClosed();
+      if (b == null) {
+        throw new NullPointerException();
+      }
+      if ((off < 0) || (off > b.capacity()) || (len < 0)
+          || ((off + len) > b.capacity()) || ((off + len) < 0)) {
+        throw new IndexOutOfBoundsException();
+      }
+      if (len == 0) {
+        return;
+      }
+
+      doInWriteLock(() -> {
+        doWrite((entry, o, writeLen) -> entry.write(b, o, writeLen),
+            off, len, false);
+        writeOffset += len;
+      });
+    } finally {
+      getRequestSemaphore().release();
+    }
+  }
+
+  /**
+   * {@link ByteBufferWritable} entry point. Writes all remaining bytes of
+   * {@code buf} and advances its position to its limit. Lets a wrapping
+   * {@link org.apache.hadoop.crypto.CryptoOutputStream} hand its direct
+   * ciphertext buffer straight down without a heap copy.
+   */
+  @Override
+  public void write(ByteBuffer buf) throws IOException {
+    final int off = buf.position();
+    final int len = buf.remaining();
+    write(buf, off, len);
+    buf.position(buf.limit());
+  }
+
+  @Override
+  public boolean hasCapability(String capability) {
+    return StreamCapabilities.WRITEBYTEBUFFER.equalsIgnoreCase(capability)
+        || StreamCapabilities.HFLUSH.equalsIgnoreCase(capability)
+        || StreamCapabilities.HSYNC.equalsIgnoreCase(capability);
+  }
+
   private <E extends Throwable> void doInWriteLock(CheckedRunnable<E> block) throws E {
     writeLock.lock();
     try {
@@ -257,6 +312,25 @@ private <E extends Throwable> void doInWriteLock(CheckedRunnable<E> block) throw
   @VisibleForTesting
   void handleWrite(byte[] b, int off, long len, boolean retry)
       throws IOException {
+    doWrite((entry, o, writeLen) -> entry.write(b, o, writeLen), off, len,
+        retry);
+  }
+
+  /**
+   * Writes a chunk of a data source into a {@link BlockOutputStreamEntry}.
+   * The source (byte[] or {@link ByteBuffer}) is captured by the implementation
+   * so that the block-allocation, retry and exception bookkeeping below stays
+   * source-agnostic. Not invoked on the retry path, which replays from the
+   * block's own buffered data.
+   */
+  @FunctionalInterface
+  private interface ChunkWriter {
+    void write(BlockOutputStreamEntry entry, int off, int writeLen)
+        throws IOException;
+  }
+
+  private void doWrite(ChunkWriter writer, int off, long len, boolean retry)
+      throws IOException {
     while (len > 0) {
       try {
         BlockOutputStreamEntry current =
@@ -271,8 +345,8 @@ void handleWrite(byte[] b, int off, long len, boolean retry)
         // or if it sees an exception, how much the actual write was
         // acknowledged.
         int writtenLength =
-                writeToOutputStream(current, retry, len, b, expectedWriteLen,
-                off, currentPos);
+                writeToOutputStream(current, retry, len, writer,
+                expectedWriteLen, off, currentPos);
         if (current.getRemaining() <= 0) {
           // since the current block is already written close the stream.
           handleFlushOrClose(StreamAction.FULL);
@@ -287,15 +361,15 @@ void handleWrite(byte[] b, int off, long len, boolean retry)
   }
 
   private int writeToOutputStream(BlockOutputStreamEntry current,
-      boolean retry, long len, byte[] b, int writeLen, int off, long currentPos)
-      throws IOException {
+      boolean retry, long len, ChunkWriter writer, int writeLen, int off,
+      long currentPos) throws IOException {
     try {
       current.registerCallReceived();
       if (retry) {
         current.writeOnRetry(len);
       } else {
         current.waitForRetryHandling(retryHandlingCondition);
-        current.write(b, off, writeLen);
+        writer.write(current, off, writeLen);
         offset += writeLen;
       }
       current.registerCallFinished();

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/OzoneOutputStream.java

@@ -19,6 +19,7 @@
 
 import java.io.IOException;
 import java.io.OutputStream;
+import java.nio.ByteBuffer;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
@@ -93,6 +94,20 @@ public void write(byte[] b, int off, int len) throws IOException {
     outputStream.write(b, off, len);
   }
 
+  @Override
+  public void write(ByteBuffer buffer, int off, int len) throws IOException {
+    // The direct ByteBuffer write path only applies to a plain replicated key.
+    // Encryption (CryptoOutputStream / CipherOutputStreamOzone) needs the
+    // plaintext on the heap first, and EC uses a different striping path, so
+    // both fall back to the byte[] path in ByteArrayStreamOutput.
+    if (outputStream instanceof KeyOutputStream
+        && !(outputStream instanceof ECKeyOutputStream)) {
+      ((KeyOutputStream) outputStream).write(buffer, off, len);
+    } else {
+      super.write(buffer, off, len);
+    }
+  }
+
   @Override
   public synchronized void flush() throws IOException {
     outputStream.flush();

hadoop-ozone/client/src/test/java/org/apache/hadoop/ozone/client/io/TestKeyOutputStream.java

@@ -61,7 +61,7 @@ void testConcurrentWriteLimitOne() throws Exception {
     doAnswer(invocation -> {
       countWrite.getAndIncrement();
       return invocation.callRealMethod();
-    }).when(keyOutputStream).write(any(), anyInt(), anyInt());
+    }).when(keyOutputStream).write(any(byte[].class), anyInt(), anyInt());
 
     final ConcurrentHashMap<Long, CountDownLatch> mapNotifiers = new ConcurrentHashMap<>();