[NEMO-350] Implement Off-heap SerializedMemoryStore & [NEMO-384] Implement DirectByteBufferInputStream for Off-heap SerializedMemoryStore

common/src/main/java/org/apache/nemo/common/DirectByteBufferInputStream.java

@@ -0,0 +1,71 @@
+/*
+ * 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.nemo.common;
+
+import java.io.EOFException;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.util.List;
+
+/**
+ * This class is a customized input stream implementation which reads data from
+ * list of {@link ByteBuffer}.
+ */
+public class DirectByteBufferInputStream extends InputStream {
+  private List<ByteBuffer> bufList;
+  private int current = 0;
+
+  /**
+   * Default Constructor.
+   *
+   * @param bufList is the target data to read.
+   */
+  public DirectByteBufferInputStream(final List<ByteBuffer> bufList) {
+   this.bufList = bufList;
+  }
+
+  /**
+   * Reads data from the list of {@code ByteBuffer}s.
+   *
+   * @return integer.
+   * @throws IOException
+   */
+  @Override
+  public int read() throws IOException {
+    return getBuffer().get() & 0xff;
+  }
+
+  /**
+   * Return next non-empty @code{ByteBuffer}.
+   *
+   * @return @code{ByteBuffer} to write the data
+   * @throws IOException when fail to retrieve buffer.
+   */
+  public ByteBuffer getBuffer() throws IOException {
+    while (current < bufList.size()) {
+      ByteBuffer buffer = bufList.get(current);
+      if (buffer.hasRemaining()) {
+        return buffer;
+      }
+      current += 1;
+    }
+    throw new EOFException();
+  }
+}

common/src/main/java/org/apache/nemo/common/DirectByteBufferOutputStream.java

@@ -20,6 +20,7 @@
 
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
 import java.util.LinkedList;
 import java.util.List;
 
