From 48fb42bc85e93cf33ccdd7a8a79ff8d0e7703bac Mon Sep 17 00:00:00 2001 From: Peter Alfonsi Date: Mon, 9 Oct 2023 09:59:41 -0700 Subject: [PATCH] Adds and tests took time for QuerySearchResult Signed-off-by: Peter Alfonsi --- .../opensearch/search/query/QueryPhase.java | 4 +- .../search/query/QuerySearchResult.java | 11 ++ .../search/query/QueryPhaseTests.java | 116 ++++++++++++++++++ .../search/query/QuerySearchResultTests.java | 2 + 4 files changed, 132 insertions(+), 1 deletion(-) diff --git a/server/src/main/java/org/opensearch/search/query/QueryPhase.java b/server/src/main/java/org/opensearch/search/query/QueryPhase.java index f3cf2c13ecdef..14b239d3270d4 100644 --- a/server/src/main/java/org/opensearch/search/query/QueryPhase.java +++ b/server/src/main/java/org/opensearch/search/query/QueryPhase.java @@ -131,6 +131,7 @@ public void preProcess(SearchContext context) { } public void execute(SearchContext searchContext) throws QueryPhaseExecutionException { + long startTime = System.nanoTime(); if (searchContext.hasOnlySuggest()) { suggestProcessor.process(searchContext); searchContext.queryResult() @@ -138,6 +139,7 @@ public void execute(SearchContext searchContext) throws QueryPhaseExecutionExcep new TopDocsAndMaxScore(new TopDocs(new TotalHits(0, TotalHits.Relation.EQUAL_TO), Lucene.EMPTY_SCORE_DOCS), Float.NaN), new DocValueFormat[0] ); + searchContext.queryResult().setTookTimeNanos(System.nanoTime() - startTime); return; } @@ -165,6 +167,7 @@ public void execute(SearchContext searchContext) throws QueryPhaseExecutionExcep ); searchContext.queryResult().profileResults(shardResults); } + searchContext.queryResult().setTookTimeNanos(System.nanoTime() - startTime); } // making public for testing @@ -292,7 +295,6 @@ static boolean executeInternal(SearchContext searchContext, QueryPhaseSearcher q queryResult.nodeQueueSize(rExecutor.getCurrentQueueSize()); queryResult.serviceTimeEWMA((long) rExecutor.getTaskExecutionEWMA()); } - return shouldRescore; } finally { // Search phase has finished, no longer need to check for timeout diff --git a/server/src/main/java/org/opensearch/search/query/QuerySearchResult.java b/server/src/main/java/org/opensearch/search/query/QuerySearchResult.java index f8a1e99ff585f..cde1fdf381fcb 100644 --- a/server/src/main/java/org/opensearch/search/query/QuerySearchResult.java +++ b/server/src/main/java/org/opensearch/search/query/QuerySearchResult.java @@ -87,6 +87,7 @@ public final class QuerySearchResult extends SearchPhaseResult { private int nodeQueueSize = -1; private final boolean isNull; + private long tookTimeNanos; public QuerySearchResult() { this(false); @@ -364,6 +365,7 @@ public void readFromWithId(ShardSearchContextId id, StreamInput in) throws IOExc nodeQueueSize = in.readInt(); setShardSearchRequest(in.readOptionalWriteable(ShardSearchRequest::new)); setRescoreDocIds(new RescoreDocIds(in)); + tookTimeNanos = in.readVLong(); } @Override @@ -406,6 +408,7 @@ public void writeToNoId(StreamOutput out) throws IOException { out.writeInt(nodeQueueSize); out.writeOptionalWriteable(getShardSearchRequest()); getRescoreDocIds().writeTo(out); + out.writeVLong(tookTimeNanos); // VLong as took time should always be positive } public TotalHits getTotalHits() { @@ -415,4 +418,12 @@ public TotalHits getTotalHits() { public float getMaxScore() { return maxScore; } + + public long getTookTimeNanos() { + return tookTimeNanos; + } + + public void setTookTimeNanos(long tookTime) { + tookTimeNanos = tookTime; + } } diff --git a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java index 39126a607f968..93c54f93801ab 100644 --- a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java +++ b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java @@ -85,9 +85,14 @@ import org.apache.lucene.tests.index.RandomIndexWriter; import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.FixedBitSet; +import org.opensearch.action.OriginalIndices; +import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchShardTask; import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.common.Strings; +import org.opensearch.core.index.Index; +import org.opensearch.core.index.shard.ShardId; import org.opensearch.core.tasks.TaskCancelledException; import org.opensearch.index.mapper.DateFieldMapper; import org.opensearch.index.mapper.MappedFieldType; @@ -103,9 +108,11 @@ import org.opensearch.lucene.queries.MinDocQuery; import org.opensearch.search.DocValueFormat; import org.opensearch.search.collapse.CollapseBuilder; +import org.opensearch.search.internal.AliasFilter; import org.opensearch.search.internal.ContextIndexSearcher; import org.opensearch.search.internal.ScrollContext; import org.opensearch.search.internal.SearchContext; +import org.opensearch.search.internal.ShardSearchRequest; import org.opensearch.search.sort.SortAndFormats; import org.opensearch.test.TestSearchContext; import org.opensearch.threadpool.ThreadPool; @@ -115,6 +122,7 @@ import java.util.Arrays; import java.util.Collection; import java.util.Collections; +import java.util.LinkedList; import java.util.List; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -1145,6 +1153,114 @@ public void testQueryTimeoutChecker() throws Exception { createTimeoutCheckerThenWaitThenRun(timeCacheLifespan / 4, timeCacheLifespan / 2 + timeTolerance, false, true); } + public void testQuerySearchResultTookTime() throws IOException { + int sleepMillis = randomIntBetween(500, 4000); // between 0.5 and 4 sec + DelayedQueryPhaseSearcher delayedQueryPhaseSearcher = new DelayedQueryPhaseSearcher(sleepMillis); + + // we need to test queryPhase.execute(), not executeInternal(), since that's what the timer wraps around + // for that we must set up a searchContext with more functionality than the TestSearchContext, + // which requires a bit of complexity with test classes + + Directory dir = newDirectory(); + final Sort sort = new Sort(new SortField("rank", SortField.Type.INT)); + IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); + RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); + Document doc = new Document(); + for (int i = 0; i < 10; i++) { + doc.add(new StringField("foo", Integer.toString(i), Store.NO)); + } + w.addDocument(doc); + w.close(); + IndexReader reader = DirectoryReader.open(dir); + + QueryShardContext queryShardContext = mock(QueryShardContext.class); + when(queryShardContext.fieldMapper("user")).thenReturn( + new NumberFieldType("user", NumberType.INTEGER, true, false, true, false, null, Collections.emptyMap()) + ); + + Index index = new Index("IndexName", "UUID"); + ShardId shardId = new ShardId(index, 0); + long nowInMillis = System.currentTimeMillis(); + String clusterAlias = randomBoolean() ? null : randomAlphaOfLengthBetween(3, 10); + SearchRequest searchRequest = new SearchRequest(); + searchRequest.allowPartialSearchResults(randomBoolean()); + ShardSearchRequest request = new ShardSearchRequest( + OriginalIndices.NONE, + searchRequest, + shardId, + 1, + AliasFilter.EMPTY, + 1f, + nowInMillis, + clusterAlias, + Strings.EMPTY_ARRAY + ); + TestSearchContextWithRequest searchContext = new TestSearchContextWithRequest( + queryShardContext, + indexShard, + newEarlyTerminationContextSearcher(reader, 0, executor), + request + ); + + QueryPhase queryPhase = new QueryPhase(delayedQueryPhaseSearcher); + queryPhase.execute(searchContext); + long tookTime = searchContext.queryResult().getTookTimeNanos(); + assertTrue(tookTime >= (long) sleepMillis * 1000000); + reader.close(); + dir.close(); + } + + private class TestSearchContextWithRequest extends TestSearchContext { + ShardSearchRequest request; + Query query; + + public TestSearchContextWithRequest( + QueryShardContext queryShardContext, + IndexShard indexShard, + ContextIndexSearcher searcher, + ShardSearchRequest request + ) { + super(queryShardContext, indexShard, searcher); + this.request = request; + this.query = new TermQuery(new Term("foo", "bar")); + } + + @Override + public ShardSearchRequest request() { + return request; + } + + @Override + public Query query() { + return this.query; + } + } + + private class DelayedQueryPhaseSearcher extends QueryPhase.DefaultQueryPhaseSearcher implements QueryPhaseSearcher { + // add delay into searchWith + private final int sleepMillis; + + public DelayedQueryPhaseSearcher(int sleepMillis) { + super(); + this.sleepMillis = sleepMillis; + } + + @Override + public boolean searchWith( + SearchContext searchContext, + ContextIndexSearcher searcher, + Query query, + LinkedList collectors, + boolean hasFilterCollector, + boolean hasTimeout + ) throws IOException { + try { + Thread.sleep(sleepMillis); + } catch (Exception ignored) {} + return super.searchWith(searchContext, searcher, query, collectors, hasFilterCollector, hasTimeout); + } + } + private void createTimeoutCheckerThenWaitThenRun( long timeout, long sleepAfterCreation, diff --git a/server/src/test/java/org/opensearch/search/query/QuerySearchResultTests.java b/server/src/test/java/org/opensearch/search/query/QuerySearchResultTests.java index 41e4e1ae45a73..0d83777dbd4cf 100644 --- a/server/src/test/java/org/opensearch/search/query/QuerySearchResultTests.java +++ b/server/src/test/java/org/opensearch/search/query/QuerySearchResultTests.java @@ -99,6 +99,7 @@ private static QuerySearchResult createTestInstance() throws Exception { if (randomBoolean()) { result.aggregations(InternalAggregationsTests.createTestInstance()); } + assertEquals(0, result.getTookTimeNanos()); return result; } @@ -118,6 +119,7 @@ public void testSerialization() throws Exception { assertEquals(aggs.asList(), deserializedAggs.asList()); } assertEquals(querySearchResult.terminatedEarly(), deserialized.terminatedEarly()); + assertEquals(querySearchResult.getTookTimeNanos(), deserialized.getTookTimeNanos()); } public void testNullResponse() throws Exception {