Repository: elastic/elasticsearch
Pull request #126988: Make DelayableWriteable compress its contents
Author: login: original-brownbear
Labels: name: >non-issue
name: :Analytics/Aggregations
description: Aggregations
name: backport pending
name: Team:Analytics
description: Meta label for analytical engine team (ESQL/Aggs/Geo)
name: v8.19.0
name: v9.1.0
Created: 2025-04-17T11:49:35Z | Merged: 2025-04-19T17:53:46Z

In light of data from recent escalations and the introduction of batched execution we can make two improvements to this logic.
For one, we should prefix with a fixed length length field so that we don't need to do any copying when serializing to account for the vint. This outright halves the memory bandwidth required relative to the previous implementation.
More importantly maybe, we should compress these bytes. The wire-format for aggregations is rather inefficient when working with nested bucket aggregations since the type strings are repeated over and over. These don't contribute to the peak heap requirements because they are translated into Java types, but blow up the message size considerably (among other things). Practically, it seems that we often get compression ratios of ~10x for aggregations.
Given that we generally have more memory issues than CPU issues during the reduce-step it seems like an easy tradeoff to trade a little CPU for compression for serious heap savings here.


Diff:
diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java
index 23a19d05866a5..c185029ed58f8 100644
--- a/server/src/main/java/org/elasticsearch/TransportVersions.java
+++ b/server/src/main/java/org/elasticsearch/TransportVersions.java
@@ -230,6 +230,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ESQL_DOCUMENTS_FOUND_AND_VALUES_LOADED = def(9_056_0_00);
     public static final TransportVersion BATCHED_QUERY_EXECUTION_DELAYABLE_WRITABLE = def(9_057_0_00);
     public static final TransportVersion SEARCH_INCREMENTAL_TOP_DOCS_NULL = def(9_058_0_00);
+    public static final TransportVersion COMPRESS_DELAYABLE_WRITEABLE = def(9_059_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,
diff --git a/server/src/main/java/org/elasticsearch/common/bytes/ReleasableBytesReference.java b/server/src/main/java/org/elasticsearch/common/bytes/ReleasableBytesReference.java
index 05c2b7cc06e89..4343892428c9a 100644
--- a/server/src/main/java/org/elasticsearch/common/bytes/ReleasableBytesReference.java
+++ b/server/src/main/java/org/elasticsearch/common/bytes/ReleasableBytesReference.java
@@ -178,6 +178,11 @@ public ReleasableBytesReference readReleasableBytesReference() throws IOExceptio
                 return retainAndSkip(len);
             }
 
