use of org.opensearch.search.asynchronous.context.AsynchronousSearchContext in project asynchronous-search by opensearch-project.
the class AsynchronousSearchService method findContext.
/**
* Tries to find an {@linkplain AsynchronousSearchActiveContext}. If not found,
* queries the {@linkplain AsynchronousSearchPersistenceService} for a hit. If a response is found, it builds and returns an
* {@linkplain AsynchronousSearchPersistenceContext}, else throws
* {@linkplain ResourceNotFoundException}
*
* @param id The asynchronous search id
* @param asynchronousSearchContextId the Async search context id
* @param user current user
* @param listener to be invoked on finding an {@linkplain AsynchronousSearchContext}
*/
public void findContext(String id, AsynchronousSearchContextId asynchronousSearchContextId, User user, ActionListener<AsynchronousSearchContext> listener) {
ActionListener<AsynchronousSearchContext> exceptionTranslationListener = getExceptionTranslationWrapper(id, listener);
Optional<AsynchronousSearchActiveContext> optionalAsynchronousSearchActiveContext = asynchronousSearchActiveStore.getContext(asynchronousSearchContextId);
// so most likely a CLOSED context is stale
if (optionalAsynchronousSearchActiveContext.isPresent() && optionalAsynchronousSearchActiveContext.get().isAlive()) {
logger.debug("Active context is present for asynchronous search ID [{}]", id);
AsynchronousSearchActiveContext asynchronousSearchActiveContext = optionalAsynchronousSearchActiveContext.get();
if (isUserValid(user, asynchronousSearchActiveContext.getUser()) == false) {
logger.debug("Invalid user requesting GET active context for asynchronous search id {}", id);
exceptionTranslationListener.onFailure(new OpenSearchSecurityException("User doesn't have necessary roles to access the asynchronous search with id " + id, RestStatus.FORBIDDEN));
} else {
exceptionTranslationListener.onResponse(asynchronousSearchActiveContext);
}
} else {
logger.debug("Active context is not present for asynchronous search ID [{}]", id);
persistenceService.getResponse(id, user, wrap((persistenceModel) -> exceptionTranslationListener.onResponse(new AsynchronousSearchPersistenceContext(id, asynchronousSearchContextId, persistenceModel, currentTimeSupplier, namedWriteableRegistry)), ex -> {
logger.debug(() -> new ParameterizedMessage("Context not found for ID in the system index {}", id), ex);
exceptionTranslationListener.onFailure(ex);
}));
}
}
use of org.opensearch.search.asynchronous.context.AsynchronousSearchContext in project asynchronous-search by opensearch-project.
the class AsynchronousSearchService method updateKeepAliveAndGetContext.
/**
* If an active context is found, a permit is acquired from
* {@linkplain AsynchronousSearchContextPermits}
* and on acquisition of permit, a check is performed to see if response has been persisted in system index. If true, we update
* expiration in index. Else we update expiration field in {@linkplain AsynchronousSearchActiveContext}.
*
* @param id asynchronous search id
* @param keepAlive the new keep alive duration
* @param asynchronousSearchContextId asynchronous search context id
* @param user current user
* @param listener listener to invoke after updating expiration.
*/
public void updateKeepAliveAndGetContext(String id, TimeValue keepAlive, AsynchronousSearchContextId asynchronousSearchContextId, User user, ActionListener<AsynchronousSearchContext> listener) {
ActionListener<AsynchronousSearchContext> exceptionTranslationWrapper = getExceptionTranslationWrapper(id, listener);
validateKeepAlive(keepAlive);
long requestedExpirationTime = currentTimeSupplier.getAsLong() + keepAlive.getMillis();
// find an active context on this node if one exists
Optional<AsynchronousSearchActiveContext> asynchronousSearchContextOptional = asynchronousSearchActiveStore.getContext(asynchronousSearchContextId);
// for all other stages we don't really care much as those contexts are destined to be discarded
if (asynchronousSearchContextOptional.isPresent()) {
AsynchronousSearchActiveContext asynchronousSearchActiveContext = asynchronousSearchContextOptional.get();
asynchronousSearchActiveContext.acquireContextPermitIfRequired(wrap(releasable -> {
ActionListener<AsynchronousSearchContext> releasableActionListener = runAfter(exceptionTranslationWrapper, releasable::close);
// At this point it's possible that the response would have been persisted to system index
if (asynchronousSearchActiveContext.isAlive() == false && asynchronousSearchActiveContext.keepOnCompletion()) {
logger.debug("Updating persistence store after state is PERSISTED asynchronous search id [{}] " + "for updating context", asynchronousSearchActiveContext.getAsynchronousSearchId());
persistenceService.updateExpirationTime(id, requestedExpirationTime, user, wrap((actionResponse) -> releasableActionListener.onResponse(new AsynchronousSearchPersistenceContext(id, asynchronousSearchContextId, actionResponse, currentTimeSupplier, namedWriteableRegistry)), releasableActionListener::onFailure));
} else {
if (isUserValid(user, asynchronousSearchActiveContext.getUser())) {
logger.debug("Updating persistence store: NO as state is NOT PERSISTED yet asynchronous search id [{}] " + "for updating context", asynchronousSearchActiveContext.getAsynchronousSearchId());
asynchronousSearchActiveContext.setExpirationTimeMillis(requestedExpirationTime);
releasableActionListener.onResponse(asynchronousSearchActiveContext);
} else {
releasableActionListener.onFailure(new OpenSearchSecurityException("User doesn't have necessary roles to access the " + "asynchronous search with id " + id, RestStatus.FORBIDDEN));
}
}
}, exception -> {
Throwable cause = ExceptionsHelper.unwrapCause(exception);
if (cause instanceof TimeoutException) {
// this should ideally not happen. This would mean we couldn't acquire permits within the timeout
logger.debug(() -> new ParameterizedMessage("Failed to acquire permits for " + "asynchronous search id [{}] for updating context within timeout 5s", asynchronousSearchActiveContext.getAsynchronousSearchId()), exception);
listener.onFailure(new OpenSearchTimeoutException(id));
} else {
// best effort we try an update the doc if one exists
if (asynchronousSearchActiveContext.keepOnCompletion()) {
logger.debug("Updating persistence store after failing to acquire permits for asynchronous search id [{}] for " + "updating context with expiration time [{}]", asynchronousSearchActiveContext.getAsynchronousSearchId(), requestedExpirationTime);
persistenceService.updateExpirationTime(id, requestedExpirationTime, user, wrap((actionResponse) -> exceptionTranslationWrapper.onResponse(new AsynchronousSearchPersistenceContext(id, asynchronousSearchContextId, actionResponse, currentTimeSupplier, namedWriteableRegistry)), exceptionTranslationWrapper::onFailure));
} else {
exceptionTranslationWrapper.onFailure(new ResourceNotFoundException(asynchronousSearchActiveContext.getAsynchronousSearchId()));
}
}
}), TimeValue.timeValueSeconds(5), "update keep alive");
} else {
// try update the doc on the index assuming there exists one.
logger.debug("Updating persistence store after active context evicted for asynchronous search id [{}] " + "for updating context", id);
persistenceService.updateExpirationTime(id, requestedExpirationTime, user, wrap((actionResponse) -> exceptionTranslationWrapper.onResponse(new AsynchronousSearchPersistenceContext(id, asynchronousSearchContextId, actionResponse, currentTimeSupplier, namedWriteableRegistry)), exceptionTranslationWrapper::onFailure));
}
}
use of org.opensearch.search.asynchronous.context.AsynchronousSearchContext 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.AsynchronousSearchContext 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.AsynchronousSearchContext in project asynchronous-search by opensearch-project.
the class AsynchronousSearchServiceTests method testFindContextsToReap.
public void testFindContextsToReap() {
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) {
@Override
public long absoluteTimeInMillis() {
// simulate search has over run)
return System.currentTimeMillis() - 24 * 3600 * 1000;
}
};
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(asService.getContextsToReap().contains(context));
} finally {
ThreadPool.terminate(testThreadPool, 30, TimeUnit.SECONDS);
}
}
Aggregations