Cleanup

server/src/main/java/org/elasticsearch/common/util/ValueWithThreadLocalOverride.java

@@ -0,0 +1,67 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.common.util;
+
+import java.util.function.Supplier;
+
+/**
+ * Stores value and provides optional functionality to set up a per-thread override value.
+ * Intended usage is a singleton value that is commonly accessed from multiple places.
+ * Having it as singleton allows to not pass instance down to every consumer.
+ * Thread-local override allows to use different value in tests even though it is a singleton.
+ * Inspired by <a href="https://docs.rs/tracing/latest/tracing/dispatcher/index.html">tracing</a>.
+ */
+public class ValueWithThreadLocalOverride<T> {
+    // Intentionally not static - different values should allow different overrides.
+    private final ThreadLocal<T> threadLocal = ThreadLocal.withInitial(() -> null);
+    private Supplier<T> supplier;
+
+    public ValueWithThreadLocalOverride(T value) {
+        this.supplier = () -> value;
+    }
+
+    /**
+     * returns stored value or an override if set
+     * @return T
+     */
+    public T get() {
+        return supplier.get();
+    }
+
+    /**
+     * Installs a thread-local override value.
+     * @param value
+     * @return an {@link AutoCloseable} that removes the override.
+     */
+    public AutoCloseable withOverride(T value) {
+        threadLocal.set(value);
+        // This is a small optimization to eliminate thread local lookup
+        // if override was never set, which is most of the time.
+        T original = supplier.get();
+        this.supplier = () -> getWithOverride(original);
+
+        return new Reset(threadLocal);
+    }
+
+    private T getWithOverride(T original) {
+        var local = threadLocal.get();
+        if (local != null) {
+            return local;
+        }
+
+        return original;
+    }
+
+    private record Reset(ThreadLocal<?> threadLocal) implements AutoCloseable {
+        @Override
+        public void close() throws Exception {
+            threadLocal.remove();
+        }
+    }
+}

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

@@ -113,7 +113,7 @@ private record LeafWithMetrics(Leaf leaf) implements Leaf {
 
             @Override
             public Source source(LeafStoredFieldLoader storedFields, int docId) throws IOException {
-                SourceFieldMetrics metrics = MapperMetrics.INSTANCE.get().getSyntheticSourceMetrics();
+                SourceFieldMetrics metrics = MapperMetrics.SOURCE_FIELD_METRICS.get();
                 long startTime = metrics.getRelativeTimeSupplier().getAsLong();
 
                 var source = leaf.source(storedFields, docId);

server/src/main/java/org/elasticsearch/indices/MapperMetrics.java

@@ -8,29 +8,23 @@
 
 package org.elasticsearch.indices;
 
-import org.elasticsearch.common.metrics.MetricAccessor;
+import org.elasticsearch.common.util.ValueWithThreadLocalOverride;
 import org.elasticsearch.index.mapper.SourceFieldMetrics;
 
 /**
  * Groups together all metrics used in mappers.
  * Main purpose of this class is to avoid verbosity of passing individual metric instances around.
  */
 public class MapperMetrics {
-    public static MapperMetrics NOOP = new MapperMetrics(SourceFieldMetrics.NOOP);
+    public static MapperMetrics NOOP = new MapperMetrics();
 
-    public static MetricAccessor<MapperMetrics> INSTANCE = new MetricAccessor<>(NOOP);
+    public static ValueWithThreadLocalOverride<SourceFieldMetrics> SOURCE_FIELD_METRICS = new ValueWithThreadLocalOverride<>(
+        SourceFieldMetrics.NOOP
+    );
 
     public static void init(SourceFieldMetrics sourceFieldMetrics) {
-        INSTANCE = new MetricAccessor<>(new MapperMetrics(sourceFieldMetrics));
+        SOURCE_FIELD_METRICS = new ValueWithThreadLocalOverride<>(sourceFieldMetrics);
     }
 
-    private final SourceFieldMetrics sourceFieldMetrics;
-
-    public MapperMetrics(SourceFieldMetrics sourceFieldMetrics) {
-        this.sourceFieldMetrics = sourceFieldMetrics;
-    }
-
-    public SourceFieldMetrics getSyntheticSourceMetrics() {
-        return sourceFieldMetrics;
-    }
+    private MapperMetrics() {}
 }

server/src/test/java/org/elasticsearch/index/mapper/SourceLoaderTests.java

@@ -136,7 +136,7 @@ public void testHideTheCopyTo() {
         assertThat(e.getMessage(), equalTo("[copy_to] may not be used to copy from a multi-field: [foo.hidden]"));
     }
 
-    public void testSyntheticSourceTelemetry() throws Exception {
+    public void testSyntheticSourceMetrics() throws Exception {
         var testTelemetry = new TestTelemetryPlugin();
         var sourceFieldMetrics = new SourceFieldMetrics(
             testTelemetry.getTelemetryProvider(Settings.EMPTY).getMeterRegistry(),
@@ -153,7 +153,7 @@ public long getAsLong() {
         DocumentMapper mapper = createDocumentMapper(
             syntheticSourceMapping(b -> { b.startObject("kwd").field("type", "keyword").endObject(); })
         );
-        try (var ignored = MapperMetrics.INSTANCE.initForTest(new MapperMetrics(sourceFieldMetrics))) {
+        try (var ignored = MapperMetrics.SOURCE_FIELD_METRICS.withOverride(sourceFieldMetrics)) {
             assertThat(syntheticSource(mapper, b -> b.field("kwd", "foo")), equalTo("""
                 {"kwd":"foo"}"""));