Examples with TransportShardUpsertAction io.crate.execution.dml.upsert.TransportShardUpsertAction used on opensource projects

Example 1 with TransportShardUpsertAction

use of io.crate.execution.dml.upsert.TransportShardUpsertAction in project crate by crate.

the class SQLIntegrationTestCase method assertNoTasksAreLeftOpen.

@After
public void assertNoTasksAreLeftOpen() throws Exception {
    final Field activeTasks = TasksService.class.getDeclaredField("activeTasks");
    final Field activeOperationsSb = TransportShardAction.class.getDeclaredField("activeOperations");
    activeTasks.setAccessible(true);
    activeOperationsSb.setAccessible(true);
    try {
        assertBusy(() -> {
            for (TasksService tasksService : internalCluster().getInstances(TasksService.class)) {
                try {
                    // noinspection unchecked
                    Map<UUID, RootTask> contexts = (Map<UUID, RootTask>) activeTasks.get(tasksService);
                    assertThat(contexts.size(), is(0));
                } catch (IllegalAccessException e) {
                    throw new RuntimeException(e);
                }
            }
            for (TransportShardUpsertAction action : internalCluster().getInstances(TransportShardUpsertAction.class)) {
                try {
                    ConcurrentHashMap<TaskId, KillableCallable<?>> operations = (ConcurrentHashMap<TaskId, KillableCallable<?>>) activeOperationsSb.get(action);
                    assertThat(operations.size(), is(0));
                } catch (IllegalAccessException e) {
                    throw new RuntimeException(e);
                }
            }
            for (TransportShardDeleteAction action : internalCluster().getInstances(TransportShardDeleteAction.class)) {
                try {
                    ConcurrentHashMap<TaskId, KillableCallable<?>> operations = (ConcurrentHashMap<TaskId, KillableCallable<?>>) activeOperationsSb.get(action);
                    assertThat(operations.size(), is(0));
                } catch (IllegalAccessException e) {
                    throw new RuntimeException(e);
                }
            }
        }, 10L, TimeUnit.SECONDS);
    } catch (AssertionError e) {
        StringBuilder errorMessageBuilder = new StringBuilder();
        String[] nodeNames = internalCluster().getNodeNames();
        for (String nodeName : nodeNames) {
            TasksService tasksService = internalCluster().getInstance(TasksService.class, nodeName);
            try {
                // noinspection unchecked
                Map<UUID, RootTask> contexts = (Map<UUID, RootTask>) activeTasks.get(tasksService);
                String contextsString = contexts.toString();
                if (!"{}".equals(contextsString)) {
                    errorMessageBuilder.append("## node: ");
                    errorMessageBuilder.append(nodeName);
                    errorMessageBuilder.append("\n");
                    errorMessageBuilder.append(contextsString);
                    errorMessageBuilder.append("\n");
                }
                contexts.clear();
            } catch (IllegalAccessException ex) {
                throw new RuntimeException(ex);
            }
        }
        throw new AssertionError(errorMessageBuilder.toString(), e);
    }
}

Example 2 with TransportShardUpsertAction

use of io.crate.execution.dml.upsert.TransportShardUpsertAction in project crate by crate.

the class InsertFromValues method execute.

private CompletableFuture<ShardResponse.CompressedResult> execute(NodeLimits nodeLimits, ClusterState state, Collection<ShardUpsertRequest> shardUpsertRequests, TransportShardUpsertAction shardUpsertAction, ScheduledExecutorService scheduler) {
    ShardResponse.CompressedResult compressedResult = new ShardResponse.CompressedResult();
    if (shardUpsertRequests.isEmpty()) {
        return CompletableFuture.completedFuture(compressedResult);
    }
    CompletableFuture<ShardResponse.CompressedResult> result = new CompletableFuture<>();
    AtomicInteger numRequests = new AtomicInteger(shardUpsertRequests.size());
    AtomicReference<Throwable> lastFailure = new AtomicReference<>(null);
    Consumer<ShardUpsertRequest> countdown = request -> {
        if (numRequests.decrementAndGet() == 0) {
            Throwable throwable = lastFailure.get();
            if (throwable == null) {
                result.complete(compressedResult);
            } else {
                throwable = SQLExceptions.unwrap(throwable, t -> t instanceof RuntimeException);
                // we want to report duplicate key exceptions
                if (!SQLExceptions.isDocumentAlreadyExistsException(throwable) && (partitionWasDeleted(throwable, request.index()) || partitionClosed(throwable, request.index()) || mixedArgumentTypesFailure(throwable))) {
                    result.complete(compressedResult);
                } else {
                    result.completeExceptionally(throwable);
                }
            }
        }
    };
    for (ShardUpsertRequest request : shardUpsertRequests) {
        String nodeId;
        try {
            nodeId = state.routingTable().shardRoutingTable(request.shardId()).primaryShard().currentNodeId();
        } catch (IndexNotFoundException e) {
            lastFailure.set(e);
            if (!IndexParts.isPartitioned(request.index())) {
                synchronized (compressedResult) {
                    compressedResult.markAsFailed(request.items());
                }
            }
            countdown.accept(request);
            continue;
        }
        final ConcurrencyLimit nodeLimit = nodeLimits.get(nodeId);
        final long startTime = nodeLimit.startSample();
        ActionListener<ShardResponse> listener = new ActionListener<>() {

            @Override
            public void onResponse(ShardResponse shardResponse) {
                Throwable throwable = shardResponse.failure();
                if (throwable == null) {
                    nodeLimit.onSample(startTime, false);
                    synchronized (compressedResult) {
                        compressedResult.update(shardResponse);
                    }
                } else {
                    nodeLimit.onSample(startTime, true);
                    lastFailure.set(throwable);
                }
                countdown.accept(request);
            }

            @Override
            public void onFailure(Exception e) {
                nodeLimit.onSample(startTime, true);
                Throwable t = SQLExceptions.unwrap(e);
                if (!partitionWasDeleted(t, request.index())) {
                    synchronized (compressedResult) {
                        compressedResult.markAsFailed(request.items());
                    }
                }
                lastFailure.set(t);
                countdown.accept(request);
            }
        };
        shardUpsertAction.execute(request, new RetryListener<>(scheduler, l -> shardUpsertAction.execute(request, l), listener, BackoffPolicy.limitedDynamic(nodeLimit)));
    }
    return result;
}