add index execution time tracking

jdconrad · Aug 14, 2024 · ec9387c · ec9387c
1 parent d3fdb36
commit ec9387c
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 10 deletions.
diff --git a/...g/elasticsearch/common/util/concurrent/TaskExecutionTimeTrackingEsThreadPoolExecutor.java b/...g/elasticsearch/common/util/concurrent/TaskExecutionTimeTrackingEsThreadPoolExecutor.java
@@ -136,6 +136,10 @@ protected void appendThreadPoolExecutorDetails(StringBuilder sb) {
             .append(", ");
     }
 
+    public boolean trackOngoingTasks() {
+        return trackOngoingTasks;
+    }
+
     /**
      * Returns the set of currently running tasks and their start timestamp.
      * <p>

diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java
@@ -45,6 +45,7 @@
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.common.util.concurrent.EsRejectedExecutionException;
+import org.elasticsearch.common.util.concurrent.TaskExecutionTimeTrackingEsThreadPoolExecutor;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.Releasable;
@@ -141,13 +142,15 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.LongAdder;
 import java.util.function.LongSupplier;
 import java.util.function.Supplier;
 
@@ -314,6 +317,9 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
 
     private final Tracer tracer;
 
+    private final boolean trackIndexExecutionTime;
+    private final ConcurrentHashMap<String, LongAdder> indexExecutionTimes = new ConcurrentHashMap<>();
+
     public SearchService(
         ClusterService clusterService,
         IndicesService indicesService,
@@ -382,6 +388,10 @@ public SearchService(
         enableQueryPhaseParallelCollection = QUERY_PHASE_PARALLEL_COLLECTION_ENABLED.get(settings);
         clusterService.getClusterSettings()
             .addSettingsUpdateConsumer(QUERY_PHASE_PARALLEL_COLLECTION_ENABLED, this::setEnableQueryPhaseParallelCollection);
+
+        Executor searchExecutor = threadPool.executor(Names.SEARCH);
+        trackIndexExecutionTime = searchExecutor instanceof TaskExecutionTimeTrackingEsThreadPoolExecutor
+            && ((TaskExecutionTimeTrackingEsThreadPoolExecutor) searchExecutor).trackOngoingTasks();
     }
 
     private void setEnableSearchWorkerThreads(boolean enableSearchWorkerThreads) {
@@ -570,7 +580,7 @@ private <T extends RefCounted> void ensureAfterSeqNoRefreshed(
         final Executor executor = getExecutor(shard);
         try {
             if (waitForCheckpoint <= UNASSIGNED_SEQ_NO) {
-                runAsync(executor, executable, listener);
+                runAsync(executor, executable, listener, shard);
                 return;
             }
             if (shard.indexSettings().getRefreshInterval().getMillis() <= 0) {
@@ -642,7 +652,7 @@ private void searchReady() {
                         if (timeoutTask != null) {
                             timeoutTask.cancel();
                         }
-                        runAsync(executor, executable, listener);
+                        runAsync(executor, executable, listener, shard);
                     }
                 }
             });
@@ -664,12 +674,30 @@ private IndexShard getShard(ShardSearchRequest request) {
         return indicesService.indexServiceSafe(request.shardId().getIndex()).getShard(request.shardId().id());
     }
 
-    private static <T extends RefCounted> void runAsync(
+    public Map<String, LongAdder> getIndexExecutionTimes() {
+        return indexExecutionTimes;
+    }
+
+    private <T extends RefCounted> void runAsync(
         Executor executor,
         CheckedSupplier<T, Exception> executable,
-        ActionListener<T> listener
+        ActionListener<T> listener,
+        IndexShard indexShard
     ) {
-        executor.execute(ActionRunnable.supplyAndDecRef(listener, executable));
+        if (trackIndexExecutionTime) {
+            String indexName = indexShard.shardId().getIndexName();
+            indexExecutionTimes.putIfAbsent(indexName, new LongAdder());
+            LongAdder indexExecutionTime = indexExecutionTimes.get(indexName);
+            executor.execute(ActionRunnable.supplyAndDecRef(listener, () -> {
+                long start = System.nanoTime();
+                var result = executable.get();
+                long total = System.nanoTime() - start;
+                indexExecutionTime.add(total);
+                return result;
+            }));
+        } else {
+            executor.execute(ActionRunnable.supplyAndDecRef(listener, executable));
+        }
     }
 
     /**
@@ -745,7 +773,7 @@ public void executeRankFeaturePhase(RankFeatureShardRequest request, SearchShard
                 // we handle the failure in the failure listener below
                 throw e;
             }
-        }, wrapFailureListener(listener, readerContext, markAsUsed));
+        }, wrapFailureListener(listener, readerContext, markAsUsed), readerContext.indexShard());
     }
 
     private QueryFetchSearchResult executeFetchPhase(ReaderContext reader, SearchContext context, long afterQueryTime) {
@@ -795,7 +823,7 @@ public void executeQueryPhase(
                 // we handle the failure in the failure listener below
                 throw e;
             }
-        }, wrapFailureListener(listener, readerContext, markAsUsed));
+        }, wrapFailureListener(listener, readerContext, markAsUsed), readerContext.indexShard());
     }
 
     /**
@@ -838,7 +866,7 @@ public void executeQueryPhase(QuerySearchRequest request, SearchShardTask task,
                     // we handle the failure in the failure listener below
                     throw e;
                 }
-            }, wrapFailureListener(l, readerContext, markAsUsed));
+            }, wrapFailureListener(l, readerContext, markAsUsed), readerContext.indexShard());
         }));
     }
 
@@ -889,7 +917,7 @@ public void executeFetchPhase(
                 // we handle the failure in the failure listener below
                 throw e;
             }
-        }, wrapFailureListener(listener, readerContext, markAsUsed));
+        }, wrapFailureListener(listener, readerContext, markAsUsed), readerContext.indexShard());
     }
 
     public void executeFetchPhase(ShardFetchRequest request, SearchShardTask task, ActionListener<FetchSearchResult> listener) {
@@ -921,7 +949,7 @@ public void executeFetchPhase(ShardFetchRequest request, SearchShardTask task, A
                 // we handle the failure in the failure listener below
                 throw e;
             }
-        }, wrapFailureListener(listener, readerContext, markAsUsed));
+        }, wrapFailureListener(listener, readerContext, markAsUsed), readerContext.indexShard());
     }
 
     protected void checkCancelled(SearchShardTask task) {