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

Example 16 with CountDown

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

the class GoogleCloudStorageBlobContainerRetriesTests method testWriteBlobWithRetries.

public void testWriteBlobWithRetries() throws Exception {
    final int maxRetries = randomIntBetween(2, 10);
    final CountDown countDown = new CountDown(maxRetries);
    final byte[] bytes = randomBlobContent();
    httpServer.createContext("/upload/storage/v1/b/bucket/o", safeHandler(exchange -> {
        assertThat(exchange.getRequestURI().getQuery(), containsString("uploadType=multipart"));
        if (countDown.countDown()) {
            Optional<Tuple<String, BytesReference>> content = parseMultipartRequestBody(exchange.getRequestBody());
            assertThat(content.isPresent(), is(true));
            assertThat(content.get().v1(), equalTo("write_blob_max_retries"));
            if (Objects.deepEquals(bytes, BytesReference.toBytes(content.get().v2()))) {
                byte[] response = ("{\"bucket\":\"bucket\",\"name\":\"" + content.get().v1() + "\"}").getBytes(UTF_8);
                exchange.getResponseHeaders().add("Content-Type", "application/json");
                exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length);
                exchange.getResponseBody().write(response);
            } else {
                exchange.sendResponseHeaders(HttpStatus.SC_BAD_REQUEST, -1);
            }
            return;
        }
        if (randomBoolean()) {
            if (randomBoolean()) {
                Streams.readFully(exchange.getRequestBody(), new byte[randomIntBetween(1, Math.max(1, bytes.length - 1))]);
            } else {
                Streams.readFully(exchange.getRequestBody());
                exchange.sendResponseHeaders(HttpStatus.SC_INTERNAL_SERVER_ERROR, -1);
            }
        }
    }));
    final BlobContainer blobContainer = createBlobContainer(maxRetries, null, null, null);
    try (InputStream stream = new InputStreamIndexInput(new ByteArrayIndexInput("desc", bytes), bytes.length)) {
        blobContainer.writeBlob("write_blob_max_retries", stream, bytes.length, false);
    }
    assertThat(countDown.isCountedDown(), is(true));
}

Example 17 with CountDown

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

the class TransportSearchAction method ccsRemoteReduce.

