|
|
@@ -19,12 +19,14 @@
|
|
|
|
|
|
package org.elasticsearch.common.compress;
|
|
|
|
|
|
+import org.elasticsearch.Assertions;
|
|
|
import org.elasticsearch.common.bytes.BytesReference;
|
|
|
import org.elasticsearch.common.io.stream.BytesStreamOutput;
|
|
|
import org.elasticsearch.common.io.stream.InputStreamStreamInput;
|
|
|
import org.elasticsearch.common.io.stream.OutputStreamStreamOutput;
|
|
|
import org.elasticsearch.common.io.stream.StreamInput;
|
|
|
import org.elasticsearch.common.io.stream.StreamOutput;
|
|
|
+import org.elasticsearch.common.lease.Releasable;
|
|
|
|
|
|
import java.io.BufferedInputStream;
|
|
|
import java.io.BufferedOutputStream;
|
|
|
@@ -32,7 +34,6 @@ import java.io.IOException;
|
|
|
import java.io.InputStream;
|
|
|
import java.io.OutputStream;
|
|
|
import java.util.Arrays;
|
|
|
-import java.util.concurrent.atomic.AtomicBoolean;
|
|
|
import java.util.zip.Deflater;
|
|
|
import java.util.zip.DeflaterOutputStream;
|
|
|
import java.util.zip.Inflater;
|
|
|
@@ -73,8 +74,75 @@ public class DeflateCompressor implements Compressor {
|
|
|
return HEADER.length;
|
|
|
}
|
|
|
|
|
|
+ // Reusable inflater reference for streaming decompression
|
|
|
+ private static final ThreadLocal<ReleasableReference<Inflater>> inflaterForStreamRef = ThreadLocal.withInitial(() -> {
|
|
|
+ final Inflater inflater = new Inflater(true);
|
|
|
+ return new ReleasableReference<>(inflater, inflater::reset);
|
|
|
+ });
|
|
|
+
|
|
|
+ // Reusable deflater reference for streaming compression
|
|
|
+ private static final ThreadLocal<ReleasableReference<Deflater>> deflaterForStreamRef = ThreadLocal.withInitial(() -> {
|
|
|
+ final Deflater deflater = new Deflater(LEVEL, true);
|
|
|
+ return new ReleasableReference<>(deflater, deflater::reset);
|
|
|
+ });
|
|
|
+
|
|
|
+ // Reference to a deflater or inflater that is used to make sure we do not use the same stream twice when nesting streams.
|
|
|
+ private static final class ReleasableReference<T> implements Releasable {
|
|
|
+
|
|
|
+ protected final T resource;
|
|
|
+
|
|
|
+ private final Releasable releasable;
|
|
|
+
|
|
|
+ // Thread that is currently using this reference. Only used for assertions and only assigned if assertions are enabled.
|
|
|
+ private Thread thread = null;
|
|
|
+
|
|
|
+ // true if this reference is currently in use and is not available for re-use
|
|
|
+ boolean inUse;
|
|
|
+
|
|
|
+ protected ReleasableReference(T resource, Releasable releasable) {
|
|
|
+ this.resource = resource;
|
|
|
+ this.releasable = releasable;
|
|
|
+ }
|
|
|
+
|
|
|
+ T get() {
|
|
|
+ if (Assertions.ENABLED) {
|
|
|
+ assert thread == null;
|
|
|
+ thread = Thread.currentThread();
|
|
|
+ }
|
|
|
+ assert inUse == false;
|
|
|
+ inUse = true;
|
|
|
+ return resource;
|
|
|
+ }
|
|
|
+
|
|
|
+ @Override
|
|
|
+ public void close() {
|
|
|
+ if (Assertions.ENABLED) {
|
|
|
+ assert thread == Thread.currentThread() :
|
|
|
+ "Opened on [" + thread.getName() + "] but closed on [" + Thread.currentThread().getName() + "]";
|
|
|
+ thread = null;
|
|
|
+ }
|
|
|
+ assert inUse;
|
|
|
+ inUse = false;
|
|
|
+ releasable.close();
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
@Override
|
|
|
- public StreamInput streamInput(StreamInput in) throws IOException {
|
|
|
+ public StreamInput threadLocalStreamInput(StreamInput in) throws IOException {
|
|
|
+ return new InputStreamStreamInput(inputStream(in, true));
|
|
|
+ }
|
|
|
+
|
|
|
+ /**
|
|
|
+ * Creates a new input stream that decompresses the contents read from the provided input stream.
|
|
|
+ * Closing the returned stream will close the provided input stream.
|
|
|
+ * Optionally uses thread-local, pooled resources to save off-heap allocations if the stream is guaranteed to not escape the current
|
|
|
+ * thread.
|
|
|
+ *
|
|
|
+ * @param in input stream to wrap
|
|
|
+ * @param threadLocal whether this stream will only be used on the current thread or not
|
|
|
+ * @return decompressing stream
|
|
|
+ */
|
|
|
+ public static InputStream inputStream(InputStream in, boolean threadLocal) throws IOException {
|
|
|
final byte[] headerBytes = new byte[HEADER.length];
|
|
|
int len = 0;
|
|
|
while (len < headerBytes.length) {
|
|
|
@@ -88,56 +156,72 @@ public class DeflateCompressor implements Compressor {
|
|
|
throw new IllegalArgumentException("Input stream is not compressed with DEFLATE!");
|
|
|
}
|
|
|
|
|
|
- final boolean nowrap = true;
|
|
|
- final Inflater inflater = new Inflater(nowrap);
|
|
|
- InputStream decompressedIn = new InflaterInputStream(in, inflater, BUFFER_SIZE);
|
|
|
- decompressedIn = new BufferedInputStream(decompressedIn, BUFFER_SIZE);
|
|
|
- return new InputStreamStreamInput(decompressedIn) {
|
|
|
- final AtomicBoolean closed = new AtomicBoolean(false);
|
|
|
-
|
|
|
+ final Releasable releasable;
|
|
|
+ final Inflater inflater;
|
|
|
+ if (threadLocal) {
|
|
|
+ final ReleasableReference<Inflater> current = inflaterForStreamRef.get();
|
|
|
+ if (current.inUse) {
|
|
|
+ // Nested de-compression streams should not happen but we still handle them safely by using a fresh Inflater
|
|
|
+ inflater = new Inflater(true);
|
|
|
+ releasable = inflater::end;
|
|
|
+ } else {
|
|
|
+ inflater = current.get();
|
|
|
+ releasable = current;
|
|
|
+ }
|
|
|
+ } else {
|
|
|
+ inflater = new Inflater(true);
|
|
|
+ releasable = inflater::end;
|
|
|
+ }
|
|
|
+ return new BufferedInputStream(new InflaterInputStream(in, inflater, BUFFER_SIZE) {
|
|
|
+ @Override
|
|
|
public void close() throws IOException {
|
|
|
try {
|
|
|
super.close();
|
|
|
} finally {
|
|
|
- if (closed.compareAndSet(false, true)) {
|
|
|
- // important to release native memory
|
|
|
- inflater.end();
|
|
|
- }
|
|
|
+ // We are ensured to only call this once since we wrap this stream in a BufferedInputStream that will only close
|
|
|
+ // its delegate once
|
|
|
+ releasable.close();
|
|
|
}
|
|
|
}
|
|
|
- };
|
|
|
+ }, BUFFER_SIZE);
|
|
|
}
|
|
|
|
|
|
@Override
|
|
|
- public StreamOutput streamOutput(OutputStream out) throws IOException {
|
|
|
+ public StreamOutput threadLocalStreamOutput(OutputStream out) throws IOException {
|
|
|
out.write(HEADER);
|
|
|
- final boolean nowrap = true;
|
|
|
- final Deflater deflater = new Deflater(LEVEL, nowrap);
|
|
|
+ final ReleasableReference<Deflater> current = deflaterForStreamRef.get();
|
|
|
+ final Releasable releasable;
|
|
|
+ final Deflater deflater;
|
|
|
+ if (current.inUse) {
|
|
|
+ // Nested compression streams should not happen but we still handle them safely by using a fresh Deflater
|
|
|
+ deflater = new Deflater(LEVEL, true);
|
|
|
+ releasable = deflater::end;
|
|
|
+ } else {
|
|
|
+ deflater = current.get();
|
|
|
+ releasable = current;
|
|
|
+ }
|
|
|
final boolean syncFlush = true;
|
|
|
- DeflaterOutputStream deflaterOutputStream = new DeflaterOutputStream(out, deflater, BUFFER_SIZE, syncFlush);
|
|
|
- OutputStream compressedOut = new BufferedOutputStream(deflaterOutputStream, BUFFER_SIZE);
|
|
|
- return new OutputStreamStreamOutput(compressedOut) {
|
|
|
- final AtomicBoolean closed = new AtomicBoolean(false);
|
|
|
-
|
|
|
+ DeflaterOutputStream deflaterOutputStream = new DeflaterOutputStream(out, deflater, BUFFER_SIZE, syncFlush) {
|
|
|
+ @Override
|
|
|
public void close() throws IOException {
|
|
|
try {
|
|
|
super.close();
|
|
|
} finally {
|
|
|
- if (closed.compareAndSet(false, true)) {
|
|
|
- // important to release native memory
|
|
|
- deflater.end();
|
|
|
- }
|
|
|
+ // We are ensured to only call this once since we wrap this stream in a BufferedOutputStream that will only close
|
|
|
+ // its delegate once below
|
|
|
+ releasable.close();
|
|
|
}
|
|
|
}
|
|
|
};
|
|
|
+ return new OutputStreamStreamOutput(new BufferedOutputStream(deflaterOutputStream, BUFFER_SIZE));
|
|
|
}
|
|
|
|
|
|
- // Reusable Inflater reference. Note: This is not used for the decompressing stream wrapper because we don't have strong guarantees
|
|
|
- // about the scope in which the stream wrapper is used.
|
|
|
- private static final ThreadLocal<Inflater> inflaterRef = ThreadLocal.withInitial(() -> new Inflater(true));
|
|
|
-
|
|
|
private static final ThreadLocal<BytesStreamOutput> baos = ThreadLocal.withInitial(BytesStreamOutput::new);
|
|
|
|
|
|
+ // Reusable Inflater reference. Note: This is a separate instance from the one used for the decompressing stream wrapper because we
|
|
|
+ // want to be able to deal with decompressing bytes references that were read from a decompressing stream.
|
|
|
+ private static final ThreadLocal<Inflater> inflaterRef = ThreadLocal.withInitial(() -> new Inflater(true));
|
|
|
+
|
|
|
@Override
|
|
|
public BytesReference uncompress(BytesReference bytesReference) throws IOException {
|
|
|
final BytesStreamOutput buffer = baos.get();
|
|
|
@@ -151,8 +235,8 @@ public class DeflateCompressor implements Compressor {
|
|
|
return res;
|
|
|
}
|
|
|
|
|
|
- // Reusable Deflater reference. Note: This is not used for the compressing stream wrapper because we don't have strong guarantees
|
|
|
- // about the scope in which the stream wrapper is used.
|
|
|
+ // Reusable Deflater reference. Note: This is a separate instance from the one used for the compressing stream wrapper because we
|
|
|
+ // want to be able to deal with compressing bytes references to a decompressing stream.
|
|
|
private static final ThreadLocal<Deflater> deflaterRef = ThreadLocal.withInitial(() -> new Deflater(LEVEL, true));
|
|
|
|
|
|
@Override
|
Armin Braun