Skip to content

Commit 30ab873

Browse files
authored
Closing a ReleasableBytesStreamOutput closes the underlying BigArray (#23941)
This commit makes closing a ReleasableBytesStreamOutput release the underlying BigArray so that we can use try-with-resources with these streams and avoid leaking memory by not returning the BigArray. As part of this change, the ReleasableBytesStreamOutput adds protection to only release the BigArray once. In order to make some of the changes cleaner, the ReleasableBytesStream interface has been removed. The BytesStream interface is changed to a abstract class so that we can use it as a useable return type for a new method, Streams#flushOnCloseStream. This new method wraps a given stream and overrides the close method so that the stream is simply flushed and not closed. This behavior is used in the TcpTransport when compression is used with a ReleasableBytesStreamOutput as we need to close the compressed stream to ensure all of the data is written from this stream. Closing the compressed stream will try to close the underlying stream but we only want to flush so that all of the written bytes are available. Additionally, an error message method added in the BytesRestResponse did not use a builder provided by the channel and instead created its own JSON builder. This changes that method to use the channel builder and in turn the bytes stream output that is managed by the channel. Note, this commit differs from 6bfecdf in that it updates ReleasableBytesStreamOutput to handle the case of the BigArray decreasing in size, which changes the reference to the BigArray. When the reference is changed, the releasable needs to be updated otherwise there could be a leak of bytes and corruption of data in unrelated streams. This reverts commit afd45c1, which reverted #23572.
1 parent e3aa2a8 commit 30ab873

File tree

17 files changed

+274
-88
lines changed

17 files changed

+274
-88
lines changed

core/src/main/java/org/elasticsearch/common/bytes/ReleasablePagedBytesReference.java

+6-2
Original file line numberDiff line numberDiff line change
@@ -30,13 +30,17 @@
3030
*/
3131
public final class ReleasablePagedBytesReference extends PagedBytesReference implements Releasable {
3232

33-
public ReleasablePagedBytesReference(BigArrays bigarrays, ByteArray byteArray, int length) {
33+
private final Releasable releasable;
34+
35+
public ReleasablePagedBytesReference(BigArrays bigarrays, ByteArray byteArray, int length,
36+
Releasable releasable) {
3437
super(bigarrays, byteArray, length);
38+
this.releasable = releasable;
3539
}
3640

3741
@Override
3842
public void close() {
39-
Releasables.close(byteArray);
43+
Releasables.close(releasable);
4044
}
4145

4246
}

core/src/main/java/org/elasticsearch/common/compress/Compressor.java

+5
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
package org.elasticsearch.common.compress;
2121

2222
import org.elasticsearch.common.bytes.BytesReference;
23+
import org.elasticsearch.common.io.stream.ReleasableBytesStreamOutput;
2324
import org.elasticsearch.common.io.stream.StreamInput;
2425
import org.elasticsearch.common.io.stream.StreamOutput;
2526

@@ -31,5 +32,9 @@ public interface Compressor {
3132

3233
StreamInput streamInput(StreamInput in) throws IOException;
3334

35+
/**
36+
* Creates a new stream output that compresses the contents and writes to the provided stream
37+
* output. Closing the returned {@link StreamOutput} will close the provided stream output.
38+
*/
3439
StreamOutput streamOutput(StreamOutput out) throws IOException;
3540
}

core/src/main/java/org/elasticsearch/common/compress/DeflateCompressor.java

+5-4
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@
2020
package org.elasticsearch.common.compress;
2121

2222
import org.elasticsearch.common.bytes.BytesReference;
23-
import org.elasticsearch.common.compress.Compressor;
2423
import org.elasticsearch.common.io.stream.InputStreamStreamInput;
2524
import org.elasticsearch.common.io.stream.OutputStreamStreamOutput;
2625
import org.elasticsearch.common.io.stream.StreamInput;
@@ -47,7 +46,7 @@ public class DeflateCompressor implements Compressor {
4746
// It needs to be different from other compressors and to not be specific
4847
// enough so that no stream starting with these bytes could be detected as
4948
// a XContent
50-
private static final byte[] HEADER = new byte[] { 'D', 'F', 'L', '\0' };
49+
private static final byte[] HEADER = new byte[]{'D', 'F', 'L', '\0'};
5150
// 3 is a good trade-off between speed and compression ratio
5251
private static final int LEVEL = 3;
5352
// We use buffering on the input and output of in/def-laters in order to
@@ -88,6 +87,7 @@ public StreamInput streamInput(StreamInput in) throws IOException {
8887
decompressedIn = new BufferedInputStream(decompressedIn, BUFFER_SIZE);
8988
return new InputStreamStreamInput(decompressedIn) {
9089
final AtomicBoolean closed = new AtomicBoolean(false);
90+
9191
public void close() throws IOException {
9292
try {
9393
super.close();
@@ -107,10 +107,11 @@ public StreamOutput streamOutput(StreamOutput out) throws IOException {
107107
final boolean nowrap = true;
108108
final Deflater deflater = new Deflater(LEVEL, nowrap);
109109
final boolean syncFlush = true;
110-
OutputStream compressedOut = new DeflaterOutputStream(out, deflater, BUFFER_SIZE, syncFlush);
111-
compressedOut = new BufferedOutputStream(compressedOut, BUFFER_SIZE);
110+
DeflaterOutputStream deflaterOutputStream = new DeflaterOutputStream(out, deflater, BUFFER_SIZE, syncFlush);
111+
OutputStream compressedOut = new BufferedOutputStream(deflaterOutputStream, BUFFER_SIZE);
112112
return new OutputStreamStreamOutput(compressedOut) {
113113
final AtomicBoolean closed = new AtomicBoolean(false);
114+
114115
public void close() throws IOException {
115116
try {
116117
super.close();

core/src/main/java/org/elasticsearch/common/io/ReleasableBytesStream.java

-32
This file was deleted.

core/src/main/java/org/elasticsearch/common/io/Streams.java

+55
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,9 @@
2020
package org.elasticsearch.common.io;
2121

2222
import org.apache.lucene.util.IOUtils;
23+
import org.elasticsearch.common.bytes.BytesReference;
24+
import org.elasticsearch.common.io.stream.BytesStream;
25+
import org.elasticsearch.common.io.stream.StreamOutput;
2326
import org.elasticsearch.common.util.Callback;
2427

2528
import java.io.BufferedReader;
@@ -236,4 +239,56 @@ public static void readAllLines(InputStream input, Callback<String> callback) th
236239
}
237240
}
238241
}
242+
243+
/**
244+
* Wraps the given {@link BytesStream} in a {@link StreamOutput} that simply flushes when
245+
* close is called.
246+
*/
247+
public static BytesStream flushOnCloseStream(BytesStream os) {
248+
return new FlushOnCloseOutputStream(os);
249+
}
250+
251+
/**
252+
* A wrapper around a {@link BytesStream} that makes the close operation a flush. This is
253+
* needed as sometimes a stream will be closed but the bytes that the stream holds still need
254+
* to be used and the stream cannot be closed until the bytes have been consumed.
255+
*/
256+
private static class FlushOnCloseOutputStream extends BytesStream {
257+
258+
private final BytesStream delegate;
259+
260+
private FlushOnCloseOutputStream(BytesStream bytesStreamOutput) {
261+
this.delegate = bytesStreamOutput;
262+
}
263+
264+
@Override
265+
public void writeByte(byte b) throws IOException {
266+
delegate.writeByte(b);
267+
}
268+
269+
@Override
270+
public void writeBytes(byte[] b, int offset, int length) throws IOException {
271+
delegate.writeBytes(b, offset, length);
272+
}
273+
274+
@Override
275+
public void flush() throws IOException {
276+
delegate.flush();
277+
}
278+
279+
@Override
280+
public void close() throws IOException {
281+
flush();
282+
}
283+
284+
@Override
285+
public void reset() throws IOException {
286+
delegate.reset();
287+
}
288+
289+
@Override
290+
public BytesReference bytes() {
291+
return delegate.bytes();
292+
}
293+
}
239294
}

core/src/main/java/org/elasticsearch/common/io/BytesStream.java core/src/main/java/org/elasticsearch/common/io/stream/BytesStream.java

+4-4
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,11 @@
1717
* under the License.
1818
*/
1919

20-
package org.elasticsearch.common.io;
20+
package org.elasticsearch.common.io.stream;
2121

2222
import org.elasticsearch.common.bytes.BytesReference;
2323

24-
public interface BytesStream {
24+
public abstract class BytesStream extends StreamOutput {
2525

26-
BytesReference bytes();
27-
}
26+
public abstract BytesReference bytes();
27+
}

core/src/main/java/org/elasticsearch/common/io/stream/BytesStreamOutput.java

+4-5
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121

2222
import org.elasticsearch.common.bytes.BytesReference;
2323
import org.elasticsearch.common.bytes.PagedBytesReference;
24-
import org.elasticsearch.common.io.BytesStream;
2524
import org.elasticsearch.common.util.BigArrays;
2625
import org.elasticsearch.common.util.ByteArray;
2726

@@ -31,7 +30,7 @@
3130
* A @link {@link StreamOutput} that uses {@link BigArrays} to acquire pages of
3231
* bytes, which avoids frequent reallocation &amp; copying of the internal data.
3332
*/
34-
public class BytesStreamOutput extends StreamOutput implements BytesStream {
33+
public class BytesStreamOutput extends BytesStream {
3534

3635
protected final BigArrays bigArrays;
3736

@@ -50,7 +49,7 @@ public BytesStreamOutput() {
5049
/**
5150
* Create a non recycling {@link BytesStreamOutput} with enough initial pages acquired
5251
* to satisfy the capacity given by expected size.
53-
*
52+
*
5453
* @param expectedSize the expected maximum size of the stream in bytes.
5554
*/
5655
public BytesStreamOutput(int expectedSize) {
@@ -129,7 +128,7 @@ public void close() {
129128

130129
/**
131130
* Returns the current size of the buffer.
132-
*
131+
*
133132
* @return the value of the <code>count</code> field, which is the number of valid
134133
* bytes in this output stream.
135134
* @see java.io.ByteArrayOutputStream#count
@@ -151,7 +150,7 @@ public long ramBytesUsed() {
151150
return bytes.ramBytesUsed();
152151
}
153152

154-
private void ensureCapacity(long offset) {
153+
void ensureCapacity(long offset) {
155154
if (offset > Integer.MAX_VALUE) {
156155
throw new IllegalArgumentException(getClass().getSimpleName() + " cannot hold more than 2GB of data");
157156
}

core/src/main/java/org/elasticsearch/common/io/stream/ReleasableBytesStreamOutput.java

+43-6
Original file line numberDiff line numberDiff line change
@@ -20,29 +20,66 @@
2020
package org.elasticsearch.common.io.stream;
2121

2222
import org.elasticsearch.common.bytes.ReleasablePagedBytesReference;
23-
import org.elasticsearch.common.io.ReleasableBytesStream;
23+
import org.elasticsearch.common.lease.Releasable;
24+
import org.elasticsearch.common.lease.Releasables;
2425
import org.elasticsearch.common.util.BigArrays;
26+
import org.elasticsearch.common.util.ByteArray;
2527

2628
/**
2729
* An bytes stream output that allows providing a {@link BigArrays} instance
2830
* expecting it to require releasing its content ({@link #bytes()}) once done.
2931
* <p>
30-
* Please note, its is the responsibility of the caller to make sure the bytes
31-
* reference do not "escape" and are released only once.
32+
* Please note, closing this stream will release the bytes that are in use by any
33+
* {@link ReleasablePagedBytesReference} returned from {@link #bytes()}, so this
34+
* stream should only be closed after the bytes have been output or copied
35+
* elsewhere.
3236
*/
33-
public class ReleasableBytesStreamOutput extends BytesStreamOutput implements ReleasableBytesStream {
37+
public class ReleasableBytesStreamOutput extends BytesStreamOutput
38+
implements Releasable {
39+
40+
private Releasable releasable;
3441

3542
public ReleasableBytesStreamOutput(BigArrays bigarrays) {
36-
super(BigArrays.PAGE_SIZE_IN_BYTES, bigarrays);
43+
this(BigArrays.PAGE_SIZE_IN_BYTES, bigarrays);
3744
}
3845

3946
public ReleasableBytesStreamOutput(int expectedSize, BigArrays bigArrays) {
4047
super(expectedSize, bigArrays);
48+
this.releasable = Releasables.releaseOnce(this.bytes);
4149
}
4250

51+
/**
52+
* Returns a {@link Releasable} implementation of a
53+
* {@link org.elasticsearch.common.bytes.BytesReference} that represents the current state of
54+
* the bytes in the stream.
55+
*/
4356
@Override
4457
public ReleasablePagedBytesReference bytes() {
45-
return new ReleasablePagedBytesReference(bigArrays, bytes, count);
58+
return new ReleasablePagedBytesReference(bigArrays, bytes, count, releasable);
4659
}
4760

61+
@Override
62+
public void close() {
63+
Releasables.close(releasable);
64+
}
65+
66+
@Override
67+
void ensureCapacity(long offset) {
68+
final ByteArray prevBytes = this.bytes;
69+
super.ensureCapacity(offset);
70+
if (prevBytes != this.bytes) {
71+
// re-create the releasable with the new reference
72+
releasable = Releasables.releaseOnce(this.bytes);
73+
}
74+
}
75+
76+
@Override
77+
public void reset() {
78+
final ByteArray prevBytes = this.bytes;
79+
super.reset();
80+
if (prevBytes != this.bytes) {
81+
// re-create the releasable with the new reference
82+
releasable = Releasables.releaseOnce(this.bytes);
83+
}
84+
}
4885
}

core/src/main/java/org/elasticsearch/common/xcontent/XContentBuilder.java

+2-3
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
import org.apache.lucene.util.BytesRef;
2323
import org.elasticsearch.common.bytes.BytesReference;
2424
import org.elasticsearch.common.geo.GeoPoint;
25-
import org.elasticsearch.common.io.BytesStream;
25+
import org.elasticsearch.common.io.stream.BytesStream;
2626
import org.elasticsearch.common.io.stream.BytesStreamOutput;
2727
import org.elasticsearch.common.lease.Releasable;
2828
import org.elasticsearch.common.text.Text;
@@ -53,7 +53,7 @@
5353
/**
5454
* A utility to build XContent (ie json).
5555
*/
56-
public final class XContentBuilder implements BytesStream, Releasable, Flushable {
56+
public final class XContentBuilder implements Releasable, Flushable {
5757

5858
/**
5959
* Create a new {@link XContentBuilder} using the given {@link XContent} content.
@@ -1041,7 +1041,6 @@ public XContentGenerator generator() {
10411041
return this.generator;
10421042
}
10431043

1044-
@Override
10451044
public BytesReference bytes() {
10461045
close();
10471046
return ((BytesStream) bos).bytes();

core/src/main/java/org/elasticsearch/index/translog/Translog.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -439,7 +439,7 @@ public Location add(final Operation operation) throws IOException {
439439
}
440440
throw new TranslogException(shardId, "Failed to write operation [" + operation + "]", e);
441441
} finally {
442-
Releasables.close(out.bytes());
442+
Releasables.close(out);
443443
}
444444
}
445445

@@ -1332,7 +1332,7 @@ public static void writeOperations(StreamOutput outStream, List<Operation> toWri
13321332
bytes.writeTo(outStream);
13331333
}
13341334
} finally {
1335-
Releasables.close(out.bytes());
1335+
Releasables.close(out);
13361336
}
13371337

13381338
}

0 commit comments

Comments
 (0)