Examples with RestClient org.opensearch.client.RestClient used on opensource projects

Example 26 with RestClient

use of org.opensearch.client.RestClient in project OpenSearch by opensearch-project.

the class SnifferTests method testSniff.

/**
 * Tests the {@link Sniffer#sniff()} method in isolation. Verifies that it uses the {@link NodesSniffer} implementation
 * to retrieve nodes and set them (when not empty) to the provided {@link RestClient} instance.
 */
public void testSniff() throws IOException {
    Node initialNode = new Node(new HttpHost("localhost", 9200));
    try (RestClient restClient = RestClient.builder(initialNode).build()) {
        Scheduler noOpScheduler = new Scheduler() {

            @Override
            public Future<?> schedule(Sniffer.Task task, long delayMillis) {
                return mock(Future.class);
            }

            @Override
            public void shutdown() {
            }
        };
        CountingNodesSniffer nodesSniffer = new CountingNodesSniffer();
        int iters = randomIntBetween(5, 30);
        try (Sniffer sniffer = new Sniffer(restClient, nodesSniffer, noOpScheduler, 1000L, -1)) {
            {
                assertEquals(1, restClient.getNodes().size());
                Node node = restClient.getNodes().get(0);
                assertEquals("localhost", node.getHost().getHostName());
                assertEquals(9200, node.getHost().getPort());
            }
            int emptyList = 0;
            int failures = 0;
            int runs = 0;
            List<Node> lastNodes = Collections.singletonList(initialNode);
            for (int i = 0; i < iters; i++) {
                try {
                    runs++;
                    sniffer.sniff();
                    if (nodesSniffer.failures.get() > failures) {
                        failures++;
                        fail("should have failed given that nodesSniffer says it threw an exception");
                    } else if (nodesSniffer.emptyList.get() > emptyList) {
                        emptyList++;
                        assertEquals(lastNodes, restClient.getNodes());
                    } else {
                        assertNotEquals(lastNodes, restClient.getNodes());
                        List<Node> expectedNodes = CountingNodesSniffer.buildNodes(runs);
                        assertEquals(expectedNodes, restClient.getNodes());
                        lastNodes = restClient.getNodes();
                    }
                } catch (IOException e) {
                    if (nodesSniffer.failures.get() > failures) {
                        failures++;
                        assertEquals("communication breakdown", e.getMessage());
                    }
                }
            }
            assertEquals(nodesSniffer.emptyList.get(), emptyList);
            assertEquals(nodesSniffer.failures.get(), failures);
            assertEquals(nodesSniffer.runs.get(), runs);
        }
    }
}

Example 27 with RestClient

use of org.opensearch.client.RestClient in project OpenSearch by opensearch-project.

the class SnifferTests method testSniffOnFailure.

/**
 * Test behaviour when a bunch of onFailure sniffing rounds are triggered in parallel. Each run will always
 * schedule a subsequent afterFailure round. Also, for each onFailure round that starts, the net scheduled round
 * (either afterFailure or ordinary) gets cancelled.
 */
