Add ability to limit fields added during Mapper#merge

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

@@ -449,19 +449,19 @@ public Builder add(FieldMapper.Builder builder) {
                 return this;
             }
 
-            public Builder add(FieldMapper mapper) {
+            private void add(FieldMapper mapper) {
                 mapperBuilders.put(mapper.simpleName(), context -> mapper);
-                return this;
             }
 
-            public Builder update(FieldMapper toMerge, MapperMergeContext context) {
+            private void update(FieldMapper toMerge, MapperMergeContext context) {
                 if (mapperBuilders.containsKey(toMerge.simpleName()) == false) {
-                    add(toMerge);
+                    if (context.decrementFieldBudgetIfPossible(toMerge.mapperSize())) {
+                        add(toMerge);
+                    }
                 } else {
                     FieldMapper existing = mapperBuilders.get(toMerge.simpleName()).apply(context.getMapperBuilderContext());
                     add(existing.merge(toMerge, context));
                 }
-                return this;
             }
 
             public boolean hasMultiFields() {

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

@@ -75,7 +75,7 @@ public final String simpleName() {
 
     /**
      * Return the merge of {@code mergeWith} into this.
-     *  Both {@code this} and {@code mergeWith} will be left unmodified.
+     * Both {@code this} and {@code mergeWith} will be left unmodified.
      */
     public abstract Mapper merge(Mapper mergeWith, MapperMergeContext mapperMergeContext);
 
@@ -135,4 +135,18 @@ public static FieldType freezeAndDeduplicateFieldType(FieldType fieldType) {
         }
         return fieldTypeDeduplicator.computeIfAbsent(fieldType, Function.identity());
     }
+
+    /**
+     * Returns the size this mapper counts against the {@linkplain MapperService#INDEX_MAPPING_TOTAL_FIELDS_LIMIT_SETTING field limit}.
+     * <p>
+     * Needs to be in sync with {@link MappingLookup#getTotalFieldsCount()}.
+     */
+    public int mapperSize() {
+        int size = 1;
+        for (Mapper mapper : this) {
+            size += mapper.mapperSize();
+        }
+        return size;
+    }
+
 }

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

@@ -16,37 +16,106 @@
 public final class MapperMergeContext {
 
     private final MapperBuilderContext mapperBuilderContext;
+    private final NewFieldsBudget newFieldsBudget;
+
+    private MapperMergeContext(MapperBuilderContext mapperBuilderContext, NewFieldsBudget newFieldsBudget) {
+        this.mapperBuilderContext = mapperBuilderContext;
+        this.newFieldsBudget = newFieldsBudget;
+    }
 
     /**
      * The root context, to be used when merging a tree of mappers
      */
-    public static MapperMergeContext root(boolean isSourceSynthetic, boolean isDataStream) {
-        return new MapperMergeContext(MapperBuilderContext.root(isSourceSynthetic, isDataStream));
+    public static MapperMergeContext root(boolean isSourceSynthetic, boolean isDataStream, long newFieldsBudget) {
+        return new MapperMergeContext(MapperBuilderContext.root(isSourceSynthetic, isDataStream), NewFieldsBudget.of(newFieldsBudget));
     }
 
     /**
      * Creates a new {@link MapperMergeContext} from a {@link MapperBuilderContext}
      * @param mapperBuilderContext the {@link MapperBuilderContext} for this {@link MapperMergeContext}
+     * @param newFieldsBudget limits how many fields can be added during the merge process
      * @return a new {@link MapperMergeContext}, wrapping the provided {@link MapperBuilderContext}
      */
-    public static MapperMergeContext from(MapperBuilderContext mapperBuilderContext) {
-        return new MapperMergeContext(mapperBuilderContext);
-    }
-
-    private MapperMergeContext(MapperBuilderContext mapperBuilderContext) {
-        this.mapperBuilderContext = mapperBuilderContext;
+    public static MapperMergeContext from(MapperBuilderContext mapperBuilderContext, long newFieldsBudget) {
+        return new MapperMergeContext(mapperBuilderContext, NewFieldsBudget.of(newFieldsBudget));
     }
 
     /**
-     * Creates a new {@link MapperMergeContext} that is a child of this context
+     * Creates a new {@link MapperMergeContext} with a child {@link MapperBuilderContext}.
+     * The child {@link MapperMergeContext} context will share the same field limit.
      * @param name the name of the child context
      * @return a new {@link MapperMergeContext} with this context as its parent
      */
-    public MapperMergeContext createChildContext(String name) {
-        return new MapperMergeContext(mapperBuilderContext.createChildContext(name));
+    MapperMergeContext createChildContext(String name) {
+        return createChildContext(mapperBuilderContext.createChildContext(name));
+    }
+
+    /**
+     * Creates a new {@link MapperMergeContext} with a given child {@link MapperBuilderContext}
+     * The child {@link MapperMergeContext} context will share the same field limit.
+     * @param childContext the child {@link MapperBuilderContext}
+     * @return a new {@link MapperMergeContext}, wrapping the provided {@link MapperBuilderContext}
+     */
+    MapperMergeContext createChildContext(MapperBuilderContext childContext) {
+        return new MapperMergeContext(childContext, newFieldsBudget);
     }
 
     MapperBuilderContext getMapperBuilderContext() {
         return mapperBuilderContext;
     }
+
+    boolean decrementFieldBudgetIfPossible(int fieldSize) {
+        return newFieldsBudget.decrementIfPossible(fieldSize);
+    }
+
+    /**
+     * Keeps track of how many new fields can be added during mapper merge.
+     * The field budget is shared across instances of {@link MapperMergeContext} that are created via
+     * {@link MapperMergeContext#createChildContext}.
+     * This ensures that fields that are consumed by one child object mapper also decrement the budget for another child object.
+     * Not thread safe.The same instance may not be modified by multiple threads.
+     */
+    private interface NewFieldsBudget {
+
+        static NewFieldsBudget of(long fieldsBudget) {
+            if (fieldsBudget == Long.MAX_VALUE) {
+                return Unlimited.INSTANCE;
+            }
+            return new Limited(fieldsBudget);
+        }
+
+        boolean decrementIfPossible(long fieldSize);
+
+        final class Unlimited implements NewFieldsBudget {
+
+            private static final Unlimited INSTANCE = new Unlimited();
+
+            private Unlimited() {}
+
+            @Override
+            public boolean decrementIfPossible(long fieldSize) {
+                return true;
+            }
+
+        }
+
+        final class Limited implements NewFieldsBudget {
+
+            private long fieldsBudget;
+
+            Limited(long fieldsBudget) {
+                this.fieldsBudget = fieldsBudget;
+            }
+
+            @Override
+            public boolean decrementIfPossible(long fieldSize) {
+                if (fieldsBudget >= fieldSize) {
+                    fieldsBudget -= fieldSize;
+                    return true;
+                }
+                return false;
+            }
+
+        }
+    }
 }

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

@@ -559,11 +559,15 @@ public Mapping parseMapping(String mappingType, Map<String, Object> mappingSourc
     }
 
     public static Mapping mergeMappings(DocumentMapper currentMapper, Mapping incomingMapping, MergeReason reason) {
+        return mergeMappings(currentMapper, incomingMapping, reason, Long.MAX_VALUE);
+    }
+
+    static Mapping mergeMappings(DocumentMapper currentMapper, Mapping incomingMapping, MergeReason reason, long newFieldsBudget) {
         Mapping newMapping;
         if (currentMapper == null) {
-            newMapping = incomingMapping;
+            newMapping = incomingMapping.withFieldsBudget(newFieldsBudget);
         } else {
-            newMapping = currentMapper.mapping().merge(incomingMapping, reason);
+            newMapping = currentMapper.mapping().merge(incomingMapping, reason, newFieldsBudget);
         }
         return newMapping;
     }

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

@@ -133,10 +133,12 @@ public SourceLoader.SyntheticFieldLoader syntheticFieldLoader() {
      *
      * @param mergeWith the new mapping to merge into this one.
      * @param reason the reason this merge was initiated.
+     * @param newFieldsBudget how many new fields can be added during the merge process
      * @return the resulting merged mapping.
      */
-    Mapping merge(Mapping mergeWith, MergeReason reason) {
-        RootObjectMapper mergedRoot = root.merge(mergeWith.root, reason, MapperMergeContext.root(isSourceSynthetic(), false));
+    Mapping merge(Mapping mergeWith, MergeReason reason, long newFieldsBudget) {
+        MapperMergeContext mergeContext = MapperMergeContext.root(isSourceSynthetic(), false, newFieldsBudget);
+        RootObjectMapper mergedRoot = root.merge(mergeWith.root, reason, mergeContext);
 
         // When merging metadata fields as part of applying an index template, new field definitions
         // completely overwrite existing ones instead of being merged. This behavior matches how we
@@ -148,7 +150,7 @@ Mapping merge(Mapping mergeWith, MergeReason reason) {
             if (mergeInto == null || reason == MergeReason.INDEX_TEMPLATE) {
                 merged = metaMergeWith;
             } else {
-                merged = (MetadataFieldMapper) mergeInto.merge(metaMergeWith, MapperMergeContext.root(isSourceSynthetic(), false));
+                merged = (MetadataFieldMapper) mergeInto.merge(metaMergeWith, mergeContext);
             }
             mergedMetadataMappers.put(merged.getClass(), merged);
         }
@@ -169,6 +171,18 @@ Mapping merge(Mapping mergeWith, MergeReason reason) {
         return new Mapping(mergedRoot, mergedMetadataMappers.values().toArray(new MetadataFieldMapper[0]), mergedMeta);
     }
 
+    /**
+     * Returns a copy of this mapper that ensures that the number of fields isn't greater than the provided fields budget.
+     * @param fieldsBudget the maximum number of fields this mapping may have
+     */
+    public Mapping withFieldsBudget(long fieldsBudget) {
+        MapperMergeContext mergeContext = MapperMergeContext.root(isSourceSynthetic(), false, fieldsBudget);
+        // get a copy of the root mapper, without any fields
+        RootObjectMapper shallowRoot = root.withoutMappers();
+        // calling merge on the shallow root to ensure we're only adding as many fields as allowed by the fields budget
+        return new Mapping(shallowRoot.merge(root, MergeReason.MAPPING_RECOVERY, mergeContext), metadataMappers, meta);
+    }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         root.toXContent(builder, params, (b, params1) -> {

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

@@ -271,7 +271,7 @@ private void checkFieldLimit(long limit) {
     }
 
     void checkFieldLimit(long limit, int additionalFieldsToAdd) {
-        if (getTotalFieldsCount() + additionalFieldsToAdd - mapping.getSortedMetadataMappers().length > limit) {
+        if (exceedsLimit(limit, additionalFieldsToAdd)) {
             throw new IllegalArgumentException(
                 "Limit of total fields ["
                     + limit
@@ -281,6 +281,10 @@ void checkFieldLimit(long limit, int additionalFieldsToAdd) {
         }
     }
 
+    boolean exceedsLimit(long limit, int additionalFieldsToAdd) {
+        return getTotalFieldsCount() + additionalFieldsToAdd - mapping.getSortedMetadataMappers().length > limit;
+    }
+
     private void checkDimensionFieldLimit(long limit) {
         long dimensionFieldCount = fieldMappers.values()
             .stream()

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

@@ -186,6 +186,21 @@ public ObjectMapper.Builder newBuilder(IndexVersion indexVersionCreated) {
         return builder;
     }
 
+    @Override
+    NestedObjectMapper withoutMappers() {
+        return new NestedObjectMapper(
+            simpleName(),
+            fullPath(),
+            Map.of(),
+            enabled,
+            dynamic,
+            includeInParent,
+            includeInRoot,
+            nestedTypePath,
+            nestedTypeFilter
+        );
+    }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject(simpleName());
@@ -260,7 +275,9 @@ protected MapperMergeContext createChildContext(MapperMergeContext mapperMergeCo
         if (mapperBuilderContext instanceof NestedMapperBuilderContext == false) {
             parentIncludedInRoot |= this.includeInParent.value();
         }
-        return MapperMergeContext.from(new NestedMapperBuilderContext(mapperBuilderContext.buildFullName(name), parentIncludedInRoot));
+        return mapperMergeContext.createChildContext(
+            new NestedMapperBuilderContext(mapperBuilderContext.buildFullName(name), parentIncludedInRoot)
+        );
     }
 
     @Override

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

@@ -162,7 +162,7 @@ protected final Map<String, Mapper> buildMappers(MapperBuilderContext mapperBuil
                     // This can also happen due to multiple index templates being merged into a single mappings definition using
                     // XContentHelper#mergeDefaults, again in case some index templates contained mappings for the same field using a
                     // mix of object notation and dot notation.
-                    mapper = existing.merge(mapper, MapperMergeContext.from(mapperBuilderContext));
+                    mapper = existing.merge(mapper, MapperMergeContext.from(mapperBuilderContext, Long.MAX_VALUE));
                 }
                 mappers.put(mapper.simpleName(), mapper);
             }
@@ -403,6 +403,14 @@ public ObjectMapper.Builder newBuilder(IndexVersion indexVersionCreated) {
         return builder;
     }
 
+    /**
+     * Returns a copy of this object mapper that doesn't have any fields and runtime fields.
+     * This is typically used in the context of a mapper merge when there's not enough budget to add the entire object.
+     */
+    ObjectMapper withoutMappers() {
+        return new ObjectMapper(simpleName(), fullPath, enabled, subobjects, dynamic, Map.of());
+    }
+
     @Override
     public String name() {
         return this.fullPath;
@@ -542,10 +550,15 @@ private static Map<String, Mapper> buildMergedMappers(
                 Mapper mergeWithMapper = iterator.next();
                 Mapper mergeIntoMapper = mergedMappers.get(mergeWithMapper.simpleName());
 
+                Mapper merged = null;
                 if (mergeIntoMapper == null) {
-                    mergedMappers.put(mergeWithMapper.simpleName(), mergeWithMapper);
+                    if (objectMergeContext.decrementFieldBudgetIfPossible(mergeWithMapper.mapperSize())) {
+                        merged = mergeWithMapper;
+                    } else if (mergeWithMapper instanceof ObjectMapper om) {
+                        merged = truncateObjectMapper(reason, objectMergeContext, om);
+                    }
                 } else if (mergeIntoMapper instanceof ObjectMapper objectMapper) {
-                    mergedMappers.put(objectMapper.simpleName(), objectMapper.merge(mergeWithMapper, reason, objectMergeContext));
+                    merged = objectMapper.merge(mergeWithMapper, reason, objectMergeContext);
                 } else {
                     assert mergeIntoMapper instanceof FieldMapper || mergeIntoMapper instanceof FieldAliasMapper;
                     if (mergeWithMapper instanceof NestedObjectMapper) {
@@ -557,14 +570,28 @@ private static Map<String, Mapper> buildMergedMappers(
                     // If we're merging template mappings when creating an index, then a field definition always
                     // replaces an existing one.
                     if (reason == MergeReason.INDEX_TEMPLATE) {
-                        mergedMappers.put(mergeWithMapper.simpleName(), mergeWithMapper);
+                        merged = mergeWithMapper;
                     } else {
-                        mergedMappers.put(mergeWithMapper.simpleName(), mergeIntoMapper.merge(mergeWithMapper, objectMergeContext));
+                        merged = mergeIntoMapper.merge(mergeWithMapper, objectMergeContext);
                     }
                 }
+                if (merged != null) {
+                    mergedMappers.put(merged.simpleName(), merged);
+                }
             }
             return Map.copyOf(mergedMappers);
         }
+
+        private static ObjectMapper truncateObjectMapper(MergeReason reason, MapperMergeContext context, ObjectMapper objectMapper) {
+            // there's not enough capacity for the whole object mapper,
+            // so we're just trying to add the shallow object, without it's sub-fields
+            ObjectMapper shallowObjectMapper = objectMapper.withoutMappers();
+            if (context.decrementFieldBudgetIfPossible(shallowObjectMapper.mapperSize())) {
+                // now trying to add the sub-fields one by one via a merge, until we hit the limit
+                return shallowObjectMapper.merge(objectMapper, reason, context);
+            }
+            return null;
+        }
     }
 
     @Override

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

@@ -160,7 +160,7 @@ public void addDynamicMappingsUpdate(Mapping update) {
         if (dynamicMappingsUpdate == null) {
             dynamicMappingsUpdate = update;
         } else {
-            dynamicMappingsUpdate = dynamicMappingsUpdate.merge(update, MergeReason.MAPPING_UPDATE);
+            dynamicMappingsUpdate = dynamicMappingsUpdate.merge(update, MergeReason.MAPPING_UPDATE, Long.MAX_VALUE);
         }
     }