Replacing SeekableByteChannelChunkedNioStream

Author: dlarge

src/main/java/com/emc/mongoose/storage/driver/coop/netty/NettyStorageDriverBase.java

@@ -21,7 +21,6 @@
 import com.emc.mongoose.base.logging.Loggers;
 import com.emc.mongoose.storage.driver.coop.CoopStorageDriverBase;
 import com.emc.mongoose.storage.driver.coop.netty.data.DataItemFileRegion;
-import com.emc.mongoose.storage.driver.coop.netty.data.PartialChunkedNioStream;
 import com.emc.mongoose.storage.driver.coop.netty.data.SeekableByteChannelChunkedNioStream;
 import com.github.akurilov.commons.collection.Range;
 import com.github.akurilov.commons.concurrent.ThreadUtil;
@@ -461,7 +460,7 @@ protected final void sendRequestData(final Channel channel, final O op) throws I
 					final var srcPath = dataOp.srcPath();
 					if (0 < dataItem.size() && (null == srcPath || srcPath.isEmpty())) {
 						if (sslFlag) {
-							channel.write(new PartialChunkedNioStream(dataItem));
+							channel.write(new SeekableByteChannelChunkedNioStream(dataItem));
 						} else {
 							channel.write(new DataItemFileRegion(dataItem));
 						}

src/main/java/com/emc/mongoose/storage/driver/coop/netty/data/PartialChunkedNioStream.java

@@ -2,36 +2,109 @@
 
 import static com.emc.mongoose.base.storage.driver.StorageDriver.BUFF_SIZE_MAX;
 
-import io.netty.handler.stream.ChunkedNioStream;
-
 import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.ReadableByteChannel;
 import java.nio.channels.SeekableByteChannel;
 
-/**
-Created by andrey on 24.04.17.
-*/
-public final class SeekableByteChannelChunkedNioStream
-				extends ChunkedNioStream {
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.handler.stream.ChunkedInput;
 
+/**
+ * Modified version of io.netty.handler.stream.ChunkedNioStream.
+ * Ensures that the final call to readChunk() writes only a partial chunk to the
+ * given byte channel and isEndOfInput() causes the stream to end.
+ */
+public final class SeekableByteChannelChunkedNioStream implements ChunkedInput<ByteBuf> {
+	private final ReadableByteChannel in;
+	private final int chunkSize;
+	private final ByteBuffer byteBuffer;
 	private final long sizeToTransfer;
 
-	public SeekableByteChannelChunkedNioStream(final SeekableByteChannel sbc)
-					throws IOException {
+	private long bytesTransferred = 0;
+
+	public SeekableByteChannelChunkedNioStream(final SeekableByteChannel sbc) throws IOException {
 		this(sbc, sbc.size());
 	}
 
 	private SeekableByteChannelChunkedNioStream(final SeekableByteChannel sbc, final long sizeToTransfer) {
-		super(sbc, sizeToTransfer > BUFF_SIZE_MAX ? BUFF_SIZE_MAX : (int) sizeToTransfer);
+		this.in = sbc;
+		this.chunkSize = (int) (sizeToTransfer > BUFF_SIZE_MAX ? BUFF_SIZE_MAX : sizeToTransfer);
+		this.byteBuffer = ByteBuffer.allocate(chunkSize);
 		this.sizeToTransfer = sizeToTransfer;
 	}
 
+	@Override
+	public final boolean isEndOfInput() {
+		return bytesTransferred == sizeToTransfer;
+	}
+
+	@Override
+	public void close() throws Exception {
+		in.close();
+	}
+
+	@Deprecated
+	@Override
+	public ByteBuf readChunk(ChannelHandlerContext ctx) throws Exception {
+		return readChunk(ctx.alloc());
+	}
+
+	@Override
+	public ByteBuf readChunk(ByteBufAllocator allocator) throws Exception {
+		if (isEndOfInput()) {
+			return null;
+		}
+
+		int nextChunkSize = chunkSize;
+		long bytesRemaining = length() - progress();
+
+		// Is there less than a chunk size of data remaining?
+		if (bytesRemaining < chunkSize) {
+			// Limit the byte buffer and next chunk size
+			byteBuffer.limit((int) bytesRemaining);
+			nextChunkSize = (int) bytesRemaining;
+		}
+
+		// Read into the byte buffer
+		int readBytes = byteBuffer.position();
+		while (true) {
+			int localReadBytes = in.read(byteBuffer);
+			if (localReadBytes < 0) {
+				break;
+			}
+			readBytes += localReadBytes;
+			bytesTransferred += localReadBytes;
+			if (readBytes == nextChunkSize) {
+				break;
+			}
+		}
+
+		// Write from the byte buffer
+		byteBuffer.flip();
+		boolean release = true;
+		ByteBuf buffer = allocator.buffer(byteBuffer.remaining());
+		try {
+			buffer.writeBytes(byteBuffer);
+			byteBuffer.clear();
+			release = false;
+			return buffer;
+		} finally {
+			if (release) {
+				buffer.release();
+			}
+		}
+	}
+
 	@Override
 	public long length() {
 		return sizeToTransfer;
 	}
 
 	@Override
-	public final boolean isEndOfInput() {
-		return sizeToTransfer == transferredBytes();
+	public long progress() {
+		return bytesTransferred;
 	}
-}
+}