Examples with TestThreadPool org.opensearch.threadpool.TestThreadPool used on opensource projects

Example 6 with TestThreadPool

use of org.opensearch.threadpool.TestThreadPool in project asynchronous-search by opensearch-project.

the class AsynchronousSearchPartialResponseIT method testCase.

private void testCase(Client client, SearchRequest request) throws Exception {
    AtomicReference<SearchResponse> responseRef = new AtomicReference<>();
    AtomicInteger reduceContextInvocation = new AtomicInteger();
    TestThreadPool threadPool = null;
    AsynchronousSearchProgressListener listener;
    try {
        threadPool = new TestThreadPool(AsynchronousSearchProgressListenerIT.class.getName());
        SearchService service = internalCluster().getInstance(SearchService.class);
        InternalAggregation.ReduceContextBuilder reduceContextBuilder = service.aggReduceContextBuilder(request);
        AtomicReference<Exception> exceptionRef = new AtomicReference<>();
        CountDownLatch latch = new CountDownLatch(1);
        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(responseRef.get(), this.partialResponse());
                latch.countDown();
            }

            @Override
            protected void onPartialReduce(List<SearchShard> shards, TotalHits totalHits, InternalAggregations aggs, int reducePhase) {
                super.onPartialReduce(shards, totalHits, aggs, reducePhase);
                Terms terms = this.partialResponse().getAggregations().get("keys");
                List<? extends Terms.Bucket> buckets = terms.getBuckets();
                assertThat(buckets.size(), lessThanOrEqualTo(aggregationSize));
                reduceContextInvocation.incrementAndGet();
            }

            @Override
            protected void onFinalReduce(List<SearchShard> shards, TotalHits totalHits, InternalAggregations aggs, int reducePhase) {
                super.onFinalReduce(shards, totalHits, aggs, reducePhase);
                Terms terms = this.partialResponse().getAggregations().get("keys");
                List<? extends Terms.Bucket> buckets = terms.getBuckets();
                assertThat(buckets.size(), equalTo(aggregationSize));
            }

            @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();
        Terms terms = responseRef.get().getAggregations().get("keys");
        assertThat(reduceContextInvocation.get(), equalTo(responseRef.get().getNumReducePhases() - 1));
        List<? extends Terms.Bucket> buckets = terms.getBuckets();
        assertThat(buckets.size(), equalTo(aggregationSize));
    } finally {
        ThreadPool.terminate(threadPool, 100, TimeUnit.MILLISECONDS);
    }
}

Example 7 with TestThreadPool

use of org.opensearch.threadpool.TestThreadPool 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);
    }
}

Example 8 with TestThreadPool

use of org.opensearch.threadpool.TestThreadPool in project asynchronous-search by opensearch-project.

the class AsynchronousSearchActiveStoreTests method testGetNonExistentContext.

public void testGetNonExistentContext() {
    DiscoveryNode node = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), Collections.emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
    ThreadPool testThreadPool = null;
    try {
        testThreadPool = new TestThreadPool(this.getClass().getName(), executorBuilder);
        ClusterService mockClusterService = ClusterServiceUtils.createClusterService(testThreadPool, node, clusterSettings);
        AsynchronousSearchActiveStore activeStore = new AsynchronousSearchActiveStore(mockClusterService);
        AsynchronousSearchContextId asContextId = new AsynchronousSearchContextId(UUID.randomUUID().toString(), randomNonNegativeLong());
        Optional<AsynchronousSearchActiveContext> optional = activeStore.getContext(asContextId);
        assertFalse(optional.isPresent());
        assertEquals(activeStore.getAllContexts().size(), 0);
    } finally {
        ThreadPool.terminate(testThreadPool, 10, TimeUnit.SECONDS);
    }
}

Example 9 with TestThreadPool

use of org.opensearch.threadpool.TestThreadPool in project asynchronous-search by opensearch-project.

the class AsynchronousSearchContextPermitsTests method setupThreadPool.

@BeforeClass
public static void setupThreadPool() {
    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("PermitsTests", settings, scalingExecutorBuilder);
}

Example 10 with TestThreadPool

use of org.opensearch.threadpool.TestThreadPool 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);
    }
}