@@ -31,11 +32,10 @@
 public final class DirectByteBufferOutputStream extends OutputStream {
 
   private LinkedList<ByteBuffer> dataList = new LinkedList<>();
-  private static final int DEFAULT_PAGE_SIZE = 4096;
+  private static final int DEFAULT_PAGE_SIZE = 32768; //32KB
   private final int pageSize;
   private ByteBuffer currentBuf;
 
-
   /**
    * Default constructor.
    * Sets the {@code pageSize} as default size of 4096 bytes.
@@ -140,38 +140,45 @@ public byte[] toByteArray() {
     final int arraySize = pageSize * (dataList.size() - 1) + lastBuf.position();
     final byte[] byteArray = new byte[arraySize];
     int start = 0;
-    int byteToWrite;
 
-    for (final ByteBuffer temp : dataList) {
+    for (final ByteBuffer buffer : dataList) {
       // ByteBuffer has to be shifted to read mode by calling ByteBuffer.flip(),
       // which sets limit to the current position and sets the position to 0.
       // Note that capacity remains unchanged.
-      temp.flip();
-      byteToWrite = temp.remaining();
-      temp.get(byteArray, start, byteToWrite);
+      final ByteBuffer dupBuffer = buffer.duplicate();
+      dupBuffer.flip();
+      final int byteToWrite = dupBuffer.remaining();
+      dupBuffer.get(byteArray, start, byteToWrite);
       start += byteToWrite;
     }
-    // The limit of the last buffer has to be set to the capacity for additional write.
-    lastBuf.limit(lastBuf.capacity());
 
     return byteArray;
   }
 
   /**
    * Returns the list of {@code ByteBuffer}s that contains the written data.
-   * Note that by calling this method, the existing list of {@code ByteBuffer}s is cleared.
    *
    * @return the {@code LinkedList} of {@code ByteBuffer}s.
    */
-  public List<ByteBuffer> getBufferListAndClear() {
-    List<ByteBuffer> result = dataList;
-    dataList = new LinkedList<>();
-    for (final ByteBuffer buffer : result) {
-      buffer.flip();
+  public List<ByteBuffer> getBufferList() {
+    List<ByteBuffer> result = new ArrayList<>(dataList.size());
+    for (final ByteBuffer buffer : dataList) {
+      final ByteBuffer dupBuffer = buffer.duplicate();
+      dupBuffer.flip();
+      result.add(dupBuffer);
     }
     return result;
   }
 
+  /**
+   * Returns the size of the data written in this output stream.
+   *
+   * @return the size of the data
+   */
+  public int size() {
+    return pageSize * (dataList.size() - 1) + dataList.getLast().position();
+  }
+
   /**
    * Closing this output stream has no effect.
    */

common/src/test/java/org/apache/nemo/common/DirectByteBufferOutputStreamTest.java

@@ -101,7 +101,7 @@ public void testGetBufferList() {
     String value = RandomStringUtils.randomAlphanumeric(10000);
     outputStream.write(value.getBytes());
     byte[] totalOutput = outputStream.toByteArray();
-    List<ByteBuffer> bufList = outputStream.getBufferListAndClear();
+    List<ByteBuffer> bufList = outputStream.getBufferList();
     int offset = 0;
     int byteToRead;
     for (final ByteBuffer temp : bufList) {

examples/beam/pom.xml

@@ -41,7 +41,6 @@ under the License.
       <groupId>org.apache.nemo</groupId>
       <artifactId>nemo-client</artifactId>
       <version>${project.version}</version>
-      <scope>test</scope>
     </dependency>
     <dependency>
       <groupId>com.github.fommil.netlib</groupId>

runtime/executor/src/main/java/org/apache/nemo/runtime/executor/bytetransfer/ByteOutputContext.java

@@ -34,9 +34,13 @@
 import javax.annotation.Nullable;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.nio.ByteBuffer;
 import java.nio.channels.FileChannel;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.List;
+
+import static io.netty.buffer.Unpooled.wrappedBuffer;
 
 /**
  * Container for multiple output streams. Represents a transfer context on sender-side.
@@ -157,12 +161,37 @@ public void write(final byte[] bytes, final int offset, final int length) throws
      * @return {@code this}
      * @throws IOException when an exception has been set or this stream was closed
      */
-    public ByteOutputStream writeSerializedPartition(final SerializedPartition serializedPartition)
+    public ByteOutputStream writeSerializedPartitionBuffer(final SerializedPartition serializedPartition)
       throws IOException {
-      write(serializedPartition.getData(), 0, serializedPartition.getLength());
+      writeBuffer(serializedPartition.getBuffer());
       return this;
     }
 
+    /**
+     * Wraps each of the {@link ByteBuffer} in the bufList to {@link ByteBuf} object
+     * to write a data frame.
+     *
+     * @param bufList
+     * @throws IOException
+     */
+    public void writeBuffer(final List<ByteBuffer> bufList) throws IOException {
+      final ByteBuf byteBuf = wrappedBuffer(bufList.toArray(new ByteBuffer[bufList.size()]));
+      writeByteBuf(byteBuf);
+    }
+
+
+    /**
+     * Writes a data frame, from {@link ByteBuf}.
+     *
+     * @param byteBuf {@link ByteBuf} to write.
+     * @throws IOException
+     */
+    private void writeByteBuf(final ByteBuf byteBuf) throws IOException {
+      if (byteBuf.readableBytes() > 0) {
+        writeDataFrame(byteBuf, byteBuf.readableBytes());
+      }
+    }
+
     /**
      * Writes a data frame from {@link FileArea}.
      *
@@ -196,17 +225,6 @@ public void close() throws IOException {
       closed = true;
     }
 
-    /**
-     * Writes a data frame, from {@link ByteBuf}.
-     *
-     * @param byteBuf {@link ByteBuf} to write.
-     */
-    private void writeByteBuf(final ByteBuf byteBuf) throws IOException {
-      if (byteBuf.readableBytes() > 0) {
-        writeDataFrame(byteBuf, byteBuf.readableBytes());
-      }
-    }
-
     /**
      * Write an element to the channel.
      *

runtime/executor/src/main/java/org/apache/nemo/runtime/executor/data/BlockManagerWorker.java

@@ -356,7 +356,7 @@ public void run() {
               final Iterable<SerializedPartition> partitions = optionalBlock.get().readSerializedPartitions(keyRange);
               for (final SerializedPartition partition : partitions) {
                 try (ByteOutputContext.ByteOutputStream os = outputContext.newOutputStream()) {
-                  os.writeSerializedPartition(partition);
+                  os.writeSerializedPartitionBuffer(partition);
                 }
               }
             }

runtime/executor/src/main/java/org/apache/nemo/runtime/executor/data/DataUtil.java

@@ -19,7 +19,8 @@
 package org.apache.nemo.runtime.executor.data;
 
 import com.google.common.io.CountingInputStream;
-import org.apache.nemo.common.DirectByteArrayOutputStream;
+import org.apache.nemo.common.DirectByteBufferInputStream;
+import org.apache.nemo.common.DirectByteBufferOutputStream;
 import org.apache.nemo.common.coder.DecoderFactory;
 import org.apache.nemo.common.coder.EncoderFactory;
 import org.apache.nemo.runtime.executor.data.partition.NonSerializedPartition;
@@ -31,6 +32,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.*;
+import java.nio.ByteBuffer;
 import java.util.*;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -112,19 +114,19 @@ public static <K extends Serializable> Iterable<SerializedPartition<K>> convertT
     final List<SerializedPartition<K>> serializedPartitions = new ArrayList<>();
     for (final NonSerializedPartition<K> partitionToConvert : partitionsToConvert) {
       try (
-        DirectByteArrayOutputStream bytesOutputStream = new DirectByteArrayOutputStream();
-        OutputStream wrappedStream = buildOutputStream(bytesOutputStream, serializer.getEncodeStreamChainers());
+        DirectByteBufferOutputStream bytesOutputStream = new DirectByteBufferOutputStream();
+        OutputStream wrappedStream = buildOutputStream(bytesOutputStream, serializer.getEncodeStreamChainers())
       ) {
         serializePartition(serializer.getEncoderFactory(), partitionToConvert, wrappedStream);
         // We need to close wrappedStream on here, because DirectByteArrayOutputStream:getBufDirectly() returns
         // inner buffer directly, which can be an unfinished(not flushed) buffer.
         wrappedStream.close();
         // Note that serializedBytes include invalid bytes.
         // So we have to use it with the actualLength by using size() whenever needed.
-        final byte[] serializedBytes = bytesOutputStream.getBufDirectly();
+        final List<ByteBuffer> serializedBufList = bytesOutputStream.getBufferList();
         final int actualLength = bytesOutputStream.size();
         serializedPartitions.add(
-          new SerializedPartition<>(partitionToConvert.getKey(), serializedBytes, actualLength));
+          new SerializedPartition<>(partitionToConvert.getKey(), serializedBufList, actualLength));
       }
     }
     return serializedPartitions;
@@ -148,10 +150,10 @@ public static <K extends Serializable> Iterable<NonSerializedPartition<K>> conve
       final K key = partitionToConvert.getKey();
 
 
-      try (ByteArrayInputStream byteArrayInputStream =
-             new ByteArrayInputStream(partitionToConvert.getData())) {
+      try (DirectByteBufferInputStream byteBufferInputStream =
+             new DirectByteBufferInputStream(partitionToConvert.getBuffer())) {
         final NonSerializedPartition<K> deserializePartition = deserializePartition(
-          partitionToConvert.getLength(), serializer, key, byteArrayInputStream);
+          partitionToConvert.getLength(), serializer, key, byteBufferInputStream);
         nonSerializedPartitions.add(deserializePartition);
       }
     }

runtime/executor/src/main/java/org/apache/nemo/runtime/executor/data/block/FileBlock.java

@@ -34,10 +34,23 @@
 import org.slf4j.LoggerFactory;
 
 import javax.annotation.concurrent.NotThreadSafe;
-import java.io.*;
+
+import java.io.ByteArrayInputStream;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.Serializable;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
 import java.nio.file.Files;
 import java.nio.file.Paths;
-import java.util.*;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 
 /**
  * This class represents a block which is stored in (local or remote) file.
@@ -83,11 +96,13 @@ public FileBlock(final String blockId,
    */
   private void writeToFile(final Iterable<SerializedPartition<K>> serializedPartitions)
     throws IOException {
-    try (FileOutputStream fileOutputStream = new FileOutputStream(filePath, true)) {
+    try (FileChannel fileOutputChannel = new FileOutputStream(filePath, true).getChannel()) {
       for (final SerializedPartition<K> serializedPartition : serializedPartitions) {
         // Reserve a partition write and get the metadata.
         metadata.writePartitionMetadata(serializedPartition.getKey(), serializedPartition.getLength());
-        fileOutputStream.write(serializedPartition.getData(), 0, serializedPartition.getLength());
+        for (final ByteBuffer buffer: serializedPartition.getBuffer()) {
+          fileOutputChannel.write(buffer);
+        }
       }
     }
   }