ensure the XContentBuilder is always closed in RestBuilderListener

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

@@ -998,6 +998,10 @@ public void close() {
         }
     }
 
+    public boolean isClosed() {
+        return generator.isClosed();
+    }
+
     public XContentGenerator generator() {
         return this.generator;
     }

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

@@ -94,4 +94,6 @@ public interface XContentGenerator extends Closeable, Flushable {
 
     void copyCurrentStructure(XContentParser parser) throws IOException;
 
+    boolean isClosed();
+
 }

core/src/main/java/org/elasticsearch/common/xcontent/json/JsonXContentGenerator.java

@@ -419,4 +419,8 @@ public void close() throws IOException {
         generator.close();
     }
 
+    @Override
+    public boolean isClosed() {
+        return generator.isClosed();
+    }
 }

core/src/main/java/org/elasticsearch/rest/action/RestBuilderListener.java

@@ -28,17 +28,25 @@
  */
 public abstract class RestBuilderListener<Response> extends RestResponseListener<Response> {
 
+    // pkg-private static boolean that enables testing of the auto-close functionality
+    static boolean assertionsEnabled = true;
+
     public RestBuilderListener(RestChannel channel) {
         super(channel);
     }
 
     @Override
     public final RestResponse buildResponse(Response response) throws Exception {
-        return buildResponse(response, channel.newBuilder());
+        try (XContentBuilder builder = channel.newBuilder()) {
+            final RestResponse restResponse = buildResponse(response, builder);
+            assert assertionsEnabled == false || builder.isClosed() : "callers should ensure the XContentBuilder is closed themselves";
+            return restResponse;
+        }
     }
 
     /**
-     * Builds a response to send back over the channel.
+     * Builds a response to send back over the channel. Implementors should ensure that they close the provided {@link XContentBuilder}
+     * using the {@link XContentBuilder#close()} method.
      */
     public abstract RestResponse buildResponse(Response response, XContentBuilder builder) throws Exception;
 }

core/src/test/java/org/elasticsearch/rest/action/RestBuilderListenerTests.java

@@ -0,0 +1,89 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.rest.action;
+
+import org.elasticsearch.common.bytes.BytesArray;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.rest.BytesRestResponse;
+import org.elasticsearch.rest.RestResponse;
+import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.test.rest.FakeRestChannel;
+import org.elasticsearch.test.rest.FakeRestRequest;
+import org.elasticsearch.transport.TransportResponse;
+import org.elasticsearch.transport.TransportResponse.Empty;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+public class RestBuilderListenerTests extends ESTestCase {
+
+    public void testXContentBuilderClosedInBuildResponse() throws Exception {
+        AtomicReference<XContentBuilder> builderAtomicReference = new AtomicReference<>();
+        RestBuilderListener<TransportResponse.Empty> builderListener =
+            new RestBuilderListener<Empty>(new FakeRestChannel(new FakeRestRequest(), randomBoolean(), 1)) {
+                @Override
+                public RestResponse buildResponse(Empty empty, XContentBuilder builder) throws Exception {
+                    builderAtomicReference.set(builder);
+                    builder.close();
+                    return new BytesRestResponse(RestStatus.OK, BytesRestResponse.TEXT_CONTENT_TYPE, BytesArray.EMPTY);
+                }
+        };
+
+        builderListener.buildResponse(Empty.INSTANCE);
+        assertNotNull(builderAtomicReference.get());
+        assertTrue(builderAtomicReference.get().isClosed());
+    }
+
+    public void testXContentBuilderNotClosedInBuildResponseAssertionsDisabled() throws Exception {
+        AtomicReference<XContentBuilder> builderAtomicReference = new AtomicReference<>();
+        RestBuilderListener<TransportResponse.Empty> builderListener =
+            new RestBuilderListener<Empty>(new FakeRestChannel(new FakeRestRequest(), randomBoolean(), 1)) {
+                @Override
+                public RestResponse buildResponse(Empty empty, XContentBuilder builder) throws Exception {
+                    builderAtomicReference.set(builder);
+                    return new BytesRestResponse(RestStatus.OK, BytesRestResponse.TEXT_CONTENT_TYPE, BytesArray.EMPTY);
+                }
+        };
+
+        RestBuilderListener.assertionsEnabled = false;
+        try {
+            builderListener.buildResponse(Empty.INSTANCE);
+            assertNotNull(builderAtomicReference.get());
+            assertTrue(builderAtomicReference.get().isClosed());
+        } finally {
+            RestBuilderListener.assertionsEnabled = true;
+        }
+    }
+
+    public void testXContentBuilderNotClosedInBuildResponseAssertionsEnabled() throws Exception {
+        assumeTrue("tests are not being run with assertions", RestBuilderListener.class.desiredAssertionStatus());
+
+        RestBuilderListener<TransportResponse.Empty> builderListener =
+            new RestBuilderListener<Empty>(new FakeRestChannel(new FakeRestRequest(), randomBoolean(), 1)) {
+                @Override
+                public RestResponse buildResponse(Empty empty, XContentBuilder builder) throws Exception {
+                    return new BytesRestResponse(RestStatus.OK, BytesRestResponse.TEXT_CONTENT_TYPE, BytesArray.EMPTY);
+                }
+        };
+
+        AssertionError error = expectThrows(AssertionError.class, () -> builderListener.buildResponse(Empty.INSTANCE));
+        assertEquals("callers should ensure the XContentBuilder is closed themselves", error.getMessage());
+    }
+}