Add slice level operation listener methods

Signed-off-by: Jay Deng <jayd0104@gmail.com>
jed326 · Aug 6, 2024 · 3bcdfa1 · 3bcdfa1
1 parent f980924
commit 3bcdfa1
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git a/server/src/main/java/org/opensearch/index/shard/SearchOperationListener.java b/server/src/main/java/org/opensearch/index/shard/SearchOperationListener.java
@@ -71,6 +71,10 @@ default void onFailedQueryPhase(SearchContext searchContext) {}
      */
     default void onQueryPhase(SearchContext searchContext, long tookInNanos) {}
 
+    default void onPreSlicePhase(SearchContext searchContext) {}
+
+    default void onPostSlicePhase(SearchContext searchContext) {}
+
     /**
      * Executed before the fetch phase is executed
      * @param searchContext the current search context
@@ -195,6 +199,28 @@ public void onQueryPhase(SearchContext searchContext, long tookInNanos) {
             }
         }
 
+        @Override
+        public void onPreSlicePhase(SearchContext searchContext) {
+            for (SearchOperationListener listener : listeners) {
+                try {
+                    listener.onPreSlicePhase(searchContext);
+                } catch (Exception e) {
+                    logger.warn(() -> new ParameterizedMessage("onPreSlicePhase listener [{}] failed", listener), e);
+                }
+            }
+        }
+
+        @Override
+        public void onPostSlicePhase(SearchContext searchContext) {
+            for (SearchOperationListener listener : listeners) {
+                try {
+                    listener.onPostSlicePhase(searchContext);
+                } catch (Exception e) {
+                    logger.warn(() -> new ParameterizedMessage("onPostSlicePhase listener [{}] failed", listener), e);
+                }
+            }
+        }
+
         @Override
         public void onPreFetchPhase(SearchContext searchContext) {
             for (SearchOperationListener listener : listeners) {

diff --git a/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java b/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java
@@ -270,6 +270,7 @@ public void search(
 
     @Override
     protected void search(List<LeafReaderContext> leaves, Weight weight, Collector collector) throws IOException {
+        searchContext.indexShard().getSearchOperationListener().onPreSlicePhase(searchContext);
         // Time series based workload by default traverses segments in desc order i.e. latest to the oldest order.
         // This is actually beneficial for search queries to start search on latest segments first for time series workload.
         // That can slow down ASC order queries on timestamp workload. So to avoid that slowdown, we will reverse leaf
@@ -284,6 +285,7 @@ protected void search(List<LeafReaderContext> leaves, Weight weight, Collector c
             }
         }
         searchContext.bucketCollectorProcessor().processPostCollection(collector);
+        searchContext.indexShard().getSearchOperationListener().onPostSlicePhase(searchContext);
     }
 
     /**