public void testSniffOnFailure() throws Exception {
    RestClient restClient = mock(RestClient.class);
    CountingNodesSniffer nodesSniffer = new CountingNodesSniffer();
    final AtomicBoolean initializing = new AtomicBoolean(true);
    final long sniffInterval = randomLongBetween(1, Long.MAX_VALUE);
    final long sniffAfterFailureDelay = randomLongBetween(1, Long.MAX_VALUE);
    int minNumOnFailureRounds = randomIntBetween(5, 10);
    final CountDownLatch initializingLatch = new CountDownLatch(1);
    final Set<Sniffer.ScheduledTask> ordinaryRoundsTasks = new CopyOnWriteArraySet<>();
    final AtomicReference<Future<?>> initializingFuture = new AtomicReference<>();
    final Set<Sniffer.ScheduledTask> onFailureTasks = new CopyOnWriteArraySet<>();
    final Set<Sniffer.ScheduledTask> afterFailureTasks = new CopyOnWriteArraySet<>();
    final AtomicBoolean onFailureCompleted = new AtomicBoolean(false);
    final CountDownLatch completionLatch = new CountDownLatch(1);
    final ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
    try {
        Scheduler scheduler = new Scheduler() {

            @Override
            public Future<?> schedule(final Sniffer.Task task, long delayMillis) {
                if (initializing.compareAndSet(true, false)) {
                    assertEquals(0L, delayMillis);
                    Future<?> future = executor.submit(new Runnable() {

                        @Override
                        public void run() {
                            try {
                                task.run();
                            } finally {
                                // we need to make sure that the sniffer is initialized, so the sniffOnFailure
                                // call does what it needs to do. Otherwise nothing happens until initialized.
                                initializingLatch.countDown();
                            }
                        }
                    });
                    assertTrue(initializingFuture.compareAndSet(null, future));
                    return future;
                }
                if (delayMillis == 0L) {
                    Future<?> future = executor.submit(task);
                    onFailureTasks.add(new Sniffer.ScheduledTask(task, future));
                    return future;
                }
                if (delayMillis == sniffAfterFailureDelay) {
                    Future<?> future = scheduleOrSubmit(task);
                    afterFailureTasks.add(new Sniffer.ScheduledTask(task, future));
                    return future;
                }
                assertEquals(sniffInterval, delayMillis);
                assertEquals(sniffInterval, task.nextTaskDelay);
                if (onFailureCompleted.get() && onFailureTasks.size() == afterFailureTasks.size()) {
                    completionLatch.countDown();
                    return mock(Future.class);
                }
                Future<?> future = scheduleOrSubmit(task);
                ordinaryRoundsTasks.add(new Sniffer.ScheduledTask(task, future));
                return future;
            }

            private Future<?> scheduleOrSubmit(Sniffer.Task task) {
                if (randomBoolean()) {
                    return executor.schedule(task, randomLongBetween(0L, 200L), TimeUnit.MILLISECONDS);
                } else {
                    return executor.submit(task);
                }
            }

            @Override
            public void shutdown() {
            }
        };
        final Sniffer sniffer = new Sniffer(restClient, nodesSniffer, scheduler, sniffInterval, sniffAfterFailureDelay);
        assertTrue("timeout waiting for sniffer to get initialized", initializingLatch.await(1000, TimeUnit.MILLISECONDS));
        ExecutorService onFailureExecutor = Executors.newFixedThreadPool(randomIntBetween(5, 20));
        Set<Future<?>> onFailureFutures = new CopyOnWriteArraySet<>();
        try {
            // as another round is already running. We retry till enough runs get through as that's what we want to test.
            while (onFailureTasks.size() < minNumOnFailureRounds) {
                onFailureFutures.add(onFailureExecutor.submit(new Runnable() {

                    @Override
                    public void run() {
                        sniffer.sniffOnFailure();
                    }
                }));
            }
            assertThat(onFailureFutures.size(), greaterThanOrEqualTo(minNumOnFailureRounds));
            for (Future<?> onFailureFuture : onFailureFutures) {
                assertNull(onFailureFuture.get());
            }
            onFailureCompleted.set(true);
        } finally {
            onFailureExecutor.shutdown();
            onFailureExecutor.awaitTermination(1000, TimeUnit.MILLISECONDS);
        }
        assertFalse(initializingFuture.get().isCancelled());
        assertTrue(initializingFuture.get().isDone());
        assertNull(initializingFuture.get().get());
        assertTrue("timeout waiting for sniffing rounds to be completed", completionLatch.await(1000, TimeUnit.MILLISECONDS));
        assertThat(onFailureTasks.size(), greaterThanOrEqualTo(minNumOnFailureRounds));
        assertEquals(onFailureTasks.size(), afterFailureTasks.size());
        for (Sniffer.ScheduledTask onFailureTask : onFailureTasks) {
            assertFalse(onFailureTask.future.isCancelled());
            assertTrue(onFailureTask.future.isDone());
            assertNull(onFailureTask.future.get());
            assertTrue(onFailureTask.task.hasStarted());
            assertFalse(onFailureTask.task.isSkipped());
        }
        int cancelledTasks = 0;
        int completedTasks = onFailureTasks.size() + 1;
        for (Sniffer.ScheduledTask afterFailureTask : afterFailureTasks) {
            if (assertTaskCancelledOrCompleted(afterFailureTask)) {
                completedTasks++;
            } else {
                cancelledTasks++;
            }
        }
        assertThat(ordinaryRoundsTasks.size(), greaterThan(0));
        for (Sniffer.ScheduledTask task : ordinaryRoundsTasks) {
            if (assertTaskCancelledOrCompleted(task)) {
                completedTasks++;
            } else {
                cancelledTasks++;
            }
        }
        assertEquals(onFailureTasks.size(), cancelledTasks);
        assertEquals(completedTasks, nodesSniffer.runs.get());
        int setNodesRuns = nodesSniffer.runs.get() - nodesSniffer.failures.get() - nodesSniffer.emptyList.get();
        verify(restClient, times(setNodesRuns)).setNodes(anyCollection());
        verifyNoMoreInteractions(restClient);
    } finally {
        executor.shutdown();
        executor.awaitTermination(1000L, TimeUnit.MILLISECONDS);
    }
}

Example 28 with RestClient

use of org.opensearch.client.RestClient in project OpenSearch by opensearch-project.

