use of org.opensearch.search.asynchronous.task.AsynchronousSearchTask in project asynchronous-search by opensearch-project.
the class AsynchronousSearchServiceTests method testUpdateExpirationTimesOutBlockedOnPersistence.
public void testUpdateExpirationTimesOutBlockedOnPersistence() throws InterruptedException {
DiscoveryNode discoveryNode = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
ThreadPool testThreadPool = null;
try {
testThreadPool = new TestThreadPool(AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME, executorBuilder);
ClusterService mockClusterService = ClusterServiceUtils.createClusterService(testThreadPool, discoveryNode, clusterSettings);
FakeClient fakeClient = new FakeClient(testThreadPool);
AsynchronousSearchActiveStore asActiveStore = new AsynchronousSearchActiveStore(mockClusterService);
AsynchronousSearchPersistenceService persistenceService = new AsynchronousSearchPersistenceService(fakeClient, mockClusterService, testThreadPool);
AsynchronousSearchService asService = new AsynchronousSearchService(persistenceService, asActiveStore, fakeClient, mockClusterService, testThreadPool, new InternalAsynchronousSearchStats(), new NamedWriteableRegistry(emptyList()));
TimeValue keepAlive = timeValueHours(9);
boolean keepOnCompletion = true;
SearchRequest searchRequest = new SearchRequest();
SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(searchRequest);
submitAsynchronousSearchRequest.keepOnCompletion(keepOnCompletion);
submitAsynchronousSearchRequest.keepAlive(keepAlive);
AsynchronousSearchContext context = asService.createAndStoreContext(submitAsynchronousSearchRequest, System.currentTimeMillis(), () -> null, null);
assertTrue(context instanceof AsynchronousSearchActiveContext);
AsynchronousSearchActiveContext asActiveContext = (AsynchronousSearchActiveContext) context;
assertNull(asActiveContext.getTask());
assertNull(asActiveContext.getAsynchronousSearchId());
assertEquals(asActiveContext.getAsynchronousSearchState(), AsynchronousSearchState.INIT);
// bootstrap search
AsynchronousSearchTask task = new AsynchronousSearchTask(randomNonNegativeLong(), "transport", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap(), (AsynchronousSearchActiveContext) context, null, (c) -> {
});
asService.bootstrapSearch(task, context.getContextId());
assertEquals(asActiveContext.getTask(), task);
assertEquals(asActiveContext.getStartTimeMillis(), task.getStartTime());
long originalExpirationTimeMillis = asActiveContext.getExpirationTimeMillis();
assertEquals(originalExpirationTimeMillis, task.getStartTime() + keepAlive.millis());
assertEquals(asActiveContext.getAsynchronousSearchState(), AsynchronousSearchState.RUNNING);
blockPersistence = true;
context.getAsynchronousSearchProgressListener().onResponse(getMockSearchResponse());
CountDownLatch updateLatch = new CountDownLatch(1);
TimeValue newKeepAlive = timeValueHours(10);
fakeClient.awaitBlock();
asService.updateKeepAliveAndGetContext(asActiveContext.getAsynchronousSearchId(), newKeepAlive, asActiveContext.getContextId(), null, new LatchedActionListener<>(wrap(r -> fail("expected update req to timeout"), e -> assertTrue("expected timeout got " + e.getClass(), e instanceof OpenSearchTimeoutException)), updateLatch));
updateLatch.await();
fakeClient.releaseBlock();
} finally {
ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
}
}
use of org.opensearch.search.asynchronous.task.AsynchronousSearchTask in project asynchronous-search by opensearch-project.
the class AsynchronousSearchServiceTests method testUpdateExpirationOnRunningSearch.
public void testUpdateExpirationOnRunningSearch() throws InterruptedException {
DiscoveryNode discoveryNode = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
ThreadPool testThreadPool = null;
try {
testThreadPool = new TestThreadPool(AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME, executorBuilder);
ClusterService mockClusterService = ClusterServiceUtils.createClusterService(testThreadPool, discoveryNode, clusterSettings);
FakeClient fakeClient = new FakeClient(testThreadPool);
AsynchronousSearchActiveStore asActiveStore = new AsynchronousSearchActiveStore(mockClusterService);
AsynchronousSearchPersistenceService persistenceService = new AsynchronousSearchPersistenceService(fakeClient, mockClusterService, testThreadPool);
AsynchronousSearchService asService = new AsynchronousSearchService(persistenceService, asActiveStore, fakeClient, mockClusterService, testThreadPool, new InternalAsynchronousSearchStats(), new NamedWriteableRegistry(emptyList()));
TimeValue keepAlive = timeValueHours(9);
boolean keepOnCompletion = false;
SearchRequest searchRequest = new SearchRequest();
SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(searchRequest);
submitAsynchronousSearchRequest.keepOnCompletion(keepOnCompletion);
submitAsynchronousSearchRequest.keepAlive(keepAlive);
AsynchronousSearchContext context = asService.createAndStoreContext(submitAsynchronousSearchRequest, System.currentTimeMillis(), () -> null, null);
assertTrue(context instanceof AsynchronousSearchActiveContext);
AsynchronousSearchActiveContext asActiveContext = (AsynchronousSearchActiveContext) context;
assertNull(asActiveContext.getTask());
assertNull(asActiveContext.getAsynchronousSearchId());
assertEquals(asActiveContext.getAsynchronousSearchState(), AsynchronousSearchState.INIT);
// bootstrap search
AsynchronousSearchTask task = new AsynchronousSearchTask(randomNonNegativeLong(), "transport", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap(), (AsynchronousSearchActiveContext) context, null, (c) -> {
});
asService.bootstrapSearch(task, context.getContextId());
assertEquals(asActiveContext.getTask(), task);
assertEquals(asActiveContext.getStartTimeMillis(), task.getStartTime());
long originalExpirationTimeMillis = asActiveContext.getExpirationTimeMillis();
assertEquals(originalExpirationTimeMillis, task.getStartTime() + keepAlive.millis());
assertEquals(asActiveContext.getAsynchronousSearchState(), AsynchronousSearchState.RUNNING);
CountDownLatch findContextLatch = new CountDownLatch(1);
asService.findContext(asActiveContext.getAsynchronousSearchId(), asActiveContext.getContextId(), null, new LatchedActionListener<>(wrap(r -> {
assertTrue(r instanceof AsynchronousSearchActiveContext);
assertEquals(r, context);
}, e -> fail("Find context shouldn't have failed")), findContextLatch));
findContextLatch.await();
CountDownLatch updateLatch = new CountDownLatch(1);
TimeValue newKeepAlive = timeValueHours(10);
asService.updateKeepAliveAndGetContext(asActiveContext.getAsynchronousSearchId(), newKeepAlive, asActiveContext.getContextId(), null, new LatchedActionListener<>(wrap(r -> {
assertTrue(r instanceof AsynchronousSearchActiveContext);
assertThat(r.getExpirationTimeMillis(), greaterThan(originalExpirationTimeMillis));
}, e -> fail()), updateLatch));
updateLatch.await();
} finally {
ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
}
}
use of org.opensearch.search.asynchronous.task.AsynchronousSearchTask in project asynchronous-search by opensearch-project.
the class AsynchronousSearchServiceTests method testFreeActiveContext.
public void testFreeActiveContext() throws InterruptedException {
DiscoveryNode discoveryNode = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
ThreadPool testThreadPool = null;
try {
testThreadPool = new TestThreadPool(AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME, executorBuilder);
ClusterService mockClusterService = ClusterServiceUtils.createClusterService(testThreadPool, discoveryNode, clusterSettings);
FakeClient fakeClient = new FakeClient(testThreadPool);
AsynchronousSearchActiveStore asActiveStore = new AsynchronousSearchActiveStore(mockClusterService);
AsynchronousSearchPersistenceService persistenceService = new AsynchronousSearchPersistenceService(fakeClient, mockClusterService, testThreadPool);
AsynchronousSearchService asService = new AsynchronousSearchService(persistenceService, asActiveStore, fakeClient, mockClusterService, testThreadPool, new InternalAsynchronousSearchStats(), new NamedWriteableRegistry(emptyList()));
TimeValue keepAlive = timeValueHours(9);
boolean keepOnCompletion = randomBoolean();
User user1 = TestClientUtils.randomUser();
SearchRequest searchRequest = new SearchRequest();
SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(searchRequest);
submitAsynchronousSearchRequest.keepOnCompletion(keepOnCompletion);
submitAsynchronousSearchRequest.keepAlive(keepAlive);
AsynchronousSearchContext context = asService.createAndStoreContext(submitAsynchronousSearchRequest, System.currentTimeMillis(), () -> null, user1);
assertTrue(context instanceof AsynchronousSearchActiveContext);
AsynchronousSearchActiveContext asActiveContext = (AsynchronousSearchActiveContext) context;
assertNull(asActiveContext.getTask());
assertNull(asActiveContext.getAsynchronousSearchId());
assertEquals(asActiveContext.getAsynchronousSearchState(), AsynchronousSearchState.INIT);
assertEquals(asActiveContext.getUser(), user1);
// bootstrap search
AsynchronousSearchTask task = new AsynchronousSearchTask(randomNonNegativeLong(), "transport", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap(), (AsynchronousSearchActiveContext) context, null, (c) -> {
}) {
@Override
public boolean isCancelled() {
return true;
}
};
asService.bootstrapSearch(task, context.getContextId());
assertEquals(asActiveContext.getTask(), task);
assertEquals(asActiveContext.getStartTimeMillis(), task.getStartTime());
assertEquals(asActiveContext.getExpirationTimeMillis(), task.getStartTime() + keepAlive.millis());
assertEquals(asActiveContext.getAsynchronousSearchState(), AsynchronousSearchState.RUNNING);
assertTrue(asService.freeActiveContext(asActiveContext));
assertTrue(asService.getAllActiveContexts().isEmpty());
CountDownLatch latch = new CountDownLatch(1);
asService.freeContext(context.getAsynchronousSearchId(), context.getContextId(), user1, new LatchedActionListener<>(wrap(r -> fail(), e -> assertTrue(e instanceof ResourceNotFoundException)), latch));
latch.await();
} finally {
ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
}
}
use of org.opensearch.search.asynchronous.task.AsynchronousSearchTask in project asynchronous-search by opensearch-project.
the class AsynchronousSearchServiceTests method testBootStrapOnClosedSearch.
public void testBootStrapOnClosedSearch() {
DiscoveryNode discoveryNode = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
ThreadPool testThreadPool = null;
try {
testThreadPool = new TestThreadPool(AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME, executorBuilder);
ClusterService mockClusterService = ClusterServiceUtils.createClusterService(testThreadPool, discoveryNode, clusterSettings);
FakeClient fakeClient = new FakeClient(testThreadPool);
AsynchronousSearchActiveStore asActiveStore = new AsynchronousSearchActiveStore(mockClusterService);
AsynchronousSearchPersistenceService persistenceService = new AsynchronousSearchPersistenceService(fakeClient, mockClusterService, testThreadPool);
AsynchronousSearchService asService = new AsynchronousSearchService(persistenceService, asActiveStore, fakeClient, mockClusterService, testThreadPool, new InternalAsynchronousSearchStats(), new NamedWriteableRegistry(emptyList()));
TimeValue keepAlive = timeValueHours(9);
boolean keepOnCompletion = false;
SearchRequest searchRequest = new SearchRequest();
SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(searchRequest);
submitAsynchronousSearchRequest.keepOnCompletion(keepOnCompletion);
submitAsynchronousSearchRequest.keepAlive(keepAlive);
AsynchronousSearchContext context = asService.createAndStoreContext(submitAsynchronousSearchRequest, System.currentTimeMillis(), () -> null, null);
assertTrue(context instanceof AsynchronousSearchActiveContext);
AsynchronousSearchActiveContext asActiveContext = (AsynchronousSearchActiveContext) context;
assertNull(asActiveContext.getTask());
assertNull(asActiveContext.getAsynchronousSearchId());
assertEquals(asActiveContext.getAsynchronousSearchState(), AsynchronousSearchState.INIT);
// close context
asActiveContext.close();
// bootstrap search
AsynchronousSearchTask task = new AsynchronousSearchTask(randomNonNegativeLong(), "transport", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap(), (AsynchronousSearchActiveContext) context, null, (c) -> {
});
asActiveContext.setState(AsynchronousSearchState.CLOSED);
expectThrows(IllegalStateException.class, () -> asService.bootstrapSearch(task, context.getContextId()));
} finally {
ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
}
}
use of org.opensearch.search.asynchronous.task.AsynchronousSearchTask in project asynchronous-search by opensearch-project.
the class AsynchronousSearchServiceUpdateContextTests method testUpdateContextWhenContextCloseAndKeepOnCompletionTrue.
public void testUpdateContextWhenContextCloseAndKeepOnCompletionTrue() throws InterruptedException, IOException {
DiscoveryNode discoveryNode = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
ThreadPool testThreadPool = null;
try {
testThreadPool = new TestThreadPool(OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME, executorBuilder);
ClusterService mockClusterService = getClusterService(discoveryNode, testThreadPool);
MockClient mockClient = new MockClient(testThreadPool);
AsynchronousSearchActiveStore mockStore = mock(AsynchronousSearchActiveStore.class);
AsynchronousSearchPersistenceService persistenceService = new AsynchronousSearchPersistenceService(mockClient, mockClusterService, testThreadPool);
AsynchronousSearchService asService = new AsynchronousSearchService(persistenceService, mockStore, mockClient, mockClusterService, testThreadPool, new InternalAsynchronousSearchStats(), new NamedWriteableRegistry(emptyList()));
TimeValue keepAlive = timeValueHours(9);
boolean keepOnCompletion = true;
User user1 = randomBoolean() ? randomUser() : null;
SearchRequest searchRequest = new SearchRequest();
SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(searchRequest);
submitAsynchronousSearchRequest.keepOnCompletion(keepOnCompletion);
submitAsynchronousSearchRequest.keepAlive(keepAlive);
AsynchronousSearchProgressListener asProgressListener = mockAsynchronousSearchProgressListener(testThreadPool);
AsynchronousSearchContextId asContextId = new AsynchronousSearchContextId(UUID.randomUUID().toString(), randomNonNegativeLong());
MockAsynchronousSearchActiveContext asActiveContext = new MockAsynchronousSearchActiveContext(asContextId, discoveryNode.getId(), keepAlive, true, testThreadPool, testThreadPool::absoluteTimeInMillis, asProgressListener, user1);
// bootstrap search
AsynchronousSearchTask task = new AsynchronousSearchTask(randomNonNegativeLong(), "transport", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap(), asActiveContext, null, (c) -> {
});
asActiveContext.setTask(task);
long oldExpirationTimeMillis = asActiveContext.getExpirationTimeMillis();
simulateIsAlive = false;
simulateTimedOut = false;
CountDownLatch updateLatch = new CountDownLatch(1);
when(mockStore.getContext(asContextId)).thenReturn(Optional.of(asActiveContext));
asService.updateKeepAliveAndGetContext(asActiveContext.getAsynchronousSearchId(), keepAlive, asActiveContext.getContextId(), user1, new LatchedActionListener<>(wrap(r -> {
assertTrue(r instanceof AsynchronousSearchPersistenceContext);
// assert active context expiration time is not updated
assertEquals(asActiveContext.getExpirationTimeMillis(), oldExpirationTimeMillis);
}, e -> {
fail("expected successful update got " + e.getMessage());
}), updateLatch));
updateLatch.await();
mockClusterService.stop();
} finally {
ThreadPool.terminate(testThreadPool, 30, TimeUnit.SECONDS);
}
}
Aggregations