use of org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener in project asynchronous-search by opensearch-project.
the class AsynchronousSearchService method createAndStoreContext.
/**
* Creates a new active asynchronous search for a newly submitted asynchronous search.
*
* @param request the SubmitAsynchronousSearchRequest
* @param relativeStartTimeMillis the relative start time of the search in millis
* @param user current user
* @param reduceContextBuilder the reference for the reduceContextBuilder
* @return the AsynchronousSearchContext for the submitted request
*/
public AsynchronousSearchContext createAndStoreContext(SubmitAsynchronousSearchRequest request, long relativeStartTimeMillis, Supplier<InternalAggregation.ReduceContextBuilder> reduceContextBuilder, User user) {
validateRequest(request);
AsynchronousSearchContextId asynchronousSearchContextId = new AsynchronousSearchContextId(UUIDs.base64UUID(), idGenerator.incrementAndGet());
contextEventListener.onNewContext(asynchronousSearchContextId);
AsynchronousSearchProgressListener progressActionListener = new AsynchronousSearchProgressListener(relativeStartTimeMillis, (response) -> asynchronousSearchPostProcessor.processSearchResponse(response, asynchronousSearchContextId), (e) -> asynchronousSearchPostProcessor.processSearchFailure(e, asynchronousSearchContextId), threadPool.executor(AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME), threadPool::relativeTimeInMillis, reduceContextBuilder);
AsynchronousSearchActiveContext asynchronousSearchContext = new AsynchronousSearchActiveContext(asynchronousSearchContextId, clusterService.localNode().getId(), request.getKeepAlive(), request.getKeepOnCompletion(), threadPool, currentTimeSupplier, progressActionListener, user, () -> persistSearchFailure);
asynchronousSearchActiveStore.putContext(asynchronousSearchContextId, asynchronousSearchContext, contextEventListener::onContextRejected);
contextEventListener.onContextInitialized(asynchronousSearchContextId);
return asynchronousSearchContext;
}
use of org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener in project asynchronous-search by opensearch-project.
the class AsynchronousSearchRejectionIT method testSearchFailures.
public void testSearchFailures() throws Exception {
for (int i = 0; i < 10; i++) {
client().prepareIndex("test").setId(Integer.toString(i)).setSource("field", "1").get();
}
int numberOfAsyncOps = randomIntBetween(100, 200);
final CountDownLatch latch = new CountDownLatch(numberOfAsyncOps);
final CopyOnWriteArrayList<Object> responses = new CopyOnWriteArrayList<>();
TestThreadPool threadPool = null;
try {
threadPool = new TestThreadPool(AsynchronousSearchProgressListenerIT.class.getName());
for (int i = 0; i < numberOfAsyncOps; i++) {
SearchRequest request = client().prepareSearch("test").setSearchType(SearchType.QUERY_THEN_FETCH).setQuery(QueryBuilders.matchQuery("field", "1")).request();
AtomicReference<SearchResponse> responseRef = new AtomicReference<>();
AtomicInteger reduceContextInvocation = new AtomicInteger();
AsynchronousSearchProgressListener listener;
SearchService service = internalCluster().getInstance(SearchService.class);
InternalAggregation.ReduceContextBuilder reduceContextBuilder = service.aggReduceContextBuilder(request);
AtomicReference<Exception> exceptionRef = new AtomicReference<>();
Function<SearchResponse, AsynchronousSearchResponse> responseFunction = (r) -> null;
Function<Exception, AsynchronousSearchResponse> failureFunction = (e) -> null;
listener = new AsynchronousSearchProgressListener(threadPool.relativeTimeInMillis(), responseFunction, failureFunction, threadPool.generic(), threadPool::relativeTimeInMillis, () -> reduceContextBuilder) {
@Override
public void onResponse(SearchResponse searchResponse) {
assertTrue(responseRef.compareAndSet(null, searchResponse));
AsynchronousSearchAssertions.assertSearchResponses(searchResponse, this.partialResponse());
latch.countDown();
}
@Override
public void onFailure(Exception exception) {
assertTrue(exceptionRef.compareAndSet(null, exception));
latch.countDown();
}
};
client().execute(SearchAction.INSTANCE, new SearchRequest(request) {
@Override
public SearchTask createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
SearchTask task = super.createTask(id, type, action, parentTaskId, headers);
task.setProgressListener(listener);
return task;
}
}, listener);
}
latch.await();
} finally {
ThreadPool.terminate(threadPool, 100, TimeUnit.MILLISECONDS);
}
}
use of org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener in project asynchronous-search by opensearch-project.
the class AsynchronousSearchActiveContextTests method testInitializeContext.
public void testInitializeContext() {
TestThreadPool threadPool = null;
try {
int writeThreadPoolSize = randomIntBetween(1, 2);
int writeThreadPoolQueueSize = randomIntBetween(1, 2);
Settings settings = Settings.builder().put("thread_pool." + AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME + ".size", writeThreadPoolSize).put("thread_pool." + AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME + ".queue_size", writeThreadPoolQueueSize).build();
final int availableProcessors = OpenSearchExecutors.allocatedProcessors(settings);
ScalingExecutorBuilder scalingExecutorBuilder = new ScalingExecutorBuilder(AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME, 1, Math.min(2 * availableProcessors, Math.max(128, 512)), TimeValue.timeValueMinutes(30));
threadPool = new TestThreadPool("Tests", settings, scalingExecutorBuilder);
String node = UUID.randomUUID().toString();
AsynchronousSearchProgressListener asProgressListener = mockAsynchronousSearchProgressListener(threadPool);
AsynchronousSearchContextId asContextId = new AsynchronousSearchContextId(UUID.randomUUID().toString(), randomNonNegativeLong());
boolean keepOnCompletion = randomBoolean();
User user = TestClientUtils.randomUser();
TimeValue keepAlive = TimeValue.timeValueDays(randomInt(100));
AsynchronousSearchActiveContext context = new AsynchronousSearchActiveContext(asContextId, node, keepAlive, keepOnCompletion, threadPool, threadPool::absoluteTimeInMillis, asProgressListener, user, () -> true);
assertEquals(AsynchronousSearchState.INIT, context.getAsynchronousSearchState());
assertNull(context.getTask());
assertNull(context.getAsynchronousSearchId());
assertEquals(context.getAsynchronousSearchState(), AsynchronousSearchState.INIT);
} finally {
ThreadPool.terminate(threadPool, 30, TimeUnit.SECONDS);
}
}
use of org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener in project asynchronous-search by opensearch-project.
the class TransportGetAsynchronousSearchAction method handleWaitForCompletion.
private void handleWaitForCompletion(AsynchronousSearchContext context, GetAsynchronousSearchRequest request, ActionListener<AsynchronousSearchResponse> listener) {
// We wait for a response only if a wait for completion is non-null and the search execution is still in progress.
if (context.isRunning() && request.getWaitForCompletionTimeout() != null) {
logger.debug("Context is running for asynchronous search id [{}]", context.getAsynchronousSearchId());
AsynchronousSearchProgressListener progressActionListener = context.getAsynchronousSearchProgressListener();
assert progressActionListener != null : "progress listener cannot be null";
PrioritizedActionListener<AsynchronousSearchResponse> wrappedListener = AsynchronousSearchTimeoutWrapper.wrapScheduledTimeout(threadPool, request.getWaitForCompletionTimeout(), OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME, listener, (actionListener) -> {
progressActionListener.searchProgressActionListener().removeListener(actionListener);
listener.onResponse(context.getAsynchronousSearchResponse());
});
progressActionListener.searchProgressActionListener().addOrExecuteListener(wrappedListener);
} else {
// we don't need to wait any further on search completion
logger.debug("Context is not running for asynchronous search id [{}]", context.getAsynchronousSearchId());
listener.onResponse(context.getAsynchronousSearchResponse());
}
}
use of org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener in project asynchronous-search by opensearch-project.
the class AsynchronousSearchServiceFreeContextTests method testFreeContextPermitAcquisitionFailure.
public void testFreeContextPermitAcquisitionFailure() 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());
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);
simulateUncheckedException = true;
persisted = false;
when(mockStore.getContext(asContextId)).thenReturn(Optional.of(asActiveContext));
CountDownLatch latch = new CountDownLatch(1);
asService.freeContext(asActiveContext.getAsynchronousSearchId(), asActiveContext.getContextId(), null, new LatchedActionListener<>(wrap(r -> fail("Expected resource_not_found_exception. Got acknowledgement " + r), e -> {
assertTrue(e.getClass().getName(), e instanceof ResourceNotFoundException);
}), latch));
latch.await();
assertEquals(1, (int) mockClient.deleteCount);
mockClusterService.stop();
} finally {
ThreadPool.terminate(testThreadPool, 30, TimeUnit.SECONDS);
}
}
Aggregations