org.elasticsearch.common.bytes.RecyclingBytesStreamOutput Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of elasticsearch Show documentation
Show all versions of elasticsearch Show documentation
Elasticsearch subproject :distribution:archives:integ-test-zip
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0 and the Server Side Public License, v 1; you may not use this file except
* in compliance with, at your election, the Elastic License 2.0 or the Server
* Side Public License, v 1.
*/
package org.elasticsearch.common.bytes;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.BytesRefIterator;
import org.elasticsearch.common.io.stream.BytesStream;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.util.BigArrays;
import org.elasticsearch.common.util.ByteArray;
import org.elasticsearch.core.Nullable;
import org.elasticsearch.core.internal.io.IOUtils;
import java.io.IOException;
import java.util.Objects;
/**
* An in-memory {@link StreamOutput} which first fills the given {@code byte[]} and then allocates more space from the given
* {@link BigArrays} if needed. The idea is that you can use this for passing data to an API that requires a single {@code byte[]} (or a
* {@link org.apache.lucene.util.BytesRef}) which you'd prefer to re-use if possible, avoiding excessive allocations, but which may not
* always be large enough.
*/
public class RecyclingBytesStreamOutput extends BytesStream {
private final byte[] buffer;
private final BigArrays bigArrays;
private int position;
@Nullable // if buffer is large enough
private ByteArray overflow;
public RecyclingBytesStreamOutput(byte[] buffer, BigArrays bigArrays) {
this.buffer = Objects.requireNonNull(buffer);
this.bigArrays = Objects.requireNonNull(bigArrays);
}
@Override
public void writeByte(byte b) {
if (position < buffer.length) {
buffer[position++] = b;
} else {
ensureCapacity(position + 1);
overflow.set(position++ - buffer.length, b);
}
}
private void ensureCapacity(int size) {
final int overflowSize = size - buffer.length;
assert overflowSize > 0 : "no need to ensureCapacity(" + size + ") with buffer of size [" + buffer.length + "]";
assert position >= buffer.length
: "no need to ensureCapacity(" + size + ") with buffer of size [" + buffer.length + "] at position [" + position + "]";
if (overflow == null) {
overflow = bigArrays.newByteArray(overflowSize, false);
} else if (overflowSize > overflow.size()) {
overflow = bigArrays.resize(overflow, overflowSize);
}
assert overflow.size() >= overflowSize;
}
@Override
public void writeBytes(byte[] b, int offset, int length) {
if (position < buffer.length) {
final int lengthForBuffer = Math.min(length, buffer.length - position);
System.arraycopy(b, offset, buffer, position, lengthForBuffer);
position += lengthForBuffer;
offset += lengthForBuffer;
length -= lengthForBuffer;
}
if (length > 0) {
ensureCapacity(position + length);
overflow.set(position - buffer.length, b, offset, length);
position += length;
}
}
@Override
public void flush() {}
@Override
public void close() throws IOException {
IOUtils.close(overflow);
}
@Override
public void reset() throws IOException {
throw new UnsupportedOperationException();
}
/**
* Return the written bytes in a {@link BytesRef}, avoiding allocating a new {@code byte[]} if the original buffer was already large
* enough. If we allocate a new (larger) buffer here then callers should typically re-use it for subsequent streams.
*/
public BytesRef toBytesRef() {
if (position <= buffer.length) {
assert overflow == null;
return new BytesRef(buffer, 0, position);
}
final byte[] newBuffer = new byte[position];
System.arraycopy(buffer, 0, newBuffer, 0, buffer.length);
int copyPos = buffer.length;
final BytesRefIterator iterator = BytesReference.fromByteArray(overflow, position - buffer.length).iterator();
BytesRef bytesRef;
try {
while ((bytesRef = iterator.next()) != null) {
assert copyPos + bytesRef.length <= position;
System.arraycopy(bytesRef.bytes, bytesRef.offset, newBuffer, copyPos, bytesRef.length);
copyPos += bytesRef.length;
}
} catch (IOException e) {
throw new AssertionError("impossible", e);
}
return new BytesRef(newBuffer, 0, position);
}
@Override
public BytesReference bytes() {
if (position <= buffer.length) {
assert overflow == null;
return new BytesArray(buffer, 0, position);
} else {
return CompositeBytesReference.of(
new BytesArray(buffer, 0, buffer.length),
BytesReference.fromByteArray(overflow, position - buffer.length)
);
}
}
}