Account for time taken to write index buffers in IndexingMemoryController (#126786)

This PR adds to the indexing write load, the time taken to flush write indexing buffers using the indexing threads (this is done here to push back on indexing)

This changes the semantics of InternalIndexingStats#recentIndexMetric and InternalIndexingStats#peakIndexMetric  to more accurately account for load on the indexing thread. Address ES-11356.

Author: ankikuma

docs/changelog/126786.yaml

@@ -0,0 +1,5 @@
+pr: 126786
+summary: Account for time taken to write index buffers in `IndexingMemoryController`
+area: Distributed
+type: enhancement
+issues: []

modules/data-streams/src/internalClusterTest/java/org/elasticsearch/datastreams/DataStreamAutoshardingIT.java

@@ -527,7 +527,9 @@ private static ShardStats getShardStats(IndexMetadata indexMeta, int shardIndex,
         CommonStats stats = new CommonStats();
         stats.docs = new DocsStats(100, 0, randomByteSizeValue().getBytes());
         stats.store = new StoreStats();
-        stats.indexing = new IndexingStats(new IndexingStats.Stats(1, 1, 1, 1, 1, 1, 1, 1, 1, false, 1, targetWriteLoad, 1, 0.123, 0.234));
+        stats.indexing = new IndexingStats(
+            new IndexingStats.Stats(1, 1, 1, 1, 1, 1, 1, 1, 1, false, 1, targetWriteLoad, targetWriteLoad, 1, 0.123, 0.234)
+        );
         return new ShardStats(shardRouting, new ShardPath(false, path, path, shardId), stats, null, null, null, false, 0);
     }
 

server/src/internalClusterTest/java/org/elasticsearch/indices/IndexingMemoryControllerIT.java

@@ -29,6 +29,7 @@
 import java.util.Optional;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;
+import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
 public class IndexingMemoryControllerIT extends ESSingleNodeTestCase {
@@ -37,7 +38,9 @@ public class IndexingMemoryControllerIT extends ESSingleNodeTestCase {
     protected Settings nodeSettings() {
         return Settings.builder()
             .put(super.nodeSettings())
-            // small indexing buffer so that we can trigger refresh after buffering 100 deletes
+            // small indexing buffer so that
+            // 1. We can trigger refresh after buffering 100 deletes
+            // 2. Indexing memory Controller writes indexing buffers in sync with indexing on the indexing thread
             .put("indices.memory.index_buffer_size", "1kb")
             .build();
     }
@@ -111,4 +114,22 @@ public void testDeletesAloneCanTriggerRefresh() throws Exception {
         }
         assertThat(shard.getEngineOrNull().getIndexBufferRAMBytesUsed(), lessThanOrEqualTo(ByteSizeUnit.KB.toBytes(1)));
     }
+
+    /* When there is memory pressure, we write indexing buffers to disk on the same thread as the indexing thread,
+     * @see org.elasticsearch.indices.IndexingMemoryController.
+     * This test verifies that we update the stats that capture the combined time for indexing + writing the
+     * indexing buffers.
+     * Note that the small indices.memory.index_buffer_size setting is required for this test to work.
+     */
+    public void testIndexingUpdatesRelevantStats() throws Exception {
+        IndexService indexService = createIndex("index", indexSettings(1, 0).put("index.refresh_interval", -1).build());
+        IndexShard shard = indexService.getShard(0);
+        prepareIndex("index").setSource("field", randomUnicodeOfCodepointLengthBetween(10, 25)).get();
+        // Check that
+        assertThat(shard.indexingStats().getTotal().getTotalIndexingExecutionTimeInMillis(), greaterThan(0L));
+        assertThat(
+            shard.indexingStats().getTotal().getTotalIndexingExecutionTimeInMillis(),
+            greaterThan(shard.indexingStats().getTotal().getIndexTime().getMillis())
+        );
+    }
 }

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -234,6 +234,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion AGGREGATE_METRIC_DOUBLE_BLOCK = def(9_067_00_0);
     public static final TransportVersion PINNED_RETRIEVER = def(9_068_0_00);
     public static final TransportVersion ML_INFERENCE_SAGEMAKER = def(9_069_0_00);
+    public static final TransportVersion WRITE_LOAD_INCLUDES_BUFFER_WRITES = def(9_070_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/index/shard/IndexShard.java

@@ -308,6 +308,7 @@ public class IndexShard extends AbstractIndexShardComponent implements IndicesCl
     private final LongSupplier relativeTimeInNanosSupplier;
     private volatile long startedRelativeTimeInNanos = -1L; // use -1 to indicate this has not yet been set to its true value
     private volatile long indexingTimeBeforeShardStartedInNanos;
+    private volatile long indexingTaskExecutionTimeBeforeShardStartedInNanos;
     private volatile double recentIndexingLoadAtShardStarted;
     private final SubscribableListener<Void> waitForEngineOrClosedShardListeners = new SubscribableListener<>();
 
@@ -569,6 +570,7 @@ public void updateShardState(
                 // unlikely case that getRelativeTimeInNanos() returns exactly -1, we advance by 1ns to avoid that special value.
                 startedRelativeTimeInNanos = (relativeTimeInNanos != -1L) ? relativeTimeInNanos : 0L;
                 indexingTimeBeforeShardStartedInNanos = internalIndexingStats.totalIndexingTimeInNanos();
+                indexingTaskExecutionTimeBeforeShardStartedInNanos = internalIndexingStats.totalIndexingExecutionTimeInNanos();
                 recentIndexingLoadAtShardStarted = internalIndexingStats.recentIndexingLoad(startedRelativeTimeInNanos);
             } else if (currentRouting.primary()
                 && currentRouting.relocating()
@@ -1401,6 +1403,7 @@ public IndexingStats indexingStats() {
             throttled,
             throttleTimeInMillis,
             indexingTimeBeforeShardStartedInNanos,
+            indexingTaskExecutionTimeBeforeShardStartedInNanos,
             timeSinceShardStartedInNanos,
             currentTimeInNanos,
             recentIndexingLoadAtShardStarted
@@ -3235,6 +3238,16 @@ public void noopUpdate() {
         internalIndexingStats.noopUpdate();
     }
 
+    /**
+     * Increment relevant stats when indexing buffers are written to disk using indexing threads,
+     * in order to apply back-pressure on indexing.
+     * @param tookInNanos  time it took to write the indexing buffers for this shard (in ns)
+     * @see IndexingMemoryController#writePendingIndexingBuffers()
+     */
+    public void addWriteIndexBuffersToIndexThreadsTime(long tookInNanos) {
+        internalIndexingStats.writeIndexingBuffersTime(tookInNanos);
+    }
+
     public void maybeCheckIndex() {
         recoveryState.setStage(RecoveryState.Stage.VERIFY_INDEX);
         if (Booleans.isTrue(checkIndexOnStartup) || "checksum".equals(checkIndexOnStartup)) {

server/src/main/java/org/elasticsearch/index/shard/IndexingStats.java

@@ -27,6 +27,7 @@
 
 import static org.elasticsearch.TransportVersions.INDEXING_STATS_INCLUDES_RECENT_WRITE_LOAD;
 import static org.elasticsearch.TransportVersions.INDEX_STATS_AND_METADATA_INCLUDE_PEAK_WRITE_LOAD;
+import static org.elasticsearch.TransportVersions.WRITE_LOAD_INCLUDES_BUFFER_WRITES;
 
 public class IndexingStats implements Writeable, ToXContentFragment {
 
@@ -45,6 +46,10 @@ public static class Stats implements Writeable, ToXContentFragment {
         private long throttleTimeInMillis;
         private boolean isThrottled;
         private long totalIndexingTimeSinceShardStartedInNanos;
+        // This is different from totalIndexingTimeSinceShardStartedInNanos, as it also includes the time taken to write indexing buffers
+        // to disk on the same thread as the indexing thread. This happens when we are running low on memory and want to push
+        // back on indexing, see IndexingMemoryController#writePendingIndexingBuffers()
+        private long totalIndexingExecutionTimeSinceShardStartedInNanos;
         private long totalActiveTimeInNanos;
         private double recentIndexingLoad;
         private double peakIndexingLoad;
@@ -87,6 +92,15 @@ public Stats(StreamInput in) throws IOException {
                     ? (double) totalIndexingTimeSinceShardStartedInNanos / totalActiveTimeInNanos
                     : 0;
             }
+            if (in.getTransportVersion().onOrAfter(WRITE_LOAD_INCLUDES_BUFFER_WRITES)) {
+                totalIndexingExecutionTimeSinceShardStartedInNanos = in.readLong();
+            } else {
+                // When getting stats from an older version which doesn't have the more accurate indexing execution time,
+                // better to fall back to the indexing time, rather that assuming zero load:
+                totalIndexingExecutionTimeSinceShardStartedInNanos = totalActiveTimeInNanos > 0
+                    ? totalIndexingTimeSinceShardStartedInNanos
+                    : 0;
+            }
         }
 
         public Stats(
@@ -102,6 +116,7 @@ public Stats(
             boolean isThrottled,
             long throttleTimeInMillis,
             long totalIndexingTimeSinceShardStartedInNanos,
+            long totalIndexingExecutionTimeSinceShardStartedInNanos,
             long totalActiveTimeInNanos,
             double recentIndexingLoad,
             double peakIndexingLoad
@@ -119,6 +134,7 @@ public Stats(
             this.throttleTimeInMillis = throttleTimeInMillis;
             // We store the raw unweighted write load values in order to avoid losing precision when we combine the shard stats
             this.totalIndexingTimeSinceShardStartedInNanos = totalIndexingTimeSinceShardStartedInNanos;
+            this.totalIndexingExecutionTimeSinceShardStartedInNanos = totalIndexingExecutionTimeSinceShardStartedInNanos;
             this.totalActiveTimeInNanos = totalActiveTimeInNanos;
             // We store the weighted write load as a double because the calculation is inherently floating point
             this.recentIndexingLoad = recentIndexingLoad;
@@ -141,8 +157,9 @@ public void add(Stats stats) {
             if (isThrottled != stats.isThrottled) {
                 isThrottled = true; // When combining if one is throttled set result to throttled.
             }
-            // N.B. getWriteLoad() returns the ratio of these sums, which is the average of the ratios weighted by active time:
             totalIndexingTimeSinceShardStartedInNanos += stats.totalIndexingTimeSinceShardStartedInNanos;
+            // N.B. getWriteLoad() returns the ratio of these sums, which is the average of the ratios weighted by active time:
+            totalIndexingExecutionTimeSinceShardStartedInNanos += stats.totalIndexingExecutionTimeSinceShardStartedInNanos;
             totalActiveTimeInNanos += stats.totalActiveTimeInNanos;
             // We want getRecentWriteLoad() and getPeakWriteLoad() for the aggregated stats to also be the average weighted by active time,
             // so we use the updating formula for a weighted mean:
@@ -237,7 +254,7 @@ public long getNoopUpdateCount() {
          * the elapsed time for each shard.
          */
         public double getWriteLoad() {
-            return totalActiveTimeInNanos > 0 ? (double) totalIndexingTimeSinceShardStartedInNanos / totalActiveTimeInNanos : 0;
+            return totalActiveTimeInNanos > 0 ? (double) totalIndexingExecutionTimeSinceShardStartedInNanos / totalActiveTimeInNanos : 0;
         }
 
         /**
@@ -271,6 +288,13 @@ public long getTotalActiveTimeInMillis() {
             return TimeUnit.NANOSECONDS.toMillis(totalActiveTimeInNanos);
         }
 
+        /**
+         * The total amount of time spend on indexing plus writing indexing buffers.
+         */
+        public long getTotalIndexingExecutionTimeInMillis() {
+            return TimeUnit.NANOSECONDS.toMillis(totalIndexingExecutionTimeSinceShardStartedInNanos);
+        }
+
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             out.writeVLong(indexCount);
@@ -296,6 +320,9 @@ public void writeTo(StreamOutput out) throws IOException {
             if (out.getTransportVersion().onOrAfter(INDEX_STATS_AND_METADATA_INCLUDE_PEAK_WRITE_LOAD)) {
                 out.writeDouble(peakIndexingLoad);
             }
+            if (out.getTransportVersion().onOrAfter(WRITE_LOAD_INCLUDES_BUFFER_WRITES)) {
+                out.writeLong(totalIndexingExecutionTimeSinceShardStartedInNanos);
+            }
         }
 
         @Override
@@ -338,6 +365,7 @@ public boolean equals(Object o) {
                 && isThrottled == that.isThrottled
                 && throttleTimeInMillis == that.throttleTimeInMillis
                 && totalIndexingTimeSinceShardStartedInNanos == that.totalIndexingTimeSinceShardStartedInNanos
+                && totalIndexingExecutionTimeSinceShardStartedInNanos == that.totalIndexingExecutionTimeSinceShardStartedInNanos
                 && totalActiveTimeInNanos == that.totalActiveTimeInNanos
                 && recentIndexingLoad == that.recentIndexingLoad
                 && peakIndexingLoad == that.peakIndexingLoad;
@@ -358,6 +386,7 @@ public int hashCode() {
                 isThrottled,
                 throttleTimeInMillis,
                 totalIndexingTimeSinceShardStartedInNanos,
+                totalIndexingExecutionTimeSinceShardStartedInNanos,
                 totalActiveTimeInNanos
             );
         }

server/src/main/java/org/elasticsearch/index/shard/InternalIndexingStats.java

@@ -22,6 +22,7 @@
 
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.atomic.LongAdder;
 import java.util.function.LongSupplier;
 
 import static org.elasticsearch.core.TimeValue.timeValueNanos;
@@ -51,6 +52,7 @@ IndexingStats stats(
         boolean isThrottled,
         long currentThrottleInMillis,
         long indexingTimeBeforeShardStartedInNanos,
+        long indexingLoadBeforeShardStartedInNanos,
         long timeSinceShardStartedInNanos,
         long currentTimeInNanos,
         double recentIndexingLoadAtShardStarted
@@ -59,6 +61,7 @@ IndexingStats stats(
             isThrottled,
             currentThrottleInMillis,
             indexingTimeBeforeShardStartedInNanos,
+            indexingLoadBeforeShardStartedInNanos,
             timeSinceShardStartedInNanos,
             currentTimeInNanos,
             recentIndexingLoadAtShardStarted
@@ -70,6 +73,10 @@ long totalIndexingTimeInNanos() {
         return totalStats.indexMetric.sum();
     }
 
+    long totalIndexingExecutionTimeInNanos() {
+        return totalStats.indexMetric.sum() + totalStats.writeIndexingBufferTime.sum();
+    }
+
     /**
      * Returns an exponentially-weighted moving rate which measures the indexing load, favoring more recent load.
      */
@@ -153,10 +160,26 @@ void noopUpdate() {
         totalStats.noopUpdates.inc();
     }
 
+    /**
+     * Increment relevant stats when indexing buffers are written to disk using indexing threads,
+     * in order to apply back-pressure on indexing.
+     * @param took   time taken to write buffers
+     * @see org.elasticsearch.indices.IndexingMemoryController
+     */
+    void writeIndexingBuffersTime(long took) {
+        totalStats.writeIndexingBufferTime.add(took);
+        totalStats.recentIndexMetric.addIncrement(took, relativeTimeInNanosSupplier.getAsLong());
+    }
+
     static class StatsHolder {
-        private final MeanMetric indexMetric = new MeanMetric(); // Used for the count and total 'took' time (in ns) of index operations
-        private final ExponentiallyWeightedMovingRate recentIndexMetric; // An EWMR of the total 'took' time of index operations (in ns)
-        private final AtomicReference<Double> peakIndexMetric; // The peak value of the EWMR observed in any stats() call
+        // Used for the count and total 'took' time (in ns) of index operations
+        private final MeanMetric indexMetric = new MeanMetric();
+        // Used for the total time taken to flush indexing buffers to disk (on indexing threads) (in ns)
+        private final LongAdder writeIndexingBufferTime = new LongAdder();
+        // An EWMR of the total 'took' time of index operations (indexMetric) plus the writeIndexingBufferTime (in ns)
+        private final ExponentiallyWeightedMovingRate recentIndexMetric;
+        // The peak value of the EWMR (recentIndexMetric) observed in any stats() call
+        private final AtomicReference<Double> peakIndexMetric;
         private final MeanMetric deleteMetric = new MeanMetric();
         private final CounterMetric indexCurrent = new CounterMetric();
         private final CounterMetric indexFailed = new CounterMetric();
@@ -179,12 +202,19 @@ IndexingStats.Stats stats(
             boolean isThrottled,
             long currentThrottleMillis,
             long indexingTimeBeforeShardStartedInNanos,
+            long indexingLoadBeforeShardStartedInNanos,
             long timeSinceShardStartedInNanos,
             long currentTimeInNanos,
             double recentIndexingLoadAtShardStarted
         ) {
             final long totalIndexingTimeInNanos = indexMetric.sum();
             final long totalIndexingTimeSinceShardStartedInNanos = totalIndexingTimeInNanos - indexingTimeBeforeShardStartedInNanos;
+            // This is different from indexing time as it also includes the time taken to write indexing buffers to disk
+            // on the same thread as the indexing thread. This happens when we are running low on memory and want to push
+            // back on indexing, see IndexingMemoryController#writePendingIndexingBuffers()
+            final long totalIndexingExecutionTimeInNanos = totalIndexingTimeInNanos + writeIndexingBufferTime.sum();
+            final long totalIndexingExecutionTimeSinceShardStartedInNanos = totalIndexingExecutionTimeInNanos
+                - indexingLoadBeforeShardStartedInNanos;
             final double recentIndexingLoadSinceShardStarted = recentIndexMetric.calculateRateSince(
                 currentTimeInNanos,
                 recentIndexMetric.getRate(currentTimeInNanos),
@@ -218,6 +248,7 @@ IndexingStats.Stats stats(
                 isThrottled,
                 TimeUnit.MILLISECONDS.toMillis(currentThrottleMillis),
                 totalIndexingTimeSinceShardStartedInNanos,
+                totalIndexingExecutionTimeSinceShardStartedInNanos,
                 timeSinceShardStartedInNanos,
                 recentIndexingLoadSinceShardStarted,
                 peakIndexingLoad

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

@@ -209,7 +209,11 @@ private boolean writePendingIndexingBuffers() {
             .pollFirst()) {
             // Remove the shard from the set first, so that multiple threads can run writeIndexingBuffer concurrently on the same shard.
             pendingWriteIndexingBufferSet.remove(shard);
+            // Calculate the time taken to write the indexing buffers so it can be accounted for in the index write load
+            long startTime = System.nanoTime();
             shard.writeIndexingBuffer();
+            long took = System.nanoTime() - startTime;
+            shard.addWriteIndexBuffersToIndexThreadsTime(took);
             wrotePendingIndexingBuffer = true;
         }
         return wrotePendingIndexingBuffer;
@@ -258,6 +262,7 @@ private void postOperation(ShardId shardId, Engine.Operation operation, Engine.R
         // be reclaimed rapidly. This has the downside of increasing the latency of _bulk requests though. Lucene does the same thing in
         // DocumentsWriter#postUpdate, flushing a segment because the size limit on the RAM buffer was reached happens on the call to
         // IndexWriter#addDocument.
+
         while (writePendingIndexingBuffers()) {
             // If we just wrote segments, then run the checker again if not already running to check if we released enough memory.
             if (statusChecker.tryRun() == false) {

server/src/test/java/org/elasticsearch/action/admin/cluster/node/stats/NodeStatsTests.java

@@ -593,6 +593,7 @@ private static CommonStats createShardLevelCommonStats() {
             ++iota,
             ++iota,
             ++iota,
+            ++iota,
             ++iota
         );
         indicesCommonStats.getIndexing().add(new IndexingStats(indexingStats));

server/src/test/java/org/elasticsearch/action/datastreams/autosharding/DataStreamAutoShardingServiceTests.java

@@ -1066,6 +1066,7 @@ private static IndexingStats createIndexingStats(double indexingLoad, double rec
                 false,
                 0,
                 totalIndexingTimeSinceShardStartedInNanos,
+                totalIndexingTimeSinceShardStartedInNanos,
                 totalActiveTimeInNanos,
                 recentIndexingLoad,
                 peakIndexingLoad