+            @Override
+            public ReleasableBytesReference readReleasableBytesReference(int len) throws IOException {
+                return retainAndSkip(len);
+            }
+
             @Override
             public ReleasableBytesReference readAllToReleasableBytesReference() throws IOException {
                 return retainAndSkip(length() - offset());
diff --git a/server/src/main/java/org/elasticsearch/common/io/Streams.java b/server/src/main/java/org/elasticsearch/common/io/Streams.java
index d2b9580efe029..60e6d36c2fc0a 100644
--- a/server/src/main/java/org/elasticsearch/common/io/Streams.java
+++ b/server/src/main/java/org/elasticsearch/common/io/Streams.java
@@ -235,6 +235,11 @@ public void close() throws IOException {
         public BytesReference bytes() {
             return delegate.bytes();
         }
+
+        @Override
+        public void seek(long position) {
+            delegate.seek(position);
+        }
     }
 
     /**
diff --git a/server/src/main/java/org/elasticsearch/common/io/stream/BytesStream.java b/server/src/main/java/org/elasticsearch/common/io/stream/BytesStream.java
index 86a0eb5f329f3..c1a03ab9f9f2f 100644
--- a/server/src/main/java/org/elasticsearch/common/io/stream/BytesStream.java
+++ b/server/src/main/java/org/elasticsearch/common/io/stream/BytesStream.java
@@ -10,8 +10,28 @@
 package org.elasticsearch.common.io.stream;
 
 import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.common.compress.CompressorFactory;
+import org.elasticsearch.core.Streams;
+
+import java.io.IOException;
 
 public abstract class BytesStream extends StreamOutput {
 
+    @Override
+    public void writeWithSizePrefix(Writeable writeable) throws IOException {
+        long pos = position();
+        seek(pos + Integer.BYTES);
+        try (var out = new OutputStreamStreamOutput(CompressorFactory.COMPRESSOR.threadLocalOutputStream(Streams.noCloseStream(this)))) {
+            out.setTransportVersion(getTransportVersion());
+            writeable.writeTo(out);
+        }
+        long newPos = position();
+        seek(pos);
+        writeInt(Math.toIntExact(newPos - pos - Integer.BYTES));
+        seek(newPos);
+    }
+
     public abstract BytesReference bytes();
+
+    public abstract void seek(long position);
 }
diff --git a/server/src/main/java/org/elasticsearch/common/io/stream/BytesStreamOutput.java b/server/src/main/java/org/elasticsearch/common/io/stream/BytesStreamOutput.java
index 317a33f4829c9..18394a26a4192 100644
--- a/server/src/main/java/org/elasticsearch/common/io/stream/BytesStreamOutput.java
+++ b/server/src/main/java/org/elasticsearch/common/io/stream/BytesStreamOutput.java
@@ -106,6 +106,7 @@ public void flush() {
         // nothing to do
     }
 
+    @Override
     public void seek(long position) {
         ensureCapacity(position);
         count = (int) position;
diff --git a/server/src/main/java/org/elasticsearch/common/io/stream/DelayableWriteable.java b/server/src/main/java/org/elasticsearch/common/io/stream/DelayableWriteable.java
index e1161f03cb8a7..a2b6be6ab0363 100644
--- a/server/src/main/java/org/elasticsearch/common/io/stream/DelayableWriteable.java
+++ b/server/src/main/java/org/elasticsearch/common/io/stream/DelayableWriteable.java
@@ -10,8 +10,10 @@
 package org.elasticsearch.common.io.stream;
 
 import org.elasticsearch.TransportVersion;
+import org.elasticsearch.TransportVersions;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.bytes.ReleasableBytesReference;
+import org.elasticsearch.common.compress.CompressorFactory;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
 
@@ -53,11 +55,22 @@ public static <T extends Writeable> DelayableWriteable<T> referencing(T referenc
      * when {@link #expand()} is called.
      */
     public static <T extends Writeable> DelayableWriteable<T> delayed(Writeable.Reader<T> reader, StreamInput in) throws IOException {
-        return new Serialized<>(reader, in.getTransportVersion(), in.namedWriteableRegistry(), in.readReleasableBytesReference());
+        return new Serialized<>(
+            reader,
+            in.getTransportVersion(),
+            in.namedWriteableRegistry(),
+            in.getTransportVersion().onOrAfter(TransportVersions.COMPRESS_DELAYABLE_WRITEABLE)
+                ? in.readReleasableBytesReference(in.readInt())
+                : in.readReleasableBytesReference()
+        );
     }
 
     public static <T extends Writeable> DelayableWriteable<T> referencing(Writeable.Reader<T> reader, StreamInput in) throws IOException {
-        try (ReleasableBytesReference serialized = in.readReleasableBytesReference()) {
+        try (
+            ReleasableBytesReference serialized = in.getTransportVersion().onOrAfter(TransportVersions.COMPRESS_DELAYABLE_WRITEABLE)
+                ? in.readReleasableBytesReference(in.readInt())
+                : in.readReleasableBytesReference()
+        ) {
             return new Referencing<>(deserialize(reader, in.getTransportVersion(), in.namedWriteableRegistry(), serialized));
         }
     }
@@ -95,7 +108,11 @@ private Referencing(T reference) {
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
-            out.writeWithSizePrefix(reference);
+            if (out.getTransportVersion().onOrAfter(TransportVersions.COMPRESS_DELAYABLE_WRITEABLE)) {
+                out.writeWithSizePrefix(reference);
+            } else {
+                out.legacyWriteWithSizePrefix(reference);
+            }
         }
 
         @Override
@@ -105,13 +122,14 @@ public T expand() {
 
         @Override
         public Serialized<T> asSerialized(Reader<T> reader, NamedWriteableRegistry registry) {
-            BytesStreamOutput buffer;
-            try {
-                buffer = writeToBuffer(TransportVersion.current());
+            // TODO: this path is currently not used in production code, if it ever is this should start using pooled buffers
+            BytesStreamOutput buffer = new BytesStreamOutput();
+            try (var out = new OutputStreamStreamOutput(CompressorFactory.COMPRESSOR.threadLocalOutputStream(buffer))) {
+                out.setTransportVersion(TransportVersion.current());
+                reference.writeTo(out);
             } catch (IOException e) {
                 throw new RuntimeException("unexpected error writing writeable to buffer", e);
             }
-            // TODO: this path is currently not used in production code, if it ever is this should start using pooled buffers
             return new Serialized<>(reader, TransportVersion.current(), registry, ReleasableBytesReference.wrap(buffer.bytes()));
         }
 
@@ -125,14 +143,6 @@ public long getSerializedSize() {
             return DelayableWriteable.getSerializedSize(reference);
         }
 
-        private BytesStreamOutput writeToBuffer(TransportVersion version) throws IOException {
-            try (BytesStreamOutput buffer = new BytesStreamOutput()) {
-                buffer.setTransportVersion(version);
-                reference.writeTo(buffer);
-                return buffer;
-            }
-        }
-
         @Override
         public void close() {
             // noop
@@ -169,7 +179,12 @@ public void writeTo(StreamOutput out) throws IOException {
                  * which is good because this is how shard request caching
                  * works.
                  */
-                out.writeBytesReference(serialized);
+                if (out.getTransportVersion().onOrAfter(TransportVersions.COMPRESS_DELAYABLE_WRITEABLE)) {
+                    out.writeInt(serialized.length());
+                    serialized.writeTo(out);
+                } else {
+                    out.writeBytesReference(serialized);
+                }
             } else {
                 /*
                  * If the version doesn't line up then we have to deserialize
@@ -211,6 +226,7 @@ public long getSerializedSize() {
         public void close() {
             serialized.close();
         }
+
     }
 
     /**
@@ -232,7 +248,11 @@ private static <T> T deserialize(
         NamedWriteableRegistry registry,
         BytesReference serialized
     ) throws IOException {
-        try (StreamInput in = serialized.streamInput()) {
+        try (
+            StreamInput in = serializedAtVersion.onOrAfter(TransportVersions.COMPRESS_DELAYABLE_WRITEABLE)
+                ? CompressorFactory.COMPRESSOR.threadLocalStreamInput(serialized.streamInput())
+                : serialized.streamInput()
+        ) {
             return reader.read(wrapWithDeduplicatorStreamInput(in, serializedAtVersion, registry));
         }
     }
diff --git a/server/src/main/java/org/elasticsearch/common/io/stream/FilterStreamInput.java b/server/src/main/java/org/elasticsearch/common/io/stream/FilterStreamInput.java
index 6e9fae6e17bb8..34de4faa69f62 100644
--- a/server/src/main/java/org/elasticsearch/common/io/stream/FilterStreamInput.java
+++ b/server/src/main/java/org/elasticsearch/common/io/stream/FilterStreamInput.java
@@ -47,6 +47,11 @@ public ReleasableBytesReference readReleasableBytesReference() throws IOExceptio
         return delegate.readReleasableBytesReference();
     }
 
+    @Override
+    public ReleasableBytesReference readReleasableBytesReference(int length) throws IOException {
+        return delegate.readReleasableBytesReference(length);
+    }
+
     @Override
     public boolean supportReadAllToReleasableBytesReference() {
         return delegate.supportReadAllToReleasableBytesReference();
diff --git a/server/src/main/java/org/elasticsearch/common/io/stream/RecyclerBytesStreamOutput.java b/server/src/main/java/org/elasticsearch/common/io/stream/RecyclerBytesStreamOutput.java
index 0aacc466cce7a..ddf97a5151690 100644
--- a/server/src/main/java/org/elasticsearch/common/io/stream/RecyclerBytesStreamOutput.java
+++ b/server/src/main/java/org/elasticsearch/common/io/stream/RecyclerBytesStreamOutput.java
@@ -155,7 +155,7 @@ public void writeLongLE(long i) throws IOException {
     }
 
     @Override
-    public void writeWithSizePrefix(Writeable writeable) throws IOException {
+    public void legacyWriteWithSizePrefix(Writeable writeable) throws IOException {
         // TODO: do this without copying the bytes from tmp by calling writeBytes and just use the pages in tmp directly through
         // manipulation of the offsets on the pages after writing to tmp. This will require adjustments to the places in this class
         // that make assumptions about the page size
@@ -214,6 +214,7 @@ public void flush() {
         // nothing to do
     }
 
+    @Override
     public void seek(long position) {
         ensureCapacityFromPosition(position);
         int offsetInPage = (int) (position % pageSize);
diff --git a/server/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java b/server/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java
index e07861ba05433..8b54c7a78907c 100644
--- a/server/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java
+++ b/server/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java
@@ -127,6 +127,14 @@ public ReleasableBytesReference readReleasableBytesReference() throws IOExceptio
         return ReleasableBytesReference.wrap(readBytesReference());
     }
 
+    /**
+     * Same as {@link #readBytesReference()} but with an explicitly provided length.
+     * @param length number of bytes to read
+     */
+    public ReleasableBytesReference readReleasableBytesReference(int length) throws IOException {
+        return ReleasableBytesReference.wrap(readBytesReference(length));
+    }
+
     /**
      * Reads the same bytes returned by {@link #readReleasableBytesReference()} but does not retain a reference to these bytes.
      * The returned {@link BytesReference} thus only contains valid content as long as the underlying buffer has not been released.
@@ -138,6 +146,10 @@ public BytesReference readSlicedBytesReference() throws IOException {
         return readBytesReference();
     }
 
+    public BytesReference readSlicedBytesReference(int bytes) throws IOException {
+        return readBytesReference(bytes);
+    }
+
     /**
      * Checks if this {@link InputStream} supports {@link #readAllToReleasableBytesReference()}.
      */
diff --git a/server/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java b/server/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java
index 4573c5a1fea57..0621709683ba3 100644
--- a/server/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java
+++ b/server/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java
@@ -18,6 +18,7 @@
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.common.compress.CompressorFactory;
 import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.common.io.stream.Writeable.Writer;
 import org.elasticsearch.common.settings.SecureString;
@@ -131,14 +132,34 @@ public void writeByteArray(byte[] b) throws IOException {
      * Serializes a writable just like {@link Writeable#writeTo(StreamOutput)} would but prefixes it with the serialized size of the result.
      *
      * @param writeable {@link Writeable} to serialize
+     * @deprecated use {@link #writeWithSizePrefix} instead
      */
-    public void writeWithSizePrefix(Writeable writeable) throws IOException {
+    @Deprecated
+    public void legacyWriteWithSizePrefix(Writeable writeable) throws IOException {
         final BytesStreamOutput tmp = new BytesStreamOutput();
         tmp.setTransportVersion(version);
         writeable.writeTo(tmp);
         writeBytesReference(tmp.bytes());
     }
 
+    /**
+     * Serializes a writable just like {@link Writeable#writeTo(StreamOutput)} would but also compresses and prefixes it with the serialized
+     * size of the result.
+
+     *
+     * @param writeable {@link Writeable} to serialize
+     */
+    public void writeWithSizePrefix(Writeable writeable) throws IOException {
+        final BytesStreamOutput tmp = new BytesStreamOutput();
+        try (var o = new OutputStreamStreamOutput(CompressorFactory.COMPRESSOR.threadLocalOutputStream(tmp))) {
+            o.setTransportVersion(version);
+            writeable.writeTo(o);
+        }
+        var bytes = tmp.bytes();
+        this.writeInt(bytes.length());
+        bytes.writeTo(this);
+    }
+
     /**
      * Writes the bytes reference, including a length header.
      */
diff --git a/server/src/test/java/org/elasticsearch/common/io/stream/BytesStreamsTests.java b/server/src/test/java/org/elasticsearch/common/io/stream/BytesStreamsTests.java
index f51cc7e1de6b8..811f9b692399a 100644
--- a/server/src/test/java/org/elasticsearch/common/io/stream/BytesStreamsTests.java
+++ b/server/src/test/java/org/elasticsearch/common/io/stream/BytesStreamsTests.java
@@ -1001,6 +1001,11 @@ public BytesReference bytes() {
             return bytesReference;
         }
 
+        @Override
+        public void seek(long position) {
+            output.seek(position);
+        }
+
         public int size() {
             int size = output.size();
             assertThat((long) size, equalTo(counting.size()));
diff --git a/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java
index fa77aebb80b53..79723d074c45d 100644
--- a/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java
+++ b/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java
@@ -67,7 +67,6 @@
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.bytes.CompositeBytesReference;
 import org.elasticsearch.common.io.stream.BytesStreamOutput;
-import org.elasticsearch.common.io.stream.DelayableWriteable;
 import org.elasticsearch.common.io.stream.NamedWriteable;
 import org.elasticsearch.common.io.stream.NamedWriteableAwareStreamInput;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
@@ -1971,11 +1970,10 @@ protected static <T extends Writeable> T copyInstance(
             } else {
                 BytesReference bytesReference = output.copyBytes();
                 output.reset();
-                output.writeBytesReference(bytesReference);
+                bytesReference.writeTo(output);
                 try (StreamInput in = new NamedWriteableAwareStreamInput(output.bytes().streamInput(), namedWriteableRegistry)) {
                     in.setTransportVersion(version);
-                    DelayableWriteable<T> delayableWriteable = DelayableWriteable.delayed(reader, in);
-                    return delayableWriteable.expand();
+                    return reader.read(in);
                 }
             }
         }


Reviews:
user: login: rjernst
body: The concepts here seem good (using fixed-size sizes and compressing). I have one question.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-19T14:03:38Z
user: login: original-brownbear
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-19T15:02:41Z
user: login: rjernst
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-19T16:57:28Z
user: login: rjernst
body: LGTM
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-19T16:57:39Z

Review comments:
diff_hunk: @@ -169,7 +179,12 @@ public void writeTo(StreamOutput out) throws IOException {
                  * which is good because this is how shard request caching
                  * works.
                  */
-                out.writeBytesReference(serialized);
+                if (out.getTransportVersion().onOrAfter(TransportVersions.COMPRESS_DELAYABLE_WRITEABLE)) {
+                    out.writeInt(serialized.length());
+                    serialized.writeTo(out);
+                } else {
+                    out.writeBytesReference(serialized);
path: server/src/main/java/org/elasticsearch/common/io/stream/DelayableWriteable.java
user: login: rjernst
body: It seems the above comment is no longer accurate with the new branch since we will no longer copy the raw bytes, but instead reencode?
created_at: 2025-04-19T13:57:13Z
author_association: MEMBER
diff_hunk: @@ -169,7 +179,12 @@ public void writeTo(StreamOutput out) throws IOException {
                  * which is good because this is how shard request caching
                  * works.
                  */
-                out.writeBytesReference(serialized);
+                if (out.getTransportVersion().onOrAfter(TransportVersions.COMPRESS_DELAYABLE_WRITEABLE)) {
+                    out.writeInt(serialized.length());
+                    serialized.writeTo(out);
+                } else {
+                    out.writeBytesReference(serialized);
path: server/src/main/java/org/elasticsearch/common/io/stream/DelayableWriteable.java
user: login: original-brownbear
body: We still copy the raw bytes, it's just that the prefix is a full int now.
`writeBytesReference` is nothing else but:

```java
                 out.writeVInt(serialized.length());
                 serialized.writeTo(out);
```

all that happens here is going from a vint to an int?
created_at: 2025-04-19T15:02:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -169,7 +179,12 @@ public void writeTo(StreamOutput out) throws IOException {
                  * which is good because this is how shard request caching
                  * works.
                  */
-                out.writeBytesReference(serialized);
+                if (out.getTransportVersion().onOrAfter(TransportVersions.COMPRESS_DELAYABLE_WRITEABLE)) {
+                    out.writeInt(serialized.length());
+                    serialized.writeTo(out);
+                } else {
+                    out.writeBytesReference(serialized);
path: server/src/main/java/org/elasticsearch/common/io/stream/DelayableWriteable.java
user: login: rjernst
body: Ah ok, I understand now.
created_at: 2025-04-19T16:57:28Z
author_association: MEMBER

PR comments:
user: login: elasticsearchmachine
created_at: 2025-04-17T11:49:58Z
body: Pinging @elastic/es-analytical-engine (Team:Analytics)
author_association: COLLABORATOR
user: login: original-brownbear
created_at: 2025-04-19T17:53:33Z
body: Thanks Ryan!
author_association: CONTRIBUTOR