org.apache.druid.segment.writeout.ByteBufferWriteOutBytes Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of druid-processing Show documentation
Show all versions of druid-processing Show documentation
A module that is everything required to understands Druid Segments
/*
* 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.druid.segment.writeout;
import com.google.common.base.Preconditions;
import com.google.common.io.ByteSource;
import com.google.common.primitives.Ints;
import org.apache.druid.io.ByteBufferInputStream;
import org.apache.druid.io.Channels;
import org.apache.druid.java.util.common.IAE;
import java.io.IOException;
import java.io.InputStream;
import java.nio.BufferUnderflowException;
import java.nio.ByteBuffer;
import java.nio.channels.WritableByteChannel;
import java.util.ArrayList;
import java.util.function.Function;
import java.util.stream.Collectors;
public abstract class ByteBufferWriteOutBytes extends WriteOutBytes
{
/**
* There is no reason why 64K is chosen. Other power of 2 between 4K and 64K could be more reasonable.
*/
static final int BUFFER_SIZE = 64 * 1024;
final ArrayList buffers = new ArrayList<>();
int headBufferIndex;
ByteBuffer headBuffer;
long size;
long capacity;
boolean open = true;
ByteBufferWriteOutBytes()
{
size = 0;
headBufferIndex = 0;
headBuffer = allocateBuffer();
buffers.add(headBuffer);
capacity = BUFFER_SIZE;
}
@Override
public long size()
{
return size;
}
protected abstract ByteBuffer allocateBuffer();
private void ensureCapacity(int len)
{
long remaining = capacity - size;
for (long toAllocate = len - remaining; toAllocate >= 0; toAllocate -= BUFFER_SIZE) {
buffers.add(allocateBuffer());
capacity += BUFFER_SIZE;
}
if (headBuffer.remaining() == 0) {
nextHead();
}
}
private void nextHead()
{
headBufferIndex++;
headBuffer = buffers.get(headBufferIndex);
}
@Override
public void write(int b)
{
checkOpen();
if (headBuffer.remaining() == 0) {
ensureCapacity(1);
}
headBuffer.put((byte) b);
size += 1;
}
@Override
public void writeInt(int v)
{
checkOpen();
if (headBuffer.remaining() >= Integer.BYTES) {
headBuffer.putInt(v);
size += Integer.BYTES;
} else {
ensureCapacity(Integer.BYTES);
if (headBuffer.remaining() >= Integer.BYTES) {
headBuffer.putInt(v);
size += Integer.BYTES;
} else {
write(v >> 24);
write(v >> 16);
write(v >> 8);
write(v);
}
}
}
@Override
public void write(byte[] b)
{
write0(b, 0, b.length);
}
@Override
public void write(byte[] b, int off, int len)
{
Preconditions.checkPositionIndexes(off, off + len, b.length);
write0(b, off, len);
}
private void write0(byte[] b, int off, int len)
{
checkOpen();
if (headBuffer.remaining() < len) {
ensureCapacity(len);
}
int headRemaining = headBuffer.remaining();
if (len <= headRemaining) {
headBuffer.put(b, off, len);
} else {
headBuffer.put(b, off, headRemaining);
int bytesLeft = len - headRemaining;
off += headRemaining;
for (; bytesLeft > 0; bytesLeft -= BUFFER_SIZE, off += BUFFER_SIZE) {
nextHead();
headBuffer.put(b, off, Math.min(BUFFER_SIZE, bytesLeft));
}
}
size += len;
}
@Override
public int write(ByteBuffer src)
{
checkOpen();
int len = src.remaining();
if (headBuffer.remaining() < len) {
ensureCapacity(len);
}
int headRemaining = headBuffer.remaining();
src.limit(src.position() + Math.min(headRemaining, len));
headBuffer.put(src);
for (int bytesLeft = len - headRemaining; bytesLeft > 0; bytesLeft -= BUFFER_SIZE) {
nextHead();
src.limit(src.position() + Math.min(BUFFER_SIZE, bytesLeft));
headBuffer.put(src);
}
size += len;
return len;
}
@Override
public void writeTo(WritableByteChannel channel) throws IOException
{
checkOpen();
for (int i = 0; i <= headBufferIndex; i++) {
ByteBuffer buffer = buffers.get(i);
buffer.flip();
Channels.writeFully(channel, buffer);
// switch back to the initial state
buffer.limit(buffer.capacity());
}
}
/**
* Takes all bytes that are written to this WriteOutBytes so far and writes them into the given ByteBuffer. This method
* changes the position of the out buffer by the {@link #size()} of this WriteOutBytes.
*
* @throws java.nio.BufferOverflowException if the {@link ByteBuffer#remaining()} capacity of the given buffer is
* smaller than the size of this WriteOutBytes
*/
public void writeTo(ByteBuffer out)
{
checkOpen();
for (int i = 0; i <= headBufferIndex; i++) {
ByteBuffer buffer = buffers.get(i);
buffer.flip();
out.put(buffer);
// switch back to the initial state
buffer.limit(buffer.capacity());
}
}
@Override
public void readFully(long pos, ByteBuffer buffer)
{
checkOpen();
if (pos < 0 || pos > size) {
throw new IAE("pos %d out of range [%d, %d]", pos, 0, size);
}
int ourBufferIndex = Ints.checkedCast(pos / BUFFER_SIZE);
int ourBufferOffset = Ints.checkedCast(pos % BUFFER_SIZE);
for (int bytesLeft = buffer.remaining(); bytesLeft > 0;) {
int bytesToWrite = Math.min(BUFFER_SIZE - ourBufferOffset, bytesLeft);
ByteBuffer ourBuffer = buffers.get(ourBufferIndex);
int ourBufferPosition = ourBuffer.position();
if (bytesToWrite > ourBufferPosition - ourBufferOffset) {
throw new BufferUnderflowException();
}
try {
ourBuffer.position(ourBufferOffset);
ourBuffer.limit(ourBufferOffset + bytesToWrite);
buffer.put(ourBuffer);
}
finally {
// switch back to the initial state
ourBuffer.limit(ourBuffer.capacity());
ourBuffer.position(ourBufferPosition);
}
ourBufferIndex++;
ourBufferOffset = 0;
bytesLeft -= bytesToWrite;
}
}
@Override
public InputStream asInputStream() throws IOException
{
checkOpen();
Function byteBufferToByteSource = buf -> new ByteSource()
{
@Override
public InputStream openStream()
{
ByteBuffer inputBuf = buf.duplicate();
inputBuf.flip();
return new ByteBufferInputStream(inputBuf);
}
};
return ByteSource.concat(buffers.stream().map(byteBufferToByteSource).collect(Collectors.toList())).openStream();
}
@Override
public boolean isOpen()
{
return open;
}
public void free()
{
open = false;
buffers.clear();
headBufferIndex = -1;
headBuffer = null;
size = 0;
capacity = 0;
}
private void checkOpen()
{
if (!isOpen()) {
throw new IllegalStateException();
}
}
}