Examples with PerDatasourceShuffleMetrics org.apache.druid.indexing.worker.shuffle.ShuffleMetrics.PerDatasourceShuffleMetrics used on opensource projects

Example 6 with PerDatasourceShuffleMetrics

use of org.apache.druid.indexing.worker.shuffle.ShuffleMetrics.PerDatasourceShuffleMetrics in project druid by druid-io.

the class ShuffleMetricsTest method testConcurrency.

@Test(timeout = 5000L)
public void testConcurrency() throws ExecutionException, InterruptedException {
    // 2 for write, 1 for read
    final ExecutorService exec = Execs.multiThreaded(3, "shuffle-metrics-test-%d");
    try {
        final ShuffleMetrics metrics = new ShuffleMetrics();
        final String supervisorTask1 = "supervisor1";
        final String supervisorTask2 = "supervisor2";
        final CountDownLatch firstUpdatelatch = new CountDownLatch(2);
        final List<Future<Void>> futures = new ArrayList<>();
        futures.add(exec.submit(() -> {
            metrics.shuffleRequested(supervisorTask1, 1024);
            metrics.shuffleRequested(supervisorTask2, 30);
            firstUpdatelatch.countDown();
            Thread.sleep(ThreadLocalRandom.current().nextInt(10));
            metrics.shuffleRequested(supervisorTask2, 10);
            return null;
        }));
        futures.add(exec.submit(() -> {
            metrics.shuffleRequested(supervisorTask2, 30);
            metrics.shuffleRequested(supervisorTask1, 1024);
            firstUpdatelatch.countDown();
            Thread.sleep(ThreadLocalRandom.current().nextInt(10));
            metrics.shuffleRequested(supervisorTask1, 32);
            return null;
        }));
        final Map<String, PerDatasourceShuffleMetrics> firstSnapshot = exec.submit(() -> {
            firstUpdatelatch.await();
            Thread.sleep(ThreadLocalRandom.current().nextInt(10));
            return metrics.snapshotAndReset();
        }).get();
        int expectedSecondSnapshotSize = 0;
        boolean task1ShouldBeInSecondSnapshot = false;
        boolean task2ShouldBeInSecondSnapshot = false;
        Assert.assertEquals(2, firstSnapshot.size());
        Assert.assertNotNull(firstSnapshot.get(supervisorTask1));
        Assert.assertTrue(2048 == firstSnapshot.get(supervisorTask1).getShuffleBytes() || 2080 == firstSnapshot.get(supervisorTask1).getShuffleBytes());
        Assert.assertTrue(2 == firstSnapshot.get(supervisorTask1).getShuffleRequests() || 3 == firstSnapshot.get(supervisorTask1).getShuffleRequests());
        if (firstSnapshot.get(supervisorTask1).getShuffleRequests() == 2) {
            expectedSecondSnapshotSize++;
            task1ShouldBeInSecondSnapshot = true;
        }
        Assert.assertNotNull(firstSnapshot.get(supervisorTask2));
        Assert.assertTrue(60 == firstSnapshot.get(supervisorTask2).getShuffleBytes() || 70 == firstSnapshot.get(supervisorTask2).getShuffleBytes());
        Assert.assertTrue(2 == firstSnapshot.get(supervisorTask2).getShuffleRequests() || 3 == firstSnapshot.get(supervisorTask2).getShuffleRequests());
        if (firstSnapshot.get(supervisorTask2).getShuffleRequests() == 2) {
            expectedSecondSnapshotSize++;
            task2ShouldBeInSecondSnapshot = true;
        }
        for (Future<Void> future : futures) {
            future.get();
        }
        final Map<String, PerDatasourceShuffleMetrics> secondSnapshot = metrics.snapshotAndReset();
        Assert.assertEquals(expectedSecondSnapshotSize, secondSnapshot.size());
        Assert.assertEquals(task1ShouldBeInSecondSnapshot, secondSnapshot.containsKey(supervisorTask1));
        if (task1ShouldBeInSecondSnapshot) {
            Assert.assertEquals(32, secondSnapshot.get(supervisorTask1).getShuffleBytes());
            Assert.assertEquals(1, secondSnapshot.get(supervisorTask1).getShuffleRequests());
        }
        Assert.assertEquals(task2ShouldBeInSecondSnapshot, secondSnapshot.containsKey(supervisorTask2));
        if (task2ShouldBeInSecondSnapshot) {
            Assert.assertEquals(10, secondSnapshot.get(supervisorTask2).getShuffleBytes());
            Assert.assertEquals(1, secondSnapshot.get(supervisorTask2).getShuffleRequests());
        }
    } finally {
        exec.shutdown();
    }
}