Simple thread local store metric capture

A simple thread local store metric capture to allow exposing this data to queries ultimately at the coordinating level. Exemplified here through bytes read but we will likely have more for stateless too.

Author: henningandersen

server/src/main/java/org/elasticsearch/index/IndexModule.java

@@ -57,6 +57,8 @@
 import org.elasticsearch.index.shard.ShardPath;
 import org.elasticsearch.index.similarity.SimilarityService;
 import org.elasticsearch.index.store.FsDirectoryFactory;
+import org.elasticsearch.index.store.MetricHolder;
+import org.elasticsearch.index.store.StoreMetrics;
 import org.elasticsearch.indices.IndicesQueryCache;
 import org.elasticsearch.indices.breaker.CircuitBreakerService;
 import org.elasticsearch.indices.fielddata.cache.IndicesFieldDataCache;
@@ -182,6 +184,7 @@ public interface DirectoryWrapper {
     private final IndexingStatsSettings indexingStatsSettings;
     private final SearchStatsSettings searchStatsSettings;
     private final MergeMetrics mergeMetrics;
+    private MetricHolder<StoreMetrics> metricHolder;
 
     /**
      * Construct the index module for the index with the specified index settings. The index module contains extension points for plugins
@@ -206,7 +209,8 @@ public IndexModule(
         final List<SearchOperationListener> searchOperationListeners,
         final IndexingStatsSettings indexingStatsSettings,
         final SearchStatsSettings searchStatsSettings,
-        final MergeMetrics mergeMetrics
+        final MergeMetrics mergeMetrics,
+        MetricHolder<StoreMetrics> metricHolder
     ) {
         this.indexSettings = indexSettings;
         this.analysisRegistry = analysisRegistry;
@@ -224,6 +228,7 @@ public IndexModule(
         this.indexingStatsSettings = indexingStatsSettings;
         this.searchStatsSettings = searchStatsSettings;
         this.mergeMetrics = mergeMetrics;
+        this.metricHolder = metricHolder;
     }
 
     /**
@@ -560,7 +565,8 @@ public IndexService newIndexService(
                 mapperMetrics,
                 indexingStatsSettings,
                 searchStatsSettings,
-                mergeMetrics
+                mergeMetrics,
+                metricHolder
             );
             success = true;
             return indexService;

server/src/main/java/org/elasticsearch/index/IndexService.java

@@ -82,7 +82,9 @@
 import org.elasticsearch.index.shard.ShardNotInPrimaryModeException;
 import org.elasticsearch.index.shard.ShardPath;
 import org.elasticsearch.index.similarity.SimilarityService;
+import org.elasticsearch.index.store.MetricHolder;
 import org.elasticsearch.index.store.Store;
+import org.elasticsearch.index.store.StoreMetrics;
 import org.elasticsearch.index.translog.Translog;
 import org.elasticsearch.indices.breaker.CircuitBreakerService;
 import org.elasticsearch.indices.cluster.IndexRemovalReason;
@@ -172,6 +174,7 @@ public class IndexService extends AbstractIndexComponent implements IndicesClust
     private final IndexingStatsSettings indexingStatsSettings;
     private final SearchStatsSettings searchStatsSettings;
     private final MergeMetrics mergeMetrics;
+    private final MetricHolder<StoreMetrics> metricHolder;
 
     @SuppressWarnings("this-escape")
     public IndexService(
@@ -210,7 +213,8 @@ public IndexService(
         MapperMetrics mapperMetrics,
         IndexingStatsSettings indexingStatsSettings,
         SearchStatsSettings searchStatsSettings,
-        MergeMetrics mergeMetrics
+        MergeMetrics mergeMetrics,
+        MetricHolder<StoreMetrics> metricHolder
     ) {
         super(indexSettings);
         assert indexCreationContext != IndexCreationContext.RELOAD_ANALYZERS
@@ -297,6 +301,7 @@ public IndexService(
         this.indexingStatsSettings = indexingStatsSettings;
         this.searchStatsSettings = searchStatsSettings;
         this.mergeMetrics = mergeMetrics;
+        this.metricHolder = metricHolder;
         updateFsyncTaskIfNecessary();
     }
 
@@ -559,7 +564,8 @@ public synchronized IndexShard createShard(
                 directory,
                 lock,
                 new StoreCloseListener(shardId, () -> eventListener.onStoreClosed(shardId)),
-                this.indexSettings.getIndexSortConfig().hasIndexSort()
+                this.indexSettings.getIndexSortConfig().hasIndexSort(),
+                metricHolder
             );
             eventListener.onStoreCreated(shardId);
             indexShard = new IndexShard(

server/src/main/java/org/elasticsearch/index/store/MetricHolder.java

@@ -0,0 +1,31 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", 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", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.store;
+
+public interface MetricHolder<M> {
+    static <M> MetricHolder<M> noop(M noopData) {
+        return new MetricHolder<>() {
+
+            @Override
+            public MetricHolder<M> singleThreaded() {
+                return this;
+            }
+
+            @Override
+            public M instance() {
+                return noopData;
+            }
+        };
+    }
+
+    M instance();
+
+    MetricHolder<M> singleThreaded();
+}

server/src/main/java/org/elasticsearch/index/store/SingleThreadMetricHolder.java

@@ -0,0 +1,42 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", 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", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.store;
+
+public class SingleThreadMetricHolder<M> implements MetricHolder<M> {
+    private final ThreadLocalMetricHolder<M> delegate;
+    private Thread owner;
+    private M current;
+
+    public SingleThreadMetricHolder(ThreadLocalMetricHolder<M> delegate) {
+        this.delegate = delegate;
+    }
+
+    public SingleThreadMetricHolder(SingleThreadMetricHolder<M> source) {
+        this.delegate = source.delegate;
+        this.owner = source.owner;
+        this.current = source.current;
+    }
+
+    @Override
+    public MetricHolder<M> singleThreaded() {
+        return new SingleThreadMetricHolder<>(this);
+    }
+
+    public M instance() {
+        Thread thread = Thread.currentThread();
+        if (owner == thread) {
+            return current;
+        } else {
+            current = delegate.instance();
+            owner = thread;
+            return current;
+        }
+    }
+}

server/src/main/java/org/elasticsearch/index/store/Store.java

@@ -183,10 +183,22 @@ public Store(
         ShardLock shardLock,
         OnClose onClose,
         boolean hasIndexSort
+    ) {
+        this(shardId, indexSettings, directory, shardLock, onClose, hasIndexSort, new ThreadLocalMetricHolder<>(StoreMetrics::new));
+    }
+
+    public Store(
+        ShardId shardId,
+        IndexSettings indexSettings,
+        Directory directory,
+        ShardLock shardLock,
+        OnClose onClose,
+        boolean hasIndexSort,
+        MetricHolder<StoreMetrics> metricHolder
     ) {
         super(shardId, indexSettings);
         this.directory = new StoreDirectory(
-            byteSizeDirectory(directory, indexSettings, logger),
+            byteSizeDirectory(new StoreMetricsDirectory(directory, metricHolder), indexSettings, logger),
             Loggers.getLogger("index.store.deletes", shardId)
         );
         this.shardLock = shardLock;

server/src/main/java/org/elasticsearch/index/store/StoreMetricIndexInput.java

@@ -0,0 +1,100 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", 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", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.store;
+
+import org.apache.lucene.store.FilterIndexInput;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.RandomAccessInput;
+
+import java.io.IOException;
+
+public class StoreMetricIndexInput extends FilterIndexInput {
+    private final MetricHolder<StoreMetrics> metricHolder;
+
+    public StoreMetricIndexInput(String resourceDescription, IndexInput in, MetricHolder<StoreMetrics> metricHolder) {
+        super(resourceDescription, in);
+        this.metricHolder = metricHolder;
+    }
+
+    @Override
+    public byte readByte() throws IOException {
+        byte result = super.readByte();
+        metricHolder.instance().addBytesRead(1);
+        return result;
+    }
+
+    @Override
+    public void readBytes(byte[] b, int offset, int len) throws IOException {
+        super.readBytes(b, offset, len);
+        metricHolder.instance().addBytesRead(len);
+    }
+
+    @Override
+    public IndexInput clone() {
+        return new StoreMetricIndexInput(toString(), in.clone(), metricHolder.singleThreaded());
+    }
+
+    @Override
+    public IndexInput slice(String sliceDescription, long offset, long length) throws IOException {
+        return new StoreMetricIndexInput(sliceDescription, super.slice(sliceDescription, offset, length), metricHolder.singleThreaded());
+    }
+
+    @Override
+    public IndexInput slice(String sliceDescription, long offset, long length, IOContext context) throws IOException {
+        return new StoreMetricIndexInput(
+            sliceDescription,
+            super.slice(sliceDescription, offset, length, context),
+            metricHolder.singleThreaded()
+        );
+    }
+
+    @Override
+    public RandomAccessInput randomAccessSlice(long offset, long length) throws IOException {
+        RandomAccessInput delegate = in.randomAccessSlice(offset, length);
+
+        return new RandomAccessInput() {
+            private final MetricHolder<StoreMetrics> metricHolder = StoreMetricIndexInput.this.metricHolder;
+
+            @Override
+            public long length() {
+                return delegate.length();
+            }
+
+            @Override
+            public byte readByte(long pos) throws IOException {
+                byte result = delegate.readByte(pos);
+                metricHolder.instance().addBytesRead(1);
+                return result;
+            }
+
+            @Override
+            public short readShort(long pos) throws IOException {
+                short result = delegate.readShort(pos);
+                metricHolder.instance().addBytesRead(2);
+                return result;
+            }
+
+            @Override
+            public int readInt(long pos) throws IOException {
+                int result = delegate.readInt(pos);
+                metricHolder.instance().addBytesRead(4);
+                return result;
+            }
+
+            @Override
+            public long readLong(long pos) throws IOException {
+                long result = delegate.readLong(pos);
+                metricHolder.instance().addBytesRead(8);
+                return result;
+            }
+        };
+    }
+}

server/src/main/java/org/elasticsearch/index/store/StoreMetrics.java

@@ -0,0 +1,34 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", 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", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.store;
+
+public class StoreMetrics {
+    public static final MetricHolder<StoreMetrics> NOOP_HOLDER = MetricHolder.noop(new StoreMetrics() {
+        @Override
+        public void addBytesRead(long amount) {
+        }
+    });
+
+    private long bytesRead;
+
+    public long getBytesRead() {
+        return bytesRead;
+    }
+
+    public StoreMetrics snapshot() {
+        StoreMetrics metrics = new StoreMetrics();
+        metrics.bytesRead = bytesRead;
+        return metrics;
+    }
+
+    public void addBytesRead(long amount) {
+        bytesRead += amount;
+    }
+}

server/src/main/java/org/elasticsearch/index/store/StoreMetricsDirectory.java

@@ -0,0 +1,32 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", 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", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.store;
+
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.FilterDirectory;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexInput;
+
+import java.io.IOException;
+
+public class StoreMetricsDirectory extends FilterDirectory {
+    private final MetricHolder<StoreMetrics> metricHandler;
+
+    public StoreMetricsDirectory(Directory in, MetricHolder<StoreMetrics> metricHandler) {
+        super(in);
+        this.metricHandler = metricHandler;
+    }
+
+    @Override
+    public IndexInput openInput(String name, IOContext context) throws IOException {
+        return new StoreMetricIndexInput(name, super.openInput(name, context), metricHandler.singleThreaded());
+    }
+
+}

server/src/main/java/org/elasticsearch/index/store/ThreadLocalMetricHolder.java

@@ -0,0 +1,35 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", 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", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.store;
+
+import java.util.function.Supplier;
+
+public class ThreadLocalMetricHolder<M> implements MetricHolder<M> {
+    private final Supplier<? extends M> metricsSupplier;
+    private final ThreadLocal<M> threadLocal = new ThreadLocal<>() {
+        @Override
+        protected M initialValue() {
+            return metricsSupplier.get();
+        }
+    };
+
+    public ThreadLocalMetricHolder(Supplier<? extends M> metricsSupplier) {
+        this.metricsSupplier = metricsSupplier;
+    }
+
+    @Override
+    public MetricHolder<M> singleThreaded() {
+        return new SingleThreadMetricHolder<>(this);
+    }
+
+    public M instance() {
+        return threadLocal.get();
+    }
+}