Examples with AtomicArray org.opensearch.common.util.concurrent.AtomicArray used on opensource projects

Example 16 with AtomicArray

use of org.opensearch.common.util.concurrent.AtomicArray in project OpenSearch by opensearch-project.

the class MultiSearchActionTookTests method createTransportMultiSearchAction.

private TransportMultiSearchAction createTransportMultiSearchAction(boolean controlledClock, AtomicLong expected) {
    Settings settings = Settings.builder().put("node.name", TransportMultiSearchActionTests.class.getSimpleName()).build();
    TransportService transportService = new TransportService(Settings.EMPTY, mock(Transport.class), null, TransportService.NOOP_TRANSPORT_INTERCEPTOR, boundAddress -> DiscoveryNode.createLocal(settings, boundAddress.publishAddress(), UUIDs.randomBase64UUID()), null, Collections.emptySet()) {

        @Override
        public TaskManager getTaskManager() {
            return taskManager;
        }
    };
    ActionFilters actionFilters = new ActionFilters(new HashSet<>());
    ClusterService clusterService = mock(ClusterService.class);
    when(clusterService.state()).thenReturn(ClusterState.builder(new ClusterName("test")).build());
    final int availableProcessors = Runtime.getRuntime().availableProcessors();
    AtomicInteger counter = new AtomicInteger();
    final List<String> threadPoolNames = Arrays.asList(ThreadPool.Names.GENERIC, ThreadPool.Names.SAME);
    Randomness.shuffle(threadPoolNames);
    final ExecutorService commonExecutor = threadPool.executor(threadPoolNames.get(0));
    final Set<SearchRequest> requests = Collections.newSetFromMap(Collections.synchronizedMap(new IdentityHashMap<>()));
    NodeClient client = new NodeClient(settings, threadPool) {

        @Override
        public void search(final SearchRequest request, final ActionListener<SearchResponse> listener) {
            requests.add(request);
            commonExecutor.execute(() -> {
                counter.decrementAndGet();
                listener.onResponse(new SearchResponse(InternalSearchResponse.empty(), null, 0, 0, 0, 0L, ShardSearchFailure.EMPTY_ARRAY, SearchResponse.Clusters.EMPTY));
            });
        }

        @Override
        public String getLocalNodeId() {
            return "local_node_id";
        }
    };
    if (controlledClock) {
        return new TransportMultiSearchAction(threadPool, actionFilters, transportService, clusterService, availableProcessors, expected::get, client) {

            @Override
            void executeSearch(final Queue<SearchRequestSlot> requests, final AtomicArray<MultiSearchResponse.Item> responses, final AtomicInteger responseCounter, final ActionListener<MultiSearchResponse> listener, long startTimeInNanos) {
                expected.set(1000000);
                super.executeSearch(requests, responses, responseCounter, listener, startTimeInNanos);
            }
        };
    } else {
        return new TransportMultiSearchAction(threadPool, actionFilters, transportService, clusterService, availableProcessors, System::nanoTime, client) {

            @Override
            void executeSearch(final Queue<SearchRequestSlot> requests, final AtomicArray<MultiSearchResponse.Item> responses, final AtomicInteger responseCounter, final ActionListener<MultiSearchResponse> listener, long startTimeInNanos) {
                long elapsed = spinForAtLeastNMilliseconds(randomIntBetween(0, 10));
                expected.set(elapsed);
                super.executeSearch(requests, responses, responseCounter, listener, startTimeInNanos);
            }
        };
    }
}

Example 17 with AtomicArray

use of org.opensearch.common.util.concurrent.AtomicArray in project OpenSearch by opensearch-project.

the class DfsQueryPhaseTests method testDfsWith1ShardFailed.

