use of org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext in project asynchronous-search by opensearch-project.
the class AsynchronousSearchServiceFreeContextTests method testFreeActiveContextWithCancelledTask.
public void testFreeActiveContextWithCancelledTask() throws InterruptedException {
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 = timeValueDays(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());
AsynchronousSearchActiveContext asActiveContext = new AsynchronousSearchActiveContext(asContextId, discoveryNode.getId(), keepAlive, true, testThreadPool, testThreadPool::absoluteTimeInMillis, asProgressListener, user1, () -> true);
// bootstrap search
AsynchronousSearchTask task = new AsynchronousSearchTask(randomNonNegativeLong(), "transport", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap(), asActiveContext, null, (c) -> {
}) {
@Override
public boolean isCancelled() {
return true;
}
};
asActiveContext.setTask(task);
asActiveContext.setState(AsynchronousSearchState.RUNNING);
when(mockStore.getContext(any())).thenReturn(Optional.of(asActiveContext));
persisted = false;
// task cancellation fails
CountDownLatch latch = new CountDownLatch(1);
// user
User user = randomBoolean() ? null : user1;
asService.freeContext(asActiveContext.getAsynchronousSearchId(), asActiveContext.getContextId(), user, new LatchedActionListener<>(ActionListener.wrap(Assert::assertTrue, e -> {
fail("active context should have been deleted");
}), latch));
latch.await();
mockClusterService.stop();
} finally {
ThreadPool.terminate(testThreadPool, 30, TimeUnit.SECONDS);
}
}
use of org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext in project asynchronous-search by opensearch-project.
the class AsynchronousSearchServiceTests method testFreeContext.
public void testFreeContext() 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);
CountDownLatch latch = new CountDownLatch(1);
asService.freeContext(context.getAsynchronousSearchId(), context.getContextId(), user1, new LatchedActionListener<>(wrap(Assert::assertTrue, e -> fail()), latch));
latch.await();
} finally {
ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
}
}
use of org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext in project asynchronous-search by opensearch-project.
the class AsynchronousSearchServiceTests method testFindContext.
public void testFindContext() 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 = TestUtils.createClusterService(settings, 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();
User user2 = 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) -> {
});
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);
CountDownLatch findContextLatch = new CountDownLatch(3);
ActionListener<AsynchronousSearchContext> expectedSuccessfulActive = new LatchedActionListener<>(wrap(r -> {
assertTrue(r instanceof AsynchronousSearchActiveContext);
assertEquals(r, context);
}, e -> fail("Find context shouldn't have failed. " + e.getMessage())), findContextLatch);
ActionListener<AsynchronousSearchContext> expectedSecurityException = new LatchedActionListener<>(wrap(r -> fail("Expecting security exception"), e -> assertTrue(e instanceof ResourceNotFoundException)), findContextLatch);
asService.findContext(asActiveContext.getAsynchronousSearchId(), asActiveContext.getContextId(), user1, expectedSuccessfulActive);
asService.findContext(asActiveContext.getAsynchronousSearchId(), asActiveContext.getContextId(), user2, expectedSecurityException);
asService.findContext(asActiveContext.getAsynchronousSearchId(), asActiveContext.getContextId(), null, expectedSuccessfulActive);
findContextLatch.await();
AsynchronousSearchProgressListener asProgressListener = asActiveContext.getAsynchronousSearchProgressListener();
boolean success = randomBoolean();
if (success) {
// successful search response
asProgressListener.onResponse(getMockSearchResponse());
} else {
// exception occurred in search
asProgressListener.onFailure(new RuntimeException("test"));
}
waitUntil(() -> asService.getAllActiveContexts().isEmpty());
if (keepOnCompletion) {
// persist to disk
assertEquals(1, fakeClient.persistenceCount.intValue());
} else {
assertEquals(fakeClient.persistenceCount, Integer.valueOf(0));
CountDownLatch freeContextLatch = new CountDownLatch(1);
asService.findContext(context.getAsynchronousSearchId(), context.getContextId(), null, new LatchedActionListener<>(wrap(r -> fail("No context should have been found but found " + asService.getAllActiveContexts().size()), e -> assertTrue(e instanceof ResourceNotFoundException)), freeContextLatch));
freeContextLatch.await();
}
} finally {
ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
}
}
use of org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext in project asynchronous-search by opensearch-project.
the class AsynchronousSearchPostProcessorTests method testProcessSearchResponseForExpiredContext.
public void testProcessSearchResponseForExpiredContext() throws AsynchronousSearchStateMachineClosedException {
DiscoveryNode discoveryNode = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), Collections.emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
AtomicBoolean activeContextCleanUpConsumerInvocation = new AtomicBoolean();
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()));
AsynchronousSearchStateMachine asStateMachine = asService.getStateMachine();
ClusterService clusterService = ClusterServiceUtils.createClusterService(testThreadPool, discoveryNode, clusterSettings);
AsynchronousSearchPostProcessor postProcessor = new AsynchronousSearchPostProcessor(persistenceService, asActiveStore, asStateMachine, (context) -> activeContextCleanUpConsumerInvocation.compareAndSet(false, true), testThreadPool, clusterService);
SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(new SearchRequest());
submitAsynchronousSearchRequest.keepOnCompletion(true);
submitAsynchronousSearchRequest.keepAlive(TimeValue.timeValueMillis(1));
AsynchronousSearchActiveContext context = (AsynchronousSearchActiveContext) asService.createAndStoreContext(submitAsynchronousSearchRequest, System.currentTimeMillis(), () -> InternalAggregationTestCase.emptyReduceContextBuilder(), null);
asStateMachine.trigger(new SearchStartedEvent(context, new SearchTask(0, "n/a", "n/a", () -> "test", null, Collections.emptyMap())));
SearchResponse mockSearchResponse = getMockSearchResponse();
waitUntil(() -> context.isExpired());
AsynchronousSearchResponse asResponse = postProcessor.processSearchResponse(mockSearchResponse, context.getContextId());
assertNotNull(asResponse.getId());
assertNull(asResponse.getError());
assertEquals(AsynchronousSearchState.SUCCEEDED, asResponse.getState());
AsynchronousSearchAssertions.assertSearchResponses(mockSearchResponse, asResponse.getSearchResponse());
assertTrue(activeContextCleanUpConsumerInvocation.get());
assertEquals(0, fakeClient.persistenceCount);
} catch (InterruptedException e) {
fail("Interrupted exception" + e.getMessage());
} finally {
ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
}
}
use of org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext in project asynchronous-search by opensearch-project.
the class AsynchronousSearchPostProcessorTests method testProcessSearchFailureOnDeletedContext.
public void testProcessSearchFailureOnDeletedContext() throws AsynchronousSearchStateMachineClosedException {
DiscoveryNode discoveryNode = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), Collections.emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
AtomicBoolean activeContextCleanUpConsumerInvocation = new AtomicBoolean();
ThreadPool testThreadPool = null;
try {
testThreadPool = new TestThreadPool(this.getClass().getName(), 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()));
AsynchronousSearchStateMachine asStateMachine = asService.getStateMachine();
ClusterService clusterService = ClusterServiceUtils.createClusterService(testThreadPool, discoveryNode, clusterSettings);
AsynchronousSearchPostProcessor postProcessor = new AsynchronousSearchPostProcessor(persistenceService, asActiveStore, asStateMachine, (context) -> activeContextCleanUpConsumerInvocation.compareAndSet(false, true), testThreadPool, clusterService);
SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(new SearchRequest());
submitAsynchronousSearchRequest.keepOnCompletion(true);
submitAsynchronousSearchRequest.keepAlive(TimeValue.timeValueHours(1));
AsynchronousSearchActiveContext context = (AsynchronousSearchActiveContext) asService.createAndStoreContext(submitAsynchronousSearchRequest, System.currentTimeMillis(), () -> InternalAggregationTestCase.emptyReduceContextBuilder(), null);
ShardSearchFailure shardSearchFailure = new ShardSearchFailure(new RuntimeException("runtime-exception"));
SearchPhaseExecutionException exception = new SearchPhaseExecutionException("phase", "msg", new NullPointerException(), new ShardSearchFailure[] { shardSearchFailure });
asStateMachine.trigger(new SearchStartedEvent(context, new SearchTask(0, "n/a", "n/a", () -> "test", null, Collections.emptyMap())));
asStateMachine.trigger(new SearchDeletedEvent(context));
AsynchronousSearchResponse asResponse = postProcessor.processSearchFailure(exception, context.getContextId());
assertNull(asResponse.getId());
assertNull(asResponse.getSearchResponse());
assertEquals(-1L, asResponse.getExpirationTimeMillis());
assertEquals(-1L, asResponse.getStartTimeMillis());
assertEquals(AsynchronousSearchState.FAILED, asResponse.getState());
assertThat(asResponse.getError(), instanceOf(SearchPhaseExecutionException.class));
assertFalse(activeContextCleanUpConsumerInvocation.get());
assertEquals(0, fakeClient.persistenceCount);
} finally {
ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
}
}
Aggregations