Examples with DataSize io.airlift.units.DataSize used on opensource projects

Example 26 with DataSize

use of io.airlift.units.DataSize in project presto by prestodb.

the class TestSqlTask method createInitialTask.

public SqlTask createInitialTask() {
    TaskId taskId = new TaskId("query", 0, nextTaskId.incrementAndGet());
    URI location = URI.create("fake://task/" + taskId);
    return new SqlTask(taskId, location, new QueryContext(new QueryId("query"), new DataSize(1, MEGABYTE), new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)), new MemoryPool(new MemoryPoolId("testSystem"), new DataSize(1, GIGABYTE)), taskNotificationExecutor), sqlTaskExecutionFactory, taskNotificationExecutor, Functions.identity(), new DataSize(32, MEGABYTE));
}

Example 27 with DataSize

use of io.airlift.units.DataSize in project presto by prestodb.

the class TestSqlTaskManager method testSimpleQuery.

@Test
public void testSimpleQuery() throws Exception {
    try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) {
        TaskId taskId = TASK_ID;
        TaskInfo taskInfo = sqlTaskManager.updateTask(TEST_SESSION, taskId, Optional.of(PLAN_FRAGMENT), ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(SPLIT), true)), createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds());
        assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING);
        taskInfo = sqlTaskManager.getTaskInfo(taskId);
        assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING);
        BufferResult results = sqlTaskManager.getTaskResults(taskId, OUT, 0, new DataSize(1, Unit.MEGABYTE)).get();
        assertEquals(results.isBufferComplete(), false);
        assertEquals(results.getSerializedPages().size(), 1);
        assertEquals(results.getSerializedPages().get(0).getPositionCount(), 1);
        for (boolean moreResults = true; moreResults; moreResults = !results.isBufferComplete()) {
            results = sqlTaskManager.getTaskResults(taskId, OUT, results.getToken() + results.getSerializedPages().size(), new DataSize(1, Unit.MEGABYTE)).get();
        }
        assertEquals(results.isBufferComplete(), true);
        assertEquals(results.getSerializedPages().size(), 0);
        // complete the task by calling abort on it
        TaskInfo info = sqlTaskManager.abortTaskResults(taskId, OUT);
        assertEquals(info.getOutputBuffers().getState(), BufferState.FINISHED);
        taskInfo = sqlTaskManager.getTaskInfo(taskId, taskInfo.getTaskStatus().getState()).get(1, TimeUnit.SECONDS);
        assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED);
        taskInfo = sqlTaskManager.getTaskInfo(taskId);
        assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED);
    }
}

Example 28 with DataSize

use of io.airlift.units.DataSize in project presto by prestodb.

the class TestArbitraryOutputBuffer method testInvalidConstructorArg.

@Test
public void testInvalidConstructorArg() throws Exception {
    try {
        createArbitraryBuffer(createInitialEmptyOutputBuffers(ARBITRARY).withBuffer(FIRST, BROADCAST_PARTITION_ID).withNoMoreBufferIds(), new DataSize(0, BYTE));
        fail("Expected IllegalStateException");
    } catch (IllegalArgumentException ignored) {
    }
    try {
        createArbitraryBuffer(createInitialEmptyOutputBuffers(ARBITRARY), new DataSize(0, BYTE));
        fail("Expected IllegalStateException");
    } catch (IllegalArgumentException ignored) {
    }
}

Example 29 with DataSize

use of io.airlift.units.DataSize in project presto by prestodb.

the class TestExchangeClient method testClose.

@Test
public void testClose() throws Exception {
    DataSize maxResponseSize = new DataSize(1, Unit.BYTE);
    MockExchangeRequestProcessor processor = new MockExchangeRequestProcessor(maxResponseSize);
    URI location = URI.create("http://localhost:8080");
    processor.addPage(location, createPage(1));
    processor.addPage(location, createPage(2));
    processor.addPage(location, createPage(3));
    @SuppressWarnings("resource") ExchangeClient exchangeClient = new ExchangeClient(new DataSize(1, Unit.BYTE), maxResponseSize, 1, new Duration(1, TimeUnit.MINUTES), new Duration(1, TimeUnit.MINUTES), new TestingHttpClient(processor, newCachedThreadPool(daemonThreadsNamed("test-%s"))), executor, deltaMemoryInBytes -> {
    });
    exchangeClient.addLocation(location);
    exchangeClient.noMoreLocations();
    // fetch a page
    assertEquals(exchangeClient.isClosed(), false);
    assertPageEquals(exchangeClient.getNextPage(new Duration(1, TimeUnit.SECONDS)), createPage(1));
    // close client while pages are still available
    exchangeClient.close();
    while (!exchangeClient.isFinished()) {
        MILLISECONDS.sleep(10);
    }
    assertEquals(exchangeClient.isClosed(), true);
    assertNull(exchangeClient.getNextPage(new Duration(0, TimeUnit.SECONDS)));
    assertEquals(exchangeClient.getStatus().getBufferedPages(), 0);
    assertEquals(exchangeClient.getStatus().getBufferedBytes(), 0);
    // client should have sent only 2 requests: one to get all pages and once to get the done signal
    PageBufferClientStatus clientStatus = exchangeClient.getStatus().getPageBufferClientStatuses().get(0);
    assertEquals(clientStatus.getUri(), location);
    assertEquals(clientStatus.getState(), "closed", "status");
    assertEquals(clientStatus.getHttpRequestState(), "not scheduled", "httpRequestState");
}

Example 30 with DataSize

use of io.airlift.units.DataSize in project presto by prestodb.

the class TestExchangeClient method testHappyPath.

@Test
public void testHappyPath() throws Exception {
    DataSize maxResponseSize = new DataSize(10, Unit.MEGABYTE);
    MockExchangeRequestProcessor processor = new MockExchangeRequestProcessor(maxResponseSize);
    URI location = URI.create("http://localhost:8080");
    processor.addPage(location, createPage(1));
    processor.addPage(location, createPage(2));
    processor.addPage(location, createPage(3));
    processor.setComplete(location);
    @SuppressWarnings("resource") ExchangeClient exchangeClient = new ExchangeClient(new DataSize(32, Unit.MEGABYTE), maxResponseSize, 1, new Duration(1, TimeUnit.MINUTES), new Duration(1, TimeUnit.MINUTES), new TestingHttpClient(processor, executor), executor, deltaMemoryInBytes -> {
    });
    exchangeClient.addLocation(location);
    exchangeClient.noMoreLocations();
    assertEquals(exchangeClient.isClosed(), false);
    assertPageEquals(exchangeClient.getNextPage(new Duration(1, TimeUnit.SECONDS)), createPage(1));
    assertEquals(exchangeClient.isClosed(), false);
    assertPageEquals(exchangeClient.getNextPage(new Duration(1, TimeUnit.SECONDS)), createPage(2));
    assertEquals(exchangeClient.isClosed(), false);
    assertPageEquals(exchangeClient.getNextPage(new Duration(1, TimeUnit.SECONDS)), createPage(3));
    assertNull(exchangeClient.getNextPage(new Duration(1, TimeUnit.SECONDS)));
    assertEquals(exchangeClient.isClosed(), true);
    if (exchangeClient.getStatus().getBufferedPages() != 0) {
        assertEquals(exchangeClient.getStatus().getBufferedPages(), 0);
    }
    assertTrue(exchangeClient.getStatus().getBufferedBytes() == 0);
    // client should have sent only 2 requests: one to get all pages and once to get the done signal
    assertStatus(exchangeClient.getStatus().getPageBufferClientStatuses().get(0), location, "closed", 3, 3, 3, "not scheduled");
}