Wrap metering parser decorator around dot expanding parser if necessary (elastic#138660)

Author: mosche

server/src/main/java/org/elasticsearch/index/mapper/DocumentParser.java

@@ -28,7 +28,7 @@
 import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.indices.breaker.NoneCircuitBreakerService;
-import org.elasticsearch.plugins.internal.XContentMeteringParserDecorator;
+import org.elasticsearch.plugins.internal.XContentParserDecorator;
 import org.elasticsearch.search.lookup.LeafFieldLookupProvider;
 import org.elasticsearch.search.lookup.SearchLookup;
 import org.elasticsearch.search.lookup.Source;
@@ -37,7 +37,6 @@
 import org.elasticsearch.xcontent.XContentParseException;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xcontent.XContentParserConfiguration;
-import org.elasticsearch.xcontent.XContentType;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -85,19 +84,22 @@ public ParsedDocument parseDocument(SourceToParse source, MappingLookup mappingL
             throw new DocumentParsingException(new XContentLocation(0, 0), "failed to parse, document is empty");
         }
         final RootDocumentParserContext context;
-        final XContentType xContentType = source.getXContentType();
 
-        XContentMeteringParserDecorator meteringParserDecorator = source.getMeteringParserDecorator();
         try (
-            XContentParser parser = meteringParserDecorator.decorate(
+            // the context needs to be closed after parsing to make sure the actual parser is properly closed;
+            // closing won't impact other state of the context
+            RootDocumentParserContext ctx = new RootDocumentParserContext(
+                mappingLookup,
+                mappingParserContext,
+                source,
                 XContentHelper.createParser(
                     parserConfiguration.withIncludeSourceOnError(source.getIncludeSourceOnError()),
                     source.source(),
-                    xContentType
+                    source.getXContentType()
                 )
             )
         ) {
-            context = new RootDocumentParserContext(mappingLookup, mappingParserContext, source, parser);
+            context = ctx;
             validateStart(context.parser());
             MetadataFieldMapper[] metadataFieldsMappers = mappingLookup.getMapping().getSortedMetadataMappers();
             internalParseDocument(metadataFieldsMappers, context);
@@ -121,7 +123,7 @@ public ParsedDocument parseDocument(SourceToParse source, MappingLookup mappingL
             context.sourceToParse().source(),
             context.sourceToParse().getXContentType(),
             dynamicUpdate,
-            meteringParserDecorator.meteredDocumentSize()
+            source.getMeteringParserDecorator().meteredDocumentSize()
         ) {
             @Override
             public String documentDescription() {
@@ -1085,7 +1087,7 @@ public ObjectMapper merge(Mapper mergeWith, MapperMergeContext mapperMergeContex
      * Internal version of {@link DocumentParserContext} that is aware of implementation details like nested documents
      * and how they are stored in the lucene index.
      */
-    private static class RootDocumentParserContext extends DocumentParserContext {
+    private static class RootDocumentParserContext extends DocumentParserContext implements AutoCloseable {
         private final ContentPath path = new ContentPath();
         private final XContentParser parser;
         private final LuceneDocument document;
@@ -1120,17 +1122,23 @@ private static class RootDocumentParserContext extends DocumentParserContext {
             this.tsid = tsid;
             assert this.tsid == null || indexSettings.getMode() == IndexMode.TIME_SERIES
                 : "tsid should only be set for time series indices";
+            XContentParserDecorator parserDecorator = source.getMeteringParserDecorator();
             if (mappingLookup.getMapping().getRoot().subobjects() == ObjectMapper.Subobjects.ENABLED) {
-                this.parser = DotExpandingXContentParser.expandDots(parser, this.path);
+                this.parser = parserDecorator.decorate(DotExpandingXContentParser.expandDots(parser, this.path));
             } else {
-                this.parser = parser;
+                this.parser = parserDecorator.decorate(parser);
             }
             this.document = new LuceneDocument();
             this.documents.add(document);
             this.maxAllowedNumNestedDocs = indexSettings().getMappingNestedDocsLimit();
             this.numNestedDocs = 0L;
         }
 
+        @Override
+        public void close() throws IOException {
+            parser.close();
+        }
+
         @Override
         public Mapper getMapper(String name) {
             Mapper mapper = getMetadataMapper(name);

server/src/main/java/org/elasticsearch/index/mapper/DotExpandingXContentParser.java

@@ -389,6 +389,14 @@ public XContentString optimizedTextOrNull() throws IOException {
         return super.optimizedTextOrNull();
     }
 
+    @Override
+    public XContentString optimizedText() throws IOException {
+        if (state == State.EXPANDING_START_OBJECT) {
+            throw new IllegalStateException("Can't get text on a " + currentToken() + " at " + getTokenLocation());
+        }
+        return super.optimizedText();
+    }
+
     @Override
     public String textOrNull() throws IOException {
         if (state == State.EXPANDING_START_OBJECT) {
@@ -397,6 +405,14 @@ public String textOrNull() throws IOException {
         return super.textOrNull();
     }
 
+    @Override
+    public String text() throws IOException {
+        if (state == State.EXPANDING_START_OBJECT) {
+            throw new IllegalStateException("Can't get text on a " + currentToken() + " at " + getTokenLocation());
+        }
+        return super.text();
+    }
+
     @Override
     public Number numberValue() throws IOException {
         if (state == State.EXPANDING_START_OBJECT) {