the class SnifferTests method testSniffOnFailureNotInitialized.

public void testSniffOnFailureNotInitialized() {
    RestClient restClient = mock(RestClient.class);
    CountingNodesSniffer nodesSniffer = new CountingNodesSniffer();
    long sniffInterval = randomLongBetween(1, Long.MAX_VALUE);
    long sniffAfterFailureDelay = randomLongBetween(1, Long.MAX_VALUE);
    final AtomicInteger scheduleCalls = new AtomicInteger(0);
    Scheduler scheduler = new Scheduler() {

        @Override
        public Future<?> schedule(Sniffer.Task task, long delayMillis) {
            scheduleCalls.incrementAndGet();
            return null;
        }

        @Override
        public void shutdown() {
        }
    };
    Sniffer sniffer = new Sniffer(restClient, nodesSniffer, scheduler, sniffInterval, sniffAfterFailureDelay);
    for (int i = 0; i < 10; i++) {
        sniffer.sniffOnFailure();
    }
    assertEquals(1, scheduleCalls.get());
    int totalRuns = nodesSniffer.runs.get();
    assertEquals(0, totalRuns);
    int setNodesRuns = totalRuns - nodesSniffer.failures.get() - nodesSniffer.emptyList.get();
    verify(restClient, times(setNodesRuns)).setNodes(anyCollection());
    verifyNoMoreInteractions(restClient);
}

Example 29 with RestClient

use of org.opensearch.client.RestClient in project OpenSearch by opensearch-project.

the class SnifferTests method testDefaultSchedulerSchedule.

public void testDefaultSchedulerSchedule() {
    RestClient restClient = mock(RestClient.class);
    NodesSniffer nodesSniffer = mock(NodesSniffer.class);
    Scheduler noOpScheduler = new Scheduler() {

        @Override
        public Future<?> schedule(Sniffer.Task task, long delayMillis) {
            return mock(Future.class);
        }

        @Override
        public void shutdown() {
        }
    };
    Sniffer sniffer = new Sniffer(restClient, nodesSniffer, noOpScheduler, 0L, 0L);
    Sniffer.Task task = sniffer.new Task(randomLongBetween(1, Long.MAX_VALUE));
    ScheduledExecutorService scheduledExecutorService = mock(ScheduledExecutorService.class);
    final ScheduledFuture<?> mockedFuture = mock(ScheduledFuture.class);
    when(scheduledExecutorService.schedule(any(Runnable.class), any(Long.class), any(TimeUnit.class))).then(new Answer<ScheduledFuture<?>>() {

        @Override
        public ScheduledFuture<?> answer(InvocationOnMock invocationOnMock) {
            return mockedFuture;
        }
    });
    DefaultScheduler scheduler = new DefaultScheduler(scheduledExecutorService);
    long delay = randomLongBetween(1, Long.MAX_VALUE);
    Future<?> future = scheduler.schedule(task, delay);
    assertSame(mockedFuture, future);
    verify(scheduledExecutorService).schedule(task, delay, TimeUnit.MILLISECONDS);
    verifyNoMoreInteractions(scheduledExecutorService, mockedFuture);
}

Example 30 with RestClient

use of org.opensearch.client.RestClient in project OpenSearch by opensearch-project.

the class SnifferTests method testClose.

/**
 * Test that {@link Sniffer#close()} shuts down the underlying {@link Scheduler}, and that such calls are idempotent.
 * Also verifies that the next scheduled round gets cancelled.
 */
public void testClose() {
    final Future<?> future = mock(Future.class);
    long sniffInterval = randomLongBetween(1, Long.MAX_VALUE);
    long sniffAfterFailureDelay = randomLongBetween(1, Long.MAX_VALUE);
    RestClient restClient = mock(RestClient.class);
    final AtomicInteger shutdown = new AtomicInteger(0);
    final AtomicBoolean initialized = new AtomicBoolean(false);
    Scheduler scheduler = new Scheduler() {

        @Override
        public Future<?> schedule(Sniffer.Task task, long delayMillis) {
            if (initialized.compareAndSet(false, true)) {
                // run from the same thread so the sniffer gets for sure initialized and the scheduled task gets cancelled on close
                task.run();
            }
            return future;
        }

        @Override
        public void shutdown() {
            shutdown.incrementAndGet();
        }
    };
    Sniffer sniffer = new Sniffer(restClient, new MockNodesSniffer(), scheduler, sniffInterval, sniffAfterFailureDelay);
    assertEquals(0, shutdown.get());
    int iters = randomIntBetween(3, 10);
    for (int i = 1; i <= iters; i++) {
        sniffer.close();
        verify(future, times(i)).cancel(false);
        assertEquals(i, shutdown.get());
    }
}