/*
* Licensed 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 com.facebook.presto.rcfile;
import com.google.common.collect.ImmutableList;
import io.airlift.slice.Slice;
import io.airlift.slice.SliceOutput;
import io.airlift.slice.Slices;
import org.openjdk.jol.info.ClassLayout;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import static com.google.common.base.Preconditions.checkArgument;
import static io.airlift.slice.SizeOf.SIZE_OF_BYTE;
import static io.airlift.slice.SizeOf.SIZE_OF_INT;
import static io.airlift.slice.SizeOf.SIZE_OF_LONG;
import static io.airlift.slice.SizeOf.SIZE_OF_SHORT;
import static java.lang.Math.min;
import static java.lang.Math.multiplyExact;
import static java.lang.Math.toIntExact;
public final class ChunkedSliceOutput
extends SliceOutput
{
private static final int INSTANCE_SIZE = ClassLayout.parseClass(ChunkedSliceOutput.class).instanceSize();
private static final int MINIMUM_CHUNK_SIZE = 4096;
private static final int MAXIMUM_CHUNK_SIZE = 16 * 1024 * 1024;
// This must not be larger than MINIMUM_CHUNK_SIZE/2
private static final int MAX_UNUSED_BUFFER_SIZE = 128;
private final ChunkSupplier chunkSupplier;
private Slice slice;
private byte[] buffer;
private final List<Slice> closedSlices = new ArrayList<>();
private long closedSlicesRetainedSize;
/**
* Offset of buffer within stream.
*/
private long streamOffset;
/**
* Current position for writing in buffer.
*/
private int bufferPosition;
public ChunkedSliceOutput(int minChunkSize, int maxChunkSize)
{
this.chunkSupplier = new ChunkSupplier(minChunkSize, maxChunkSize);
this.buffer = chunkSupplier.get();
this.slice = Slices.wrappedBuffer(buffer);
}
public List<Slice> getSlices()
{
return ImmutableList.<Slice>builder()
.addAll(closedSlices)
.add(Slices.copyOf(slice, 0, bufferPosition))
.build();
}
@Override
public void reset()
{
chunkSupplier.reset();
closedSlices.clear();
buffer = chunkSupplier.get();
slice = Slices.wrappedBuffer(buffer);
closedSlicesRetainedSize = 0;
streamOffset = 0;
bufferPosition = 0;
}
@Override
public void reset(int position)
{
throw new UnsupportedOperationException();
}
@Override
public int size()
{
return toIntExact(streamOffset + bufferPosition);
}
@Override
public int getRetainedSize()
{
return toIntExact(slice.getRetainedSize() + closedSlicesRetainedSize + INSTANCE_SIZE);
}
@Override
public int writableBytes()
{
return Integer.MAX_VALUE;
}
@Override
public boolean isWritable()
{
return true;
}
@Override
public void writeByte(int value)
{
ensureWritableBytes(SIZE_OF_BYTE);
slice.setByte(bufferPosition, value);
bufferPosition += SIZE_OF_BYTE;
}
@Override
public void writeShort(int value)
{
ensureWritableBytes(SIZE_OF_SHORT);
slice.setShort(bufferPosition, value);
bufferPosition += SIZE_OF_SHORT;
}
@Override
public void writeInt(int value)
{
ensureWritableBytes(SIZE_OF_INT);
slice.setInt(bufferPosition, value);
bufferPosition += SIZE_OF_INT;
}
@Override
public void writeLong(long value)
{
ensureWritableBytes(SIZE_OF_LONG);
slice.setLong(bufferPosition, value);
bufferPosition += SIZE_OF_LONG;
}
@Override
public void writeFloat(float value)
{
writeInt(Float.floatToIntBits(value));
}
@Override
public void writeDouble(double value)
{
writeLong(Double.doubleToLongBits(value));
}
@Override
public void writeBytes(Slice source)
{
writeBytes(source, 0, source.length());
}
@Override
public void writeBytes(Slice source, int sourceIndex, int length)
{
while (length > 0) {
int batch = tryEnsureBatchSize(length);
slice.setBytes(bufferPosition, source, sourceIndex, batch);
bufferPosition += batch;
sourceIndex += batch;
length -= batch;
}
}
@Override
public void writeBytes(byte[] source)
{
writeBytes(source, 0, source.length);
}
@Override
public void writeBytes(byte[] source, int sourceIndex, int length)
{
while (length > 0) {
int batch = tryEnsureBatchSize(length);
slice.setBytes(bufferPosition, source, sourceIndex, batch);
bufferPosition += batch;
sourceIndex += batch;
length -= batch;
}
}
@Override
public void writeBytes(InputStream in, int length)
throws IOException
{
while (length > 0) {
int batch = tryEnsureBatchSize(length);
slice.setBytes(bufferPosition, in, batch);
bufferPosition += batch;
length -= batch;
}
}
@Override
public void writeZero(int length)
{
checkArgument(length >= 0, "length must be greater than or equal to 0");
while (length > 0) {
int batch = tryEnsureBatchSize(length);
Arrays.fill(buffer, bufferPosition, bufferPosition + batch, (byte) 0);
bufferPosition += batch;
length -= batch;
}
}
@Override
public SliceOutput appendLong(long value)
{
writeLong(value);
return this;
}
@Override
public SliceOutput appendDouble(double value)
{
writeDouble(value);
return this;
}
@Override
public SliceOutput appendInt(int value)
{
writeInt(value);
return this;
}
@Override
public SliceOutput appendShort(int value)
{
writeShort(value);
return this;
}
@Override
public SliceOutput appendByte(int value)
{
writeByte(value);
return this;
}
@Override
public SliceOutput appendBytes(byte[] source, int sourceIndex, int length)
{
writeBytes(source, sourceIndex, length);
return this;
}
@Override
public SliceOutput appendBytes(byte[] source)
{
writeBytes(source);
return this;
}
@Override
public SliceOutput appendBytes(Slice slice)
{
writeBytes(slice);
return this;
}
@Override
public Slice slice()
{
throw new UnsupportedOperationException();
}
@Override
public Slice getUnderlyingSlice()
{
throw new UnsupportedOperationException();
}
@Override
public String toString(Charset charset)
{
return toString();
}
@Override
public String toString()
{
StringBuilder builder = new StringBuilder("OutputStreamSliceOutputAdapter{");
builder.append("position=").append(size());
builder.append("bufferSize=").append(slice.length());
builder.append('}');
return builder.toString();
}
private int tryEnsureBatchSize(int length)
{
ensureWritableBytes(min(MAX_UNUSED_BUFFER_SIZE, length));
return min(length, slice.length() - bufferPosition);
}
private void ensureWritableBytes(int minWritableBytes)
{
checkArgument(minWritableBytes <= MAX_UNUSED_BUFFER_SIZE);
if (bufferPosition + minWritableBytes > slice.length()) {
closeChunk();
}
}
private void closeChunk()
{
// add trimmed view of slice to closed slices
closedSlices.add(slice.slice(0, bufferPosition));
// create a new buffer
// double size until we hit the max chunk size
buffer = chunkSupplier.get();
slice = Slices.wrappedBuffer(buffer);
streamOffset += bufferPosition;
bufferPosition = 0;
}
// Chunk supplier creates buffers by doubling the size from min to max chunk size.
// The supplier also tracks all created buffers and can be reset to the beginning,
// reusing the buffers.
private static class ChunkSupplier
{
private final int maxChunkSize;
private final List<byte[]> bufferPool = new ArrayList<>();
private final List<byte[]> usedBuffers = new ArrayList<>();
private int currentSize;
public ChunkSupplier(int minChunkSize, int maxChunkSize)
{
checkArgument(minChunkSize >= MINIMUM_CHUNK_SIZE, "minimum chunk size of " + MINIMUM_CHUNK_SIZE + " required");
checkArgument(maxChunkSize <= MAXIMUM_CHUNK_SIZE, "maximum chunk size of " + MAXIMUM_CHUNK_SIZE + " required");
checkArgument(minChunkSize <= maxChunkSize, "minimum chunk size must be less than maximum chunk size");
this.currentSize = minChunkSize;
this.maxChunkSize = maxChunkSize;
}
public void reset()
{
bufferPool.addAll(0, usedBuffers);
usedBuffers.clear();
}
public byte[] get()
{
byte[] buffer;
if (bufferPool.isEmpty()) {
currentSize = min(multiplyExact(currentSize, 2), maxChunkSize);
buffer = new byte[currentSize];
}
else {
buffer = bufferPool.remove(0);
currentSize = buffer.length;
}
usedBuffers.add(buffer);
return buffer;
}
}
}