public void testDfsWith1ShardFailed() throws IOException {
    AtomicArray<DfsSearchResult> results = new AtomicArray<>(2);
    AtomicReference<AtomicArray<SearchPhaseResult>> responseRef = new AtomicReference<>();
    results.set(0, newSearchResult(0, new ShardSearchContextId("", 1), new SearchShardTarget("node1", new ShardId("test", "na", 0), null, OriginalIndices.NONE)));
    results.set(1, newSearchResult(1, new ShardSearchContextId("", 2), new SearchShardTarget("node2", new ShardId("test", "na", 0), null, OriginalIndices.NONE)));
    results.get(0).termsStatistics(new Term[0], new TermStatistics[0]);
    results.get(1).termsStatistics(new Term[0], new TermStatistics[0]);
    SearchTransportService searchTransportService = new SearchTransportService(null, null) {

        @Override
        public void sendExecuteQuery(Transport.Connection connection, QuerySearchRequest request, SearchTask task, SearchActionListener<QuerySearchResult> listener) {
            if (request.contextId().getId() == 1) {
                QuerySearchResult queryResult = new QuerySearchResult(new ShardSearchContextId("", 123), new SearchShardTarget("node1", new ShardId("test", "na", 0), null, OriginalIndices.NONE), null);
                queryResult.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(1, TotalHits.Relation.EQUAL_TO), new ScoreDoc[] { new ScoreDoc(42, 1.0F) }), 2.0F), new DocValueFormat[0]);
                // the size of the result set
                queryResult.size(2);
                listener.onResponse(queryResult);
            } else if (request.contextId().getId() == 2) {
                listener.onFailure(new MockDirectoryWrapper.FakeIOException());
            } else {
                fail("no such request ID: " + request.contextId());
            }
        }
    };
    SearchPhaseController searchPhaseController = searchPhaseController();
    MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
    mockSearchPhaseContext.searchTransport = searchTransportService;
    QueryPhaseResultConsumer consumer = searchPhaseController.newSearchPhaseResults(OpenSearchExecutors.newDirectExecutorService(), new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, mockSearchPhaseContext.searchRequest, results.length(), exc -> {
    });
    DfsQueryPhase phase = new DfsQueryPhase(results.asList(), null, consumer, (response) -> new SearchPhase("test") {

        @Override
        public void run() throws IOException {
            responseRef.set(response.results);
        }
    }, mockSearchPhaseContext);
    assertEquals("dfs_query", phase.getName());
    phase.run();
    mockSearchPhaseContext.assertNoFailure();
    assertNotNull(responseRef.get());
    assertNotNull(responseRef.get().get(0));
    assertNull(responseRef.get().get(0).fetchResult());
    assertEquals(1, responseRef.get().get(0).queryResult().topDocs().topDocs.totalHits.value);
    assertEquals(42, responseRef.get().get(0).queryResult().topDocs().topDocs.scoreDocs[0].doc);
    assertNull(responseRef.get().get(1));
    assertEquals(1, mockSearchPhaseContext.numSuccess.get());
    assertEquals(1, mockSearchPhaseContext.failures.size());
    assertTrue(mockSearchPhaseContext.failures.get(0).getCause() instanceof MockDirectoryWrapper.FakeIOException);
    assertEquals(1, mockSearchPhaseContext.releasedSearchContexts.size());
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.contains(new ShardSearchContextId("", 2L)));
    assertNull(responseRef.get().get(1));
}

Example 18 with AtomicArray

use of org.opensearch.common.util.concurrent.AtomicArray in project OpenSearch by opensearch-project.

the class DfsQueryPhaseTests method testDfsWith2Shards.

