Examples with AsynchronousSearchTask org.opensearch.search.asynchronous.task.AsynchronousSearchTask used on opensource projects

Example 21 with AsynchronousSearchTask

use of org.opensearch.search.asynchronous.task.AsynchronousSearchTask in project asynchronous-search by opensearch-project.

the class AsynchronousSearchServiceUpdateContextTests method testUpdateClosedContext.

public void testUpdateClosedContext() 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 asActiveStore = new AsynchronousSearchActiveStore(mockClusterService);
        AsynchronousSearchPersistenceService persistenceService = new AsynchronousSearchPersistenceService(mockClient, mockClusterService, testThreadPool);
        AsynchronousSearchService asService = new AsynchronousSearchService(persistenceService, asActiveStore, mockClient, mockClusterService, testThreadPool, new InternalAsynchronousSearchStats(), new NamedWriteableRegistry(emptyList()));
        TimeValue keepAlive = timeValueHours(9);
        boolean keepOnCompletion = randomBoolean();
        User user1 = 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 activeContext = (AsynchronousSearchActiveContext) context;
        assertNull(activeContext.getTask());
        assertNull(activeContext.getAsynchronousSearchId());
        assertEquals(activeContext.getAsynchronousSearchState(), INIT);
        assertEquals(activeContext.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(activeContext.getTask(), task);
        assertEquals(activeContext.getStartTimeMillis(), task.getStartTime());
        assertEquals(activeContext.getExpirationTimeMillis(), task.getStartTime() + keepAlive.millis());
        assertEquals(activeContext.getAsynchronousSearchState(), RUNNING);
        context.setState(CLOSED);
        ((AsynchronousSearchActiveContext) context).close();
        CountDownLatch latch = new CountDownLatch(1);
        asService.updateKeepAliveAndGetContext(context.getAsynchronousSearchId(), timeValueHours(9), context.getContextId(), user1, new LatchedActionListener<>(wrap(r -> {
            if (keepOnCompletion) {
                assertTrue(r instanceof AsynchronousSearchPersistenceContext);
            } else {
                fail("expected resource not found exception, got result.");
            }
        }, e -> {
            if (keepOnCompletion) {
                fail("expected resource not found exception, got result");
            } else {
                assertTrue(e instanceof ResourceNotFoundException);
            }
        }), latch));
        latch.await();
        mockClusterService.stop();
    } finally {
        ThreadPool.terminate(testThreadPool, 30, TimeUnit.SECONDS);
    }
}

Example 22 with AsynchronousSearchTask

use of org.opensearch.search.asynchronous.task.AsynchronousSearchTask in project asynchronous-search by opensearch-project.

the class AsynchronousSearchActiveContextTests method testClosedContext.

public void testClosedContext() throws InterruptedException {
    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();
        TimeValue keepAlive = TimeValue.timeValueDays(randomInt(100));
        AsynchronousSearchActiveContext context = new AsynchronousSearchActiveContext(asContextId, node, keepAlive, keepOnCompletion, threadPool, threadPool::absoluteTimeInMillis, asProgressListener, null, () -> true);
        AsynchronousSearchTask task = new AsynchronousSearchTask(randomNonNegativeLong(), "transport", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap(), context, null, (c) -> {
        });
        context.setTask(task);
        assertEquals(task, context.getTask());
        assertEquals(task.getStartTime(), context.getStartTimeMillis());
        assertEquals(task.getStartTime() + keepAlive.getMillis(), context.getExpirationTimeMillis());
        assertTrue(context.isAlive());
        assertFalse(context.isExpired());
        expectThrows(SetOnce.AlreadySetException.class, () -> context.setTask(task));
        context.setState(AsynchronousSearchState.CLOSED);
        context.close();
        assertFalse(context.isAlive());
        expectThrows(AssertionError.class, () -> context.setExpirationTimeMillis(randomNonNegativeLong()));
        expectThrows(AssertionError.class, () -> context.setTask(task));
    } finally {
        ThreadPool.terminate(threadPool, 30, TimeUnit.SECONDS);
    }
}

Example 23 with AsynchronousSearchTask

use of org.opensearch.search.asynchronous.task.AsynchronousSearchTask in project asynchronous-search by opensearch-project.

the class AsynchronousSearchActiveContextTests method testTaskBootstrap.