static void ccsRemoteReduce(SearchRequest searchRequest, OriginalIndices localIndices, Map<String, OriginalIndices> remoteIndices, SearchTimeProvider timeProvider, InternalAggregation.ReduceContextBuilder aggReduceContextBuilder, RemoteClusterService remoteClusterService, ThreadPool threadPool, ActionListener<SearchResponse> listener, BiConsumer<SearchRequest, ActionListener<SearchResponse>> localSearchConsumer) {
    if (localIndices == null && remoteIndices.size() == 1) {
        // if we are searching against a single remote cluster, we simply forward the original search request to such cluster
        // and we directly perform final reduction in the remote cluster
        Map.Entry<String, OriginalIndices> entry = remoteIndices.entrySet().iterator().next();
        String clusterAlias = entry.getKey();
        boolean skipUnavailable = remoteClusterService.isSkipUnavailable(clusterAlias);
        OriginalIndices indices = entry.getValue();
        SearchRequest ccsSearchRequest = SearchRequest.subSearchRequest(searchRequest, indices.indices(), clusterAlias, timeProvider.getAbsoluteStartMillis(), true);
        Client remoteClusterClient = remoteClusterService.getRemoteClusterClient(threadPool, clusterAlias);
        remoteClusterClient.search(ccsSearchRequest, new ActionListener<SearchResponse>() {

            @Override
            public void onResponse(SearchResponse searchResponse) {
                Map<String, ProfileShardResult> profileResults = searchResponse.getProfileResults();
                SearchProfileShardResults profile = profileResults == null || profileResults.isEmpty() ? null : new SearchProfileShardResults(profileResults);
                InternalSearchResponse internalSearchResponse = new InternalSearchResponse(searchResponse.getHits(), (InternalAggregations) searchResponse.getAggregations(), searchResponse.getSuggest(), profile, searchResponse.isTimedOut(), searchResponse.isTerminatedEarly(), searchResponse.getNumReducePhases());
                listener.onResponse(new SearchResponse(internalSearchResponse, searchResponse.getScrollId(), searchResponse.getTotalShards(), searchResponse.getSuccessfulShards(), searchResponse.getSkippedShards(), timeProvider.buildTookInMillis(), searchResponse.getShardFailures(), new SearchResponse.Clusters(1, 1, 0), searchResponse.pointInTimeId()));
            }

            @Override
            public void onFailure(Exception e) {
                if (skipUnavailable) {
                    listener.onResponse(SearchResponse.empty(timeProvider::buildTookInMillis, new SearchResponse.Clusters(1, 0, 1)));
                } else {
                    listener.onFailure(wrapRemoteClusterFailure(clusterAlias, e));
                }
            }
        });
    } else {
        SearchResponseMerger searchResponseMerger = createSearchResponseMerger(searchRequest.source(), timeProvider, aggReduceContextBuilder);
        AtomicInteger skippedClusters = new AtomicInteger(0);
        final AtomicReference<Exception> exceptions = new AtomicReference<>();
        int totalClusters = remoteIndices.size() + (localIndices == null ? 0 : 1);
        final CountDown countDown = new CountDown(totalClusters);
        for (Map.Entry<String, OriginalIndices> entry : remoteIndices.entrySet()) {
            String clusterAlias = entry.getKey();
            boolean skipUnavailable = remoteClusterService.isSkipUnavailable(clusterAlias);
            OriginalIndices indices = entry.getValue();
            SearchRequest ccsSearchRequest = SearchRequest.subSearchRequest(searchRequest, indices.indices(), clusterAlias, timeProvider.getAbsoluteStartMillis(), false);
            ActionListener<SearchResponse> ccsListener = createCCSListener(clusterAlias, skipUnavailable, countDown, skippedClusters, exceptions, searchResponseMerger, totalClusters, listener);
            Client remoteClusterClient = remoteClusterService.getRemoteClusterClient(threadPool, clusterAlias);
            remoteClusterClient.search(ccsSearchRequest, ccsListener);
        }
        if (localIndices != null) {
            ActionListener<SearchResponse> ccsListener = createCCSListener(RemoteClusterAware.LOCAL_CLUSTER_GROUP_KEY, false, countDown, skippedClusters, exceptions, searchResponseMerger, totalClusters, listener);
            SearchRequest ccsLocalSearchRequest = SearchRequest.subSearchRequest(searchRequest, localIndices.indices(), RemoteClusterAware.LOCAL_CLUSTER_GROUP_KEY, timeProvider.getAbsoluteStartMillis(), false);
            localSearchConsumer.accept(ccsLocalSearchRequest, ccsListener);
        }
    }
}

Example 18 with CountDown

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

the class TransportBroadcastReplicationAction method doExecute.

@Override
protected void doExecute(Task task, Request request, ActionListener<Response> listener) {
    final ClusterState clusterState = clusterService.state();
    List<ShardId> shards = shards(request, clusterState);
    final CopyOnWriteArrayList<ShardResponse> shardsResponses = new CopyOnWriteArrayList<>();
    if (shards.size() == 0) {
        finishAndNotifyListener(listener, shardsResponses);
    }
    final CountDown responsesCountDown = new CountDown(shards.size());
    for (final ShardId shardId : shards) {
        ActionListener<ShardResponse> shardActionListener = new ActionListener<ShardResponse>() {

            @Override
            public void onResponse(ShardResponse shardResponse) {
                shardsResponses.add(shardResponse);
                logger.trace("{}: got response from {}", actionName, shardId);
                if (responsesCountDown.countDown()) {
                    finishAndNotifyListener(listener, shardsResponses);
                }
            }

            @Override
            public void onFailure(Exception e) {
                logger.trace("{}: got failure from {}", actionName, shardId);
                int totalNumCopies = clusterState.getMetadata().getIndexSafe(shardId.getIndex()).getNumberOfReplicas() + 1;
                ShardResponse shardResponse = newShardResponse();
                ReplicationResponse.ShardInfo.Failure[] failures;
                if (TransportActions.isShardNotAvailableException(e)) {
                    failures = new ReplicationResponse.ShardInfo.Failure[0];
                } else {
                    ReplicationResponse.ShardInfo.Failure failure = new ReplicationResponse.ShardInfo.Failure(shardId, null, e, ExceptionsHelper.status(e), true);
                    failures = new ReplicationResponse.ShardInfo.Failure[totalNumCopies];
                    Arrays.fill(failures, failure);
                }
                shardResponse.setShardInfo(new ReplicationResponse.ShardInfo(totalNumCopies, 0, failures));
                shardsResponses.add(shardResponse);
                if (responsesCountDown.countDown()) {
                    finishAndNotifyListener(listener, shardsResponses);
                }
            }
        };
        shardExecute(task, request, shardId, shardActionListener);
    }
}