public void testDfsWith2Shards() throws IOException {
    AtomicArray<DfsSearchResult> results = new AtomicArray<>(2);
    AtomicReference<AtomicArray<SearchPhaseResult>> responseRef = new AtomicReference<>();
    results.set(0, newSearchResult(0, new ShardSearchContextId("", 1), new SearchShardTarget("node1", new ShardId("test", "na", 0), null, OriginalIndices.NONE)));
    results.set(1, newSearchResult(1, new ShardSearchContextId("", 2), new SearchShardTarget("node2", new ShardId("test", "na", 0), null, OriginalIndices.NONE)));
    results.get(0).termsStatistics(new Term[0], new TermStatistics[0]);
    results.get(1).termsStatistics(new Term[0], new TermStatistics[0]);
    SearchTransportService searchTransportService = new SearchTransportService(null, null) {

        @Override
        public void sendExecuteQuery(Transport.Connection connection, QuerySearchRequest request, SearchTask task, SearchActionListener<QuerySearchResult> listener) {
            if (request.contextId().getId() == 1) {
                QuerySearchResult queryResult = new QuerySearchResult(new ShardSearchContextId("", 123), new SearchShardTarget("node1", new ShardId("test", "na", 0), null, OriginalIndices.NONE), null);
                queryResult.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(1, TotalHits.Relation.EQUAL_TO), new ScoreDoc[] { new ScoreDoc(42, 1.0F) }), 2.0F), new DocValueFormat[0]);
                // the size of the result set
                queryResult.size(2);
                listener.onResponse(queryResult);
            } else if (request.contextId().getId() == 2) {
                QuerySearchResult queryResult = new QuerySearchResult(new ShardSearchContextId("", 123), new SearchShardTarget("node2", new ShardId("test", "na", 0), null, OriginalIndices.NONE), null);
                queryResult.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(1, TotalHits.Relation.EQUAL_TO), new ScoreDoc[] { new ScoreDoc(84, 2.0F) }), 2.0F), new DocValueFormat[0]);
                // the size of the result set
                queryResult.size(2);
                listener.onResponse(queryResult);
            } else {
                fail("no such request ID: " + request.contextId());
            }
        }
    };
    SearchPhaseController searchPhaseController = searchPhaseController();
    MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
    mockSearchPhaseContext.searchTransport = searchTransportService;
    QueryPhaseResultConsumer consumer = searchPhaseController.newSearchPhaseResults(OpenSearchExecutors.newDirectExecutorService(), new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, mockSearchPhaseContext.searchRequest, results.length(), exc -> {
    });
    DfsQueryPhase phase = new DfsQueryPhase(results.asList(), null, consumer, (response) -> new SearchPhase("test") {

        @Override
        public void run() throws IOException {
            responseRef.set(response.results);
        }
    }, mockSearchPhaseContext);
    assertEquals("dfs_query", phase.getName());
    phase.run();
    mockSearchPhaseContext.assertNoFailure();
    assertNotNull(responseRef.get());
    assertNotNull(responseRef.get().get(0));
    assertNull(responseRef.get().get(0).fetchResult());
    assertEquals(1, responseRef.get().get(0).queryResult().topDocs().topDocs.totalHits.value);
    assertEquals(42, responseRef.get().get(0).queryResult().topDocs().topDocs.scoreDocs[0].doc);
    assertNotNull(responseRef.get().get(1));
    assertNull(responseRef.get().get(1).fetchResult());
    assertEquals(1, responseRef.get().get(1).queryResult().topDocs().topDocs.totalHits.value);
    assertEquals(84, responseRef.get().get(1).queryResult().topDocs().topDocs.scoreDocs[0].doc);
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.isEmpty());
    assertEquals(2, mockSearchPhaseContext.numSuccess.get());
}

Example 19 with AtomicArray

use of org.opensearch.common.util.concurrent.AtomicArray in project OpenSearch by opensearch-project.

the class TransportSearchHelperTests method generateQueryResults.

public static AtomicArray<SearchPhaseResult> generateQueryResults() {
    AtomicArray<SearchPhaseResult> array = new AtomicArray<>(3);
    DiscoveryNode node1 = new DiscoveryNode("node_1", buildNewFakeTransportAddress(), Version.CURRENT);
    DiscoveryNode node2 = new DiscoveryNode("node_2", buildNewFakeTransportAddress(), Version.CURRENT);
    DiscoveryNode node3 = new DiscoveryNode("node_3", buildNewFakeTransportAddress(), Version.CURRENT);
    SearchAsyncActionTests.TestSearchPhaseResult testSearchPhaseResult1 = new SearchAsyncActionTests.TestSearchPhaseResult(new ShardSearchContextId("a", 1), node1);
    testSearchPhaseResult1.setSearchShardTarget(new SearchShardTarget("node_1", new ShardId("idx", "uuid1", 2), "cluster_x", null));
    SearchAsyncActionTests.TestSearchPhaseResult testSearchPhaseResult2 = new SearchAsyncActionTests.TestSearchPhaseResult(new ShardSearchContextId("b", 12), node2);
    testSearchPhaseResult2.setSearchShardTarget(new SearchShardTarget("node_2", new ShardId("idy", "uuid2", 42), "cluster_y", null));
    SearchAsyncActionTests.TestSearchPhaseResult testSearchPhaseResult3 = new SearchAsyncActionTests.TestSearchPhaseResult(new ShardSearchContextId("c", 42), node3);
    testSearchPhaseResult3.setSearchShardTarget(new SearchShardTarget("node_3", new ShardId("idy", "uuid2", 43), null, null));
    array.setOnce(0, testSearchPhaseResult1);
    array.setOnce(1, testSearchPhaseResult2);
    array.setOnce(2, testSearchPhaseResult3);
    return array;
}