public void testTaskBootstrap() {
    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);
        SubmitAsynchronousSearchRequest request = new SubmitAsynchronousSearchRequest(new SearchRequest("test"));
        request.keepAlive(keepAlive);
        request.keepOnCompletion(keepOnCompletion);
        AsynchronousSearchTask task = new AsynchronousSearchTask(randomNonNegativeLong(), "transport", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap(), context, request, (c) -> {
        });
        context.setTask(task);
        assertEquals(task, context.getTask());
        assertEquals(task.getStartTime(), context.getStartTimeMillis());
        assertEquals(task.getStartTime() + keepAlive.getMillis(), context.getExpirationTimeMillis());
        String description = task.getDescription();
        assertThat(description, containsString("[asynchronous search]"));
        assertThat(description, containsString("indices[test]"));
        assertThat(description, containsString("keep_alive[" + keepAlive + "]"));
        assertThat(description, containsString("keep_on_completion[" + keepOnCompletion + "]"));
        assertTrue(context.isAlive());
        assertFalse(context.isExpired());
        expectThrows(SetOnce.AlreadySetException.class, () -> context.setTask(task));
        if (keepOnCompletion) {
            assertTrue(context.shouldPersist());
        } else {
            assertFalse(context.shouldPersist());
        }
        assertTrue(user.equals(context.getUser()));
    } finally {
        ThreadPool.terminate(threadPool, 30, TimeUnit.SECONDS);
    }
}

Example 24 with AsynchronousSearchTask

use of org.opensearch.search.asynchronous.task.AsynchronousSearchTask in project asynchronous-search by opensearch-project.

the class TransportSubmitAsynchronousSearchAction method doExecute.

@Override
protected void doExecute(Task task, SubmitAsynchronousSearchRequest request, ActionListener<AsynchronousSearchResponse> listener) {
    AsynchronousSearchContext asynchronousSearchContext = null;
    String userStr = threadPool.getThreadContext().getTransient(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT);
    User user = User.parse(userStr);
    try {
        final long relativeStartTimeInMillis = threadPool.relativeTimeInMillis();
        asynchronousSearchContext = asynchronousSearchService.createAndStoreContext(request, relativeStartTimeInMillis, () -> searchService.aggReduceContextBuilder(request.getSearchRequest()), user);
        assert asynchronousSearchContext.getAsynchronousSearchProgressListener() != null : "missing progress listener for an active context";
        AsynchronousSearchProgressListener progressListener = asynchronousSearchContext.getAsynchronousSearchProgressListener();
        // making it effectively final for usage in anonymous class.
        AsynchronousSearchContext context = asynchronousSearchContext;
        SearchRequest searchRequest = new SearchRequest(request.getSearchRequest()) {

            @Override
            public SearchTask createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
                AsynchronousSearchTask asynchronousSearchTask = new AsynchronousSearchTask(id, type, AsynchronousSearchTask.NAME, parentTaskId, headers, (AsynchronousSearchActiveContext) context, request, asynchronousSearchService::onCancelledFreeActiveContext);
                asynchronousSearchService.bootstrapSearch(asynchronousSearchTask, context.getContextId());
                PrioritizedActionListener<AsynchronousSearchResponse> wrappedListener = AsynchronousSearchTimeoutWrapper.wrapScheduledTimeout(threadPool, request.getWaitForCompletionTimeout(), AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME, listener, (actionListener) -> {
                    progressListener.searchProgressActionListener().removeListener(actionListener);
                    listener.onResponse(context.getAsynchronousSearchResponse());
                });
                progressListener.searchProgressActionListener().addOrExecuteListener(wrappedListener);
                return asynchronousSearchTask;
            }
        };
        // set the parent task as the submit task for cancellation on connection close
        searchRequest.setParentTask(task.taskInfo(clusterService.localNode().getId(), false).getTaskId());
        transportSearchAction.execute(searchRequest, progressListener);
    } catch (Exception e) {
        logger.error(() -> new ParameterizedMessage("Failed to submit asynchronous search request [{}]", request), e);
        if (asynchronousSearchContext != null) {
            AsynchronousSearchActiveContext asynchronousSearchActiveContext = (AsynchronousSearchActiveContext) asynchronousSearchContext;
            asynchronousSearchService.freeContext(asynchronousSearchActiveContext.getAsynchronousSearchId(), asynchronousSearchActiveContext.getContextId(), user, ActionListener.wrap((r) -> {
                logger.debug(() -> new ParameterizedMessage("Successfully cleaned up context on submit asynchronous" + " search [{}] on failure", asynchronousSearchActiveContext.getAsynchronousSearchId()), e);
                listener.onFailure(e);
            }, (ex) -> {
                logger.debug(() -> new ParameterizedMessage("Failed to cleaned up context on submit asynchronous search" + " [{}] on failure", asynchronousSearchActiveContext.getAsynchronousSearchId()), ex);
                listener.onFailure(e);
            }));
        } else {
            listener.onFailure(e);
        }
    }
}