RATIS-2568. Add Streaming Read to FileStore

ratis-examples/src/main/java/org/apache/ratis/examples/filestore/FileInfo.java

@@ -19,19 +19,17 @@
 
 import org.apache.ratis.protocol.RaftPeerId;
 import org.apache.ratis.thirdparty.com.google.protobuf.ByteString;
-import org.apache.ratis.util.CollectionUtils;
-import org.apache.ratis.util.JavaUtils;
-import org.apache.ratis.util.LogUtils;
-import org.apache.ratis.util.Preconditions;
-import org.apache.ratis.util.TaskQueue;
+import org.apache.ratis.util.*;
 import org.apache.ratis.util.function.CheckedFunction;
 import org.apache.ratis.util.function.CheckedSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
 import java.nio.channels.SeekableByteChannel;
+import java.nio.channels.WritableByteChannel;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.StandardOpenOption;
@@ -88,6 +86,28 @@ ByteString read(CheckedFunction<Path, Path, IOException> resolver, long offset,
     }
   }
 
+  void streamRead(CheckedFunction<Path, Path, IOException> resolver, long offset, long length,
+      WritableByteChannel stream) throws IOException {
+    if (offset + length > getWriteSize()) {
+      throw new IOException("Failed to read: offset (=" + offset
+          + " + length (=" + length + ") > size = " + getWriteSize()
+          + ", path=" + getRelativePath());
+    }
+
+    try (FileChannel in = FileUtils.newFileChannel(
+            resolver.apply(getRelativePath()), StandardOpenOption.READ)) {
+      long transferred = 0;
+      while (transferred < length) {
+        final long n = in.transferTo(offset + transferred, length - transferred, stream);
+        Preconditions.assertTrue(n >= 0);
+        transferred += n;
+      }
+      Preconditions.assertSame(length, transferred, "transferred");
+    } finally {
+      stream.close();
+    }
+  }
+
   UnderConstruction asUnderConstruction() {
     throw new UnsupportedOperationException(
         "File " + getRelativePath() + " is not under construction.");
@@ -121,6 +141,12 @@ static class ReadOnly extends FileInfo {
       this.writeSize = f.getWriteSize();
     }
 
+    ReadOnly(Path relativePath, long size) {
+      super(relativePath);
+      this.committedSize = size;
+      this.writeSize = size;
+    }
+
     @Override
     long getCommittedSize() {
       return committedSize;

ratis-examples/src/main/java/org/apache/ratis/examples/filestore/FileStore.java

@@ -43,6 +43,7 @@
 import java.io.IOException;
 import java.io.RandomAccessFile;
 import java.nio.ByteBuffer;
+import java.nio.channels.WritableByteChannel;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ArrayList;
@@ -102,6 +103,11 @@ void putNew(UnderConstruction uc) {
       }
     }
 
+    void putReadOnly(ReadOnly ro) {
+      LOG.trace("{}: putReadOnly {}", name, ro.getRelativePath());
+      map.put(ro.getRelativePath(), ro);
+    }
+
     ReadOnly close(UnderConstruction uc) {
       LOG.trace("{}: close {}", name, uc.getRelativePath());
       final ReadOnly ro = new ReadOnly(uc);
@@ -208,6 +214,12 @@ CompletableFuture<ReadReplyProto> read(String relative, long offset, long length
     return submit(task, reader);
   }
 
+  void streamRead(String relative, long offset, long length, WritableByteChannel stream)
+      throws IOException {
+    final FileInfo info = files.get(relative);
+    info.streamRead(this::resolve, offset, length, stream);
+  }
+
   CompletableFuture<Path> delete(long index, String relative) {
     final Supplier<String> name = () -> "delete(" + relative + ") @" + getId() + ":" + index;
     final CheckedSupplier<Path, IOException> task = LogUtils.newCheckedSupplier(LOG, () -> {
@@ -284,13 +296,18 @@ public void close() {
 
   CompletableFuture<StreamWriteReplyProto> streamCommit(String p, long bytesWritten) {
     return CompletableFuture.supplyAsync(() -> {
+      final Path relative = normalize(p);
       final long len;
-      try (RandomAccessFile file = new RandomAccessFile(resolve(normalize(p)).toFile(), "r")) {
+      try (RandomAccessFile file = new RandomAccessFile(resolve(relative).toFile(), "r")) {
         len = file.length();
-        return StreamWriteReplyProto.newBuilder().setIsSuccess(len == bytesWritten).setByteWritten(len).build();
       } catch (IOException e) {
         throw new CompletionException("Failed to commit stream " + p + " with " + bytesWritten + " B.", e);
       }
+      final boolean success = len == bytesWritten;
+      if (success) {
+        files.putReadOnly(new ReadOnly(relative, len));
+      }
+      return StreamWriteReplyProto.newBuilder().setIsSuccess(success).setByteWritten(len).build();
     }, committer);
   }
 

ratis-examples/src/main/java/org/apache/ratis/examples/filestore/FileStoreClient.java

@@ -18,6 +18,7 @@
 package org.apache.ratis.examples.filestore;
 
 import org.apache.ratis.client.RaftClient;
+import org.apache.ratis.client.api.DataStreamInput;
 import org.apache.ratis.client.api.DataStreamOutput;
 import org.apache.ratis.conf.RaftProperties;
 import org.apache.ratis.proto.ExamplesProtos.DeleteReplyProto;
@@ -29,6 +30,8 @@
 import org.apache.ratis.proto.ExamplesProtos.WriteReplyProto;
 import org.apache.ratis.proto.ExamplesProtos.WriteRequestHeaderProto;
 import org.apache.ratis.proto.ExamplesProtos.WriteRequestProto;
+import org.apache.ratis.proto.RaftProtos.DataStreamPacketHeaderProto.Type;
+import org.apache.ratis.protocol.DataStreamReply;
 import org.apache.ratis.protocol.Message;
 import org.apache.ratis.protocol.RaftClientReply;
 import org.apache.ratis.protocol.RaftGroup;
@@ -39,13 +42,15 @@
 import org.apache.ratis.util.JavaUtils;
 import org.apache.ratis.util.Preconditions;
 import org.apache.ratis.util.ProtoUtils;
+import org.apache.ratis.util.ReferenceCountedObject;
 import org.apache.ratis.util.function.CheckedFunction;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.Closeable;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.nio.channels.WritableByteChannel;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.function.Function;
@@ -197,6 +202,49 @@ public DataStreamOutput getStreamOutput(String path, long dataSize, RoutingTable
     return client.getDataStreamApi().stream(request.toByteString().asReadOnlyByteBuffer(), routingTable);
   }
 
+  public DataStreamInput getStreamInput(String path, long offset, long length) {
+    final ReadRequestProto read = ReadRequestProto.newBuilder()
+        .setPath(ProtoUtils.toByteString(path))
+        .setOffset(offset)
+        .setLength(length)
+        .build();
+    return client.getDataStreamApi().streamReadOnly(read.toByteString().asReadOnlyByteBuffer());
+  }
+
+  /**
+   * Read file data using streaming read and write it to the given channel.
+   *
+   * @return total number of bytes read.
+   */
+  public long streamRead(String path, long offset, long length, WritableByteChannel channel)
+      throws IOException {
+    long total = 0;
+    try (DataStreamInput in = getStreamInput(path, offset, length)) {
+      while (true) {
+        final ReferenceCountedObject<DataStreamReply> ref = in.readAsync().join();
+        try {
+          final DataStreamReply reply = ref.get();
+          if (reply.getType() == Type.STREAM_HEADER) {
+            Preconditions.assertTrue(reply.isSuccess(),
+                () -> "Failed to stream read " + path + ", reply=" + reply);
+            return total;
+          } else {
+            Preconditions.assertTrue(reply.isSuccess(),
+                    () -> "Failed to stream read " + path + ", reply=" + reply);
+            Preconditions.assertEquals(Type.STREAM_DATA, reply.getType(),
+                    "reply type for stream read " + path);
+            final ByteBuffer data = reply.nioBuffer();
+            while (data.hasRemaining()) {
+              total += channel.write(data);
+            }
+          }
+        } finally {
+          ref.release();
+        }
+      }
+    }
+  }
+
   public CompletableFuture<Long> writeAsync(String path, long offset, boolean close, ByteBuffer buffer, boolean sync) {
     return writeImpl(this::sendAsync, path, offset, close, buffer, sync
     ).thenApply(reply -> JavaUtils.supplyAndWrapAsCompletionException(

ratis-examples/src/main/java/org/apache/ratis/examples/filestore/FileStoreStateMachine.java

@@ -43,6 +43,7 @@
 import org.apache.ratis.util.FileUtils;
 
 import java.io.IOException;
+import java.nio.channels.WritableByteChannel;
 import java.nio.file.Path;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicReference;
@@ -106,6 +107,24 @@ public CompletableFuture<Message> query(Message request) {
         .thenApply(reply -> Message.valueOf(reply.toByteString()));
   }
 
+  @Override
+  public void query(Message request, WritableByteChannel stream) {
+    try {
+      final ReadRequestProto proto = ReadRequestProto.parseFrom(request.getContent());
+      if (proto.getIsWatch()) {
+        throw new IOException("Watch is not supported for streaming read: " + proto);
+      }
+      files.streamRead(proto.getPath().toStringUtf8(), proto.getOffset(), proto.getLength(), stream);
+    } catch (Exception e) {
+      LOG.error(getId() + ": Failed streaming read for " + request, e);
+      try {
+        stream.close();
+      } catch (IOException ignored) {
+        // ignore
+      }
+    }
+  }
+
   @Override
   public TransactionContext startTransaction(RaftClientRequest request) throws IOException {
     final ByteString content = request.getMessage().getContent();

ratis-examples/src/main/java/org/apache/ratis/examples/filestore/cli/FileStore.java

@@ -35,6 +35,7 @@ public static List<SubCommandBase> getSubCommands() {
     commands.add(new Server());
     commands.add(new LoadGen());
     commands.add(new DataStream());
+    commands.add(new Read());
     return commands;
   }
 }