diff --git a/core/src/main/java/org/elasticsearch/common/xcontent/XContentGenerator.java b/core/src/main/java/org/elasticsearch/common/xcontent/XContentGenerator.java index 8d1b8efef51..478f3a8a08f 100644 --- a/core/src/main/java/org/elasticsearch/common/xcontent/XContentGenerator.java +++ b/core/src/main/java/org/elasticsearch/common/xcontent/XContentGenerator.java @@ -94,4 +94,9 @@ public interface XContentGenerator extends Closeable, Flushable { void copyCurrentStructure(XContentParser parser) throws IOException; + /** + * Returns {@code true} if this XContentGenerator has been closed. A closed generator can not do any more output. + */ + boolean isClosed(); + } diff --git a/core/src/main/java/org/elasticsearch/common/xcontent/json/JsonXContentGenerator.java b/core/src/main/java/org/elasticsearch/common/xcontent/json/JsonXContentGenerator.java index 74e1cb5e58f..763fac4c6a3 100644 --- a/core/src/main/java/org/elasticsearch/common/xcontent/json/JsonXContentGenerator.java +++ b/core/src/main/java/org/elasticsearch/common/xcontent/json/JsonXContentGenerator.java @@ -419,4 +419,8 @@ public class JsonXContentGenerator implements XContentGenerator { generator.close(); } + @Override + public boolean isClosed() { + return generator.isClosed(); + } } diff --git a/core/src/main/java/org/elasticsearch/rest/action/RestBuilderListener.java b/core/src/main/java/org/elasticsearch/rest/action/RestBuilderListener.java index cc93e72d80d..c460331afaa 100644 --- a/core/src/main/java/org/elasticsearch/rest/action/RestBuilderListener.java +++ b/core/src/main/java/org/elasticsearch/rest/action/RestBuilderListener.java @@ -34,11 +34,22 @@ public abstract class RestBuilderListener extends RestResponseListener @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 assertBuilderClosed(builder); + 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; + + // pkg private method that we can override for testing + boolean assertBuilderClosed(XContentBuilder xContentBuilder) { + assert xContentBuilder.generator().isClosed() : "callers should ensure the XContentBuilder is closed themselves"; + return true; + } } diff --git a/core/src/test/java/org/elasticsearch/rest/action/RestBuilderListenerTests.java b/core/src/test/java/org/elasticsearch/rest/action/RestBuilderListenerTests.java new file mode 100644 index 00000000000..2bc0d0bdc81 --- /dev/null +++ b/core/src/test/java/org/elasticsearch/rest/action/RestBuilderListenerTests.java @@ -0,0 +1,90 @@ +/* + * 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 builderAtomicReference = new AtomicReference<>(); + RestBuilderListener builderListener = + new RestBuilderListener(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().generator().isClosed()); + } + + public void testXContentBuilderNotClosedInBuildResponseAssertionsDisabled() throws Exception { + AtomicReference builderAtomicReference = new AtomicReference<>(); + RestBuilderListener builderListener = + new RestBuilderListener(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); + } + + @Override + boolean assertBuilderClosed(XContentBuilder xContentBuilder) { + // don't check the actual builder being closed so we can test auto close + return true; + } + }; + + builderListener.buildResponse(Empty.INSTANCE); + assertNotNull(builderAtomicReference.get()); + assertTrue(builderAtomicReference.get().generator().isClosed()); + } + + public void testXContentBuilderNotClosedInBuildResponseAssertionsEnabled() throws Exception { + assumeTrue("tests are not being run with assertions", RestBuilderListener.class.desiredAssertionStatus()); + + RestBuilderListener builderListener = + new RestBuilderListener(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()); + } +}