Example 20 with AtomicArray

use of org.opensearch.common.util.concurrent.AtomicArray in project OpenSearch by opensearch-project.

the class IndexShardTests method testAcquirePrimaryAllOperationsPermits.

public void testAcquirePrimaryAllOperationsPermits() throws Exception {
    final IndexShard indexShard = newStartedShard(true);
    assertEquals(0, indexShard.getActiveOperationsCount());
    final CountDownLatch allPermitsAcquired = new CountDownLatch(1);
    final Thread[] threads = new Thread[randomIntBetween(2, 5)];
    final List<PlainActionFuture<Releasable>> futures = new ArrayList<>(threads.length);
    final AtomicArray<Tuple<Boolean, Exception>> results = new AtomicArray<>(threads.length);
    final CountDownLatch allOperationsDone = new CountDownLatch(threads.length);
    for (int i = 0; i < threads.length; i++) {
        final int threadId = i;
        final boolean singlePermit = randomBoolean();
        final PlainActionFuture<Releasable> future = new PlainActionFuture<Releasable>() {

            @Override
            public void onResponse(final Releasable releasable) {
                if (singlePermit) {
                    assertThat(indexShard.getActiveOperationsCount(), greaterThan(0));
                } else {
                    assertThat(indexShard.getActiveOperationsCount(), equalTo(IndexShard.OPERATIONS_BLOCKED));
                }
                releasable.close();
                super.onResponse(releasable);
                results.setOnce(threadId, Tuple.tuple(Boolean.TRUE, null));
                allOperationsDone.countDown();
            }

            @Override
            public void onFailure(final Exception e) {
                results.setOnce(threadId, Tuple.tuple(Boolean.FALSE, e));
                allOperationsDone.countDown();
            }
        };
        futures.add(threadId, future);
        threads[threadId] = new Thread(() -> {
            try {
                allPermitsAcquired.await();
            } catch (final InterruptedException e) {
                throw new RuntimeException(e);
            }
            if (singlePermit) {
                indexShard.acquirePrimaryOperationPermit(future, ThreadPool.Names.WRITE, "");
            } else {
                indexShard.acquireAllPrimaryOperationsPermits(future, TimeValue.timeValueHours(1L));
            }
        });
        threads[threadId].start();
    }
    final AtomicBoolean blocked = new AtomicBoolean();
    final CountDownLatch allPermitsTerminated = new CountDownLatch(1);
    final PlainActionFuture<Releasable> futureAllPermits = new PlainActionFuture<Releasable>() {

        @Override
        public void onResponse(final Releasable releasable) {
            try {
                blocked.set(true);
                allPermitsAcquired.countDown();
                super.onResponse(releasable);
                allPermitsTerminated.await();
            } catch (final InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
    };
    indexShard.acquireAllPrimaryOperationsPermits(futureAllPermits, TimeValue.timeValueSeconds(30L));
    allPermitsAcquired.await();
    assertTrue(blocked.get());
    assertEquals(IndexShard.OPERATIONS_BLOCKED, indexShard.getActiveOperationsCount());
    assertTrue("Expected no results, operations are blocked", results.asList().isEmpty());
    futures.forEach(future -> assertFalse(future.isDone()));
    allPermitsTerminated.countDown();
    final Releasable allPermits = futureAllPermits.get();
    assertTrue(futureAllPermits.isDone());
    assertTrue("Expected no results, operations are blocked", results.asList().isEmpty());
    futures.forEach(future -> assertFalse(future.isDone()));
    Releasables.close(allPermits);
    allOperationsDone.await();
    for (Thread thread : threads) {
        thread.join();
    }
    futures.forEach(future -> assertTrue(future.isDone()));
    assertEquals(threads.length, results.asList().size());
    results.asList().forEach(result -> {
        assertTrue(result.v1());
        assertNull(result.v2());
    });
    closeShards(indexShard);
}