|
85 | 85 | import org.apache.lucene.util.FixedBitSet;
|
86 | 86 | import org.opensearch.action.search.SearchShardTask;
|
87 | 87 | import org.opensearch.common.settings.Settings;
|
| 88 | +import org.opensearch.common.unit.TimeValue; |
88 | 89 | import org.opensearch.index.mapper.DateFieldMapper;
|
89 | 90 | import org.opensearch.index.mapper.MappedFieldType;
|
90 | 91 | import org.opensearch.index.mapper.MapperService;
|
|
105 | 106 | import org.opensearch.search.sort.SortAndFormats;
|
106 | 107 | import org.opensearch.tasks.TaskCancelledException;
|
107 | 108 | import org.opensearch.test.TestSearchContext;
|
| 109 | +import org.opensearch.threadpool.ThreadPool; |
108 | 110 |
|
109 | 111 | import java.io.IOException;
|
110 | 112 | import java.util.ArrayList;
|
|
117 | 119 | import static org.hamcrest.Matchers.greaterThanOrEqualTo;
|
118 | 120 | import static org.hamcrest.Matchers.instanceOf;
|
119 | 121 | import static org.hamcrest.Matchers.lessThanOrEqualTo;
|
| 122 | +import static org.mockito.ArgumentMatchers.eq; |
120 | 123 | import static org.mockito.Mockito.mock;
|
121 | 124 | import static org.mockito.Mockito.spy;
|
122 | 125 | import static org.mockito.Mockito.when;
|
| 126 | +import static org.mockito.Mockito.verify; |
| 127 | +import static org.mockito.Mockito.times; |
| 128 | +import static org.mockito.Mockito.atLeastOnce; |
| 129 | +import static org.mockito.Mockito.verifyNoMoreInteractions; |
123 | 130 | import static org.opensearch.search.query.TopDocsCollectorContext.hasInfMaxScore;
|
124 | 131 |
|
125 | 132 | public class QueryPhaseTests extends IndexShardTestCase {
|
@@ -1079,6 +1086,58 @@ public void testCancellationDuringPreprocess() throws IOException {
|
1079 | 1086 | }
|
1080 | 1087 | }
|
1081 | 1088 |
|
| 1089 | + public void testQueryTimeoutChecker() throws Exception { |
| 1090 | + long timeCacheLifespan = ThreadPool.ESTIMATED_TIME_INTERVAL_SETTING.get(Settings.EMPTY).millis(); |
| 1091 | + long timeTolerance = timeCacheLifespan / 20; |
| 1092 | + |
| 1093 | + // should throw time exceed exception for sure after timeCacheLifespan*2+timeTolerance (next's next cached time is available) |
| 1094 | + assertThrows( |
| 1095 | + QueryPhase.TimeExceededException.class, |
| 1096 | + () -> createTimeoutCheckerThenWaitThenRun(timeCacheLifespan, timeCacheLifespan * 2 + timeTolerance, true, false) |
| 1097 | + ); |
| 1098 | + |
| 1099 | + // should not throw time exceed exception after timeCacheLifespan+timeTolerance because new cached time - init time < timeout |
| 1100 | + createTimeoutCheckerThenWaitThenRun(timeCacheLifespan, timeCacheLifespan + timeTolerance, true, false); |
| 1101 | + |
| 1102 | + // should not throw time exceed exception after timeout < timeCacheLifespan when cached time didn't change |
| 1103 | + createTimeoutCheckerThenWaitThenRun(timeCacheLifespan / 2, timeCacheLifespan / 2 + timeTolerance, false, true); |
| 1104 | + createTimeoutCheckerThenWaitThenRun(timeCacheLifespan / 4, timeCacheLifespan / 2 + timeTolerance, false, true); |
| 1105 | + } |
| 1106 | + |
| 1107 | + private void createTimeoutCheckerThenWaitThenRun( |
| 1108 | + long timeout, |
| 1109 | + long sleepAfterCreation, |
| 1110 | + boolean checkCachedTimeChanged, |
| 1111 | + boolean checkCachedTimeHasNotChanged |
| 1112 | + ) throws Exception { |
| 1113 | + long timeCacheLifespan = ThreadPool.ESTIMATED_TIME_INTERVAL_SETTING.get(Settings.EMPTY).millis(); |
| 1114 | + long timeTolerance = timeCacheLifespan / 20; |
| 1115 | + long currentTimeDiffWithCachedTime = TimeValue.nsecToMSec(System.nanoTime()) - threadPool.relativeTimeInMillis(); |
| 1116 | + // need to run this test approximately at the start of cached time window |
| 1117 | + long timeToAlignTimeWithCachedTimeOffset = timeCacheLifespan - currentTimeDiffWithCachedTime + timeTolerance; |
| 1118 | + Thread.sleep(timeToAlignTimeWithCachedTimeOffset); |
| 1119 | + |
| 1120 | + long initialRelativeCachedTime = threadPool.relativeTimeInMillis(); |
| 1121 | + SearchContext mockedSearchContext = mock(SearchContext.class); |
| 1122 | + when(mockedSearchContext.timeout()).thenReturn(TimeValue.timeValueMillis(timeout)); |
| 1123 | + when(mockedSearchContext.getRelativeTimeInMillis()).thenAnswer(invocation -> threadPool.relativeTimeInMillis()); |
| 1124 | + when(mockedSearchContext.getRelativeTimeInMillis(eq(false))).thenCallRealMethod(); |
| 1125 | + Runnable queryTimeoutChecker = QueryPhase.createQueryTimeoutChecker(mockedSearchContext); |
| 1126 | + // make sure next time slot become available |
| 1127 | + Thread.sleep(sleepAfterCreation); |
| 1128 | + if (checkCachedTimeChanged) { |
| 1129 | + assertNotEquals(initialRelativeCachedTime, threadPool.relativeTimeInMillis()); |
| 1130 | + } |
| 1131 | + if (checkCachedTimeHasNotChanged) { |
| 1132 | + assertEquals(initialRelativeCachedTime, threadPool.relativeTimeInMillis()); |
| 1133 | + } |
| 1134 | + queryTimeoutChecker.run(); |
| 1135 | + verify(mockedSearchContext, times(1)).timeout(); |
| 1136 | + verify(mockedSearchContext, times(1)).getRelativeTimeInMillis(eq(false)); |
| 1137 | + verify(mockedSearchContext, atLeastOnce()).getRelativeTimeInMillis(); |
| 1138 | + verifyNoMoreInteractions(mockedSearchContext); |
| 1139 | + } |
| 1140 | + |
1082 | 1141 | private static class TestSearchContextWithRewriteAndCancellation extends TestSearchContext {
|
1083 | 1142 |
|
1084 | 1143 | private TestSearchContextWithRewriteAndCancellation(
|
|
0 commit comments