|
20 | 20 | package org.elasticsearch.common.io.stream;
|
21 | 21 |
|
22 | 22 | 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; |
24 | 25 | import org.elasticsearch.common.util.BigArrays;
|
| 26 | +import org.elasticsearch.common.util.ByteArray; |
25 | 27 |
|
26 | 28 | /**
|
27 | 29 | * An bytes stream output that allows providing a {@link BigArrays} instance
|
28 | 30 | * expecting it to require releasing its content ({@link #bytes()}) once done.
|
29 | 31 | * <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. |
32 | 36 | */
|
33 |
| -public class ReleasableBytesStreamOutput extends BytesStreamOutput implements ReleasableBytesStream { |
| 37 | +public class ReleasableBytesStreamOutput extends BytesStreamOutput |
| 38 | + implements Releasable { |
| 39 | + |
| 40 | + private Releasable releasable; |
34 | 41 |
|
35 | 42 | public ReleasableBytesStreamOutput(BigArrays bigarrays) {
|
36 |
| - super(BigArrays.PAGE_SIZE_IN_BYTES, bigarrays); |
| 43 | + this(BigArrays.PAGE_SIZE_IN_BYTES, bigarrays); |
37 | 44 | }
|
38 | 45 |
|
39 | 46 | public ReleasableBytesStreamOutput(int expectedSize, BigArrays bigArrays) {
|
40 | 47 | super(expectedSize, bigArrays);
|
| 48 | + this.releasable = Releasables.releaseOnce(this.bytes); |
41 | 49 | }
|
42 | 50 |
|
| 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 | + */ |
43 | 56 | @Override
|
44 | 57 | public ReleasablePagedBytesReference bytes() {
|
45 |
| - return new ReleasablePagedBytesReference(bigArrays, bytes, count); |
| 58 | + return new ReleasablePagedBytesReference(bigArrays, bytes, count, releasable); |
46 | 59 | }
|
47 | 60 |
|
| 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 | + } |
48 | 85 | }
|
0 commit comments