Examples with FakeTcpChannel org.opensearch.transport.FakeTcpChannel used on opensource projects

Example 1 with FakeTcpChannel

use of org.opensearch.transport.FakeTcpChannel in project OpenSearch by opensearch-project.

the class TaskManagerTests method testTrackingChannelTask.

public void testTrackingChannelTask() throws Exception {
    final TaskManager taskManager = new TaskManager(Settings.EMPTY, threadPool, Collections.emptySet());
    Set<Task> cancelledTasks = ConcurrentCollections.newConcurrentSet();
    taskManager.setTaskCancellationService(new TaskCancellationService(mock(TransportService.class)) {

        @Override
        void cancelTaskAndDescendants(CancellableTask task, String reason, boolean waitForCompletion, ActionListener<Void> listener) {
            assertThat(reason, equalTo("channel was closed"));
            assertFalse(waitForCompletion);
            assertTrue("task [" + task + "] was cancelled already", cancelledTasks.add(task));
        }
    });
    Map<TcpChannel, Set<Task>> pendingTasks = new HashMap<>();
    Set<Task> expectedCancelledTasks = new HashSet<>();
    FakeTcpChannel[] channels = new FakeTcpChannel[randomIntBetween(1, 10)];
    List<Releasable> stopTrackingTasks = new ArrayList<>();
    for (int i = 0; i < channels.length; i++) {
        channels[i] = new SingleThreadedTcpChannel();
    }
    int iterations = randomIntBetween(1, 200);
    for (int i = 0; i < iterations; i++) {
        final List<Releasable> subset = randomSubsetOf(stopTrackingTasks);
        stopTrackingTasks.removeAll(subset);
        Releasables.close(subset);
        final FakeTcpChannel channel = randomFrom(channels);
        final Task task = taskManager.register("transport", "test", new CancellableRequest(Integer.toString(i)));
        if (channel.isOpen() && randomBoolean()) {
            channel.close();
            expectedCancelledTasks.addAll(pendingTasks.getOrDefault(channel, Collections.emptySet()));
        }
        final Releasable stopTracking = taskManager.startTrackingCancellableChannelTask(channel, (CancellableTask) task);
        if (channel.isOpen()) {
            pendingTasks.computeIfAbsent(channel, k -> new HashSet<>()).add(task);
            stopTrackingTasks.add(() -> {
                stopTracking.close();
                assertTrue(pendingTasks.get(channel).remove(task));
                expectedCancelledTasks.remove(task);
            });
        } else {
            expectedCancelledTasks.add(task);
        }
    }
    assertBusy(() -> assertThat(expectedCancelledTasks, everyItem(in(cancelledTasks))), 30, TimeUnit.SECONDS);
    for (FakeTcpChannel channel : channels) {
        channel.close();
    }
    assertThat(taskManager.numberOfChannelPendingTaskTrackers(), equalTo(0));
}

Example 2 with FakeTcpChannel

use of org.opensearch.transport.FakeTcpChannel in project OpenSearch by opensearch-project.

the class TaskManagerTests method testTrackingTaskAndCloseChannelConcurrently.

public void testTrackingTaskAndCloseChannelConcurrently() throws Exception {
    final TaskManager taskManager = new TaskManager(Settings.EMPTY, threadPool, Collections.emptySet());
    Set<CancellableTask> cancelledTasks = ConcurrentCollections.newConcurrentSet();
    taskManager.setTaskCancellationService(new TaskCancellationService(mock(TransportService.class)) {

        @Override
        void cancelTaskAndDescendants(CancellableTask task, String reason, boolean waitForCompletion, ActionListener<Void> listener) {
            assertTrue("task [" + task + "] was cancelled already", cancelledTasks.add(task));
        }
    });
    Set<Task> expectedCancelledTasks = ConcurrentCollections.newConcurrentSet();
    FakeTcpChannel[] channels = new FakeTcpChannel[randomIntBetween(1, 10)];
    for (int i = 0; i < channels.length; i++) {
        channels[i] = new FakeTcpChannel();
    }
    Thread[] threads = new Thread[randomIntBetween(2, 8)];
    Phaser phaser = new Phaser(threads.length);
    for (int t = 0; t < threads.length; t++) {
        String threadName = "thread-" + t;
        threads[t] = new Thread(() -> {
            phaser.arriveAndAwaitAdvance();
            int iterations = randomIntBetween(50, 500);
            for (int i = 0; i < iterations; i++) {
                final FakeTcpChannel channel = randomFrom(channels);
                final Task task = taskManager.register("transport", "test", new CancellableRequest(threadName + ":" + i));
                expectedCancelledTasks.add(task);
                taskManager.startTrackingCancellableChannelTask(channel, (CancellableTask) task);
                if (randomInt(100) < 5) {
                    randomFrom(channels).close();
                }
            }
        });
        threads[t].start();
    }
    for (FakeTcpChannel channel : channels) {
        channel.close();
    }
    for (Thread thread : threads) {
        thread.join();
    }
    assertBusy(() -> assertThat(cancelledTasks, equalTo(expectedCancelledTasks)), 1, TimeUnit.MINUTES);
    assertThat(taskManager.numberOfChannelPendingTaskTrackers(), equalTo(0));
}