Stop retaining transport messages after serialization (#122981)

We should not retain messages after serialization. We currently, have no
0-copy logic in place outside of some special cases that are handled
separately. As a result of not doing zero-copy logic, there is no need
to retain messages until they have been flushed to the wire in full.
This change reduces the lifetime of things like `SearchHit`
significantly, reducing the data-node side impact of fetching large
documents/large aggregations/top-hits to name a few.
The fact that this change technically passes released request + response
instances to the `messageListener` seems irrelevant since the listeners
are only used in tests anyway. If anything, we should look to refactor
this logic to avoid holding on to the request/response objects
needlessly.

Author: original-brownbear

server/src/main/java/org/elasticsearch/transport/OutboundHandler.java

@@ -12,13 +12,13 @@
 import org.apache.logging.log4j.Level;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.apache.lucene.store.AlreadyClosedException;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.common.bytes.ReleasableBytesReference;
 import org.elasticsearch.common.io.stream.RecyclerBytesStreamOutput;
 import org.elasticsearch.common.network.CloseableChannel;
 import org.elasticsearch.common.network.HandlingTimeTracker;
@@ -113,17 +113,12 @@ void sendRequest(
             isHandshake,
             compressionScheme
         );
-        if (request.tryIncRef() == false) {
-            assert false : "request [" + request + "] has been released already";
-            throw new AlreadyClosedException("request [" + request + "] has been released already");
-        }
-        sendMessage(channel, message, ResponseStatsConsumer.NONE, () -> {
-            try {
-                messageListener.onRequestSent(node, requestId, action, request, options);
-            } finally {
-                request.decRef();
-            }
-        });
+        sendMessage(
+            channel,
+            message,
+            ResponseStatsConsumer.NONE,
+            () -> messageListener.onRequestSent(node, requestId, action, request, options)
+        );
     }
 
     /**
@@ -151,15 +146,9 @@ void sendResponse(
             isHandshake,
             compressionScheme
         );
-        response.mustIncRef();
+        assert response.hasReferences();
         try {
-            sendMessage(channel, message, responseStatsConsumer, () -> {
-                try {
-                    messageListener.onResponseSent(requestId, action, response);
-                } finally {
-                    response.decRef();
-                }
-            });
+            sendMessage(channel, message, responseStatsConsumer, () -> messageListener.onResponseSent(requestId, action, response));
         } catch (Exception ex) {
             if (isHandshake) {
                 logger.error(
@@ -222,7 +211,6 @@ private void sendMessage(
                 Releasables.closeExpectNoException(onAfter);
             }
         }
-        final Releasable release = Releasables.wrap(byteStreamOutput, onAfter);
         final BytesReference message;
         boolean serializeSuccess = false;
         try {
@@ -233,11 +221,20 @@ private void sendMessage(
             throw e;
         } finally {
             if (serializeSuccess == false) {
-                release.close();
+                Releasables.close(byteStreamOutput, onAfter);
             }
         }
         responseStatsConsumer.addResponseStats(message.length());
-        internalSend(channel, message, networkMessage, ActionListener.running(release::close));
+        internalSend(
+            channel,
+            message,
+            networkMessage,
+            ActionListener.releasing(
+                message instanceof ReleasableBytesReference r
+                    ? Releasables.wrap(byteStreamOutput, onAfter, r)
+                    : Releasables.wrap(byteStreamOutput, onAfter)
+            )
+        );
     }
 
     private void internalSend(

server/src/main/java/org/elasticsearch/transport/OutboundMessage.java

@@ -11,15 +11,16 @@
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.bytes.CompositeBytesReference;
+import org.elasticsearch.common.bytes.ReleasableBytesReference;
 import org.elasticsearch.common.compress.CompressorFactory;
 import org.elasticsearch.common.io.stream.OutputStreamStreamOutput;
 import org.elasticsearch.common.io.stream.RecyclerBytesStreamOutput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
+import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.Streams;
 
 import java.io.IOException;
@@ -56,7 +57,7 @@ BytesReference serialize(RecyclerBytesStreamOutput bytesStream) throws IOExcepti
 
         final boolean compress = TransportStatus.isCompress(status);
         final StreamOutput stream = compress ? wrapCompressed(bytesStream) : bytesStream;
-        final BytesReference zeroCopyBuffer;
+        final ReleasableBytesReference zeroCopyBuffer;
         try {
             stream.setTransportVersion(version);
             if (variableHeaderLength == -1) {
@@ -67,10 +68,10 @@ BytesReference serialize(RecyclerBytesStreamOutput bytesStream) throws IOExcepti
                 zeroCopyBuffer = bRequest.bytes;
             } else if (message instanceof RemoteTransportException) {
                 stream.writeException((RemoteTransportException) message);
-                zeroCopyBuffer = BytesArray.EMPTY;
+                zeroCopyBuffer = ReleasableBytesReference.empty();
             } else {
                 message.writeTo(stream);
-                zeroCopyBuffer = BytesArray.EMPTY;
+                zeroCopyBuffer = ReleasableBytesReference.empty();
             }
         } finally {
             // We have to close here before accessing the bytes when using compression to ensure that some marker bytes (EOS marker)
@@ -83,7 +84,8 @@ BytesReference serialize(RecyclerBytesStreamOutput bytesStream) throws IOExcepti
         if (zeroCopyBuffer.length() == 0) {
             reference = message;
         } else {
-            reference = CompositeBytesReference.of(message, zeroCopyBuffer);
+            zeroCopyBuffer.mustIncRef();
+            reference = new ReleasableBytesReference(CompositeBytesReference.of(message, zeroCopyBuffer), (RefCounted) zeroCopyBuffer);
         }
 
         bytesStream.seek(0);