Examples with QueryId io.prestosql.spi.QueryId used on opensource projects

Example 21 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestQueryContext method testSetMemoryPool.

@Test(dataProvider = "testSetMemoryPoolOptions")
public void testSetMemoryPool(boolean useReservedPool) {
    QueryId secondQuery = new QueryId("second");
    MemoryPool reservedPool = new MemoryPool(RESERVED_POOL, new DataSize(10, BYTE));
    long secondQueryMemory = reservedPool.getMaxBytes() - 1;
    if (useReservedPool) {
        assertTrue(reservedPool.reserve(secondQuery, "test", secondQueryMemory).isDone());
    }
    try (LocalQueryRunner localQueryRunner = new LocalQueryRunner(TEST_SESSION)) {
        QueryContext queryContext = new QueryContext(new QueryId("query"), new DataSize(10, BYTE), new DataSize(20, BYTE), new MemoryPool(GENERAL_POOL, new DataSize(10, BYTE)), new TestingGcMonitor(), localQueryRunner.getExecutor(), localQueryRunner.getScheduler(), new DataSize(0, BYTE), new SpillSpaceTracker(new DataSize(0, BYTE)), NOOP_SNAPSHOT_UTILS);
        // Use memory
        queryContext.getQueryMemoryContext().initializeLocalMemoryContexts("test");
        LocalMemoryContext userMemoryContext = queryContext.getQueryMemoryContext().localUserMemoryContext();
        LocalMemoryContext revocableMemoryContext = queryContext.getQueryMemoryContext().localRevocableMemoryContext();
        assertTrue(userMemoryContext.setBytes(3).isDone());
        assertTrue(revocableMemoryContext.setBytes(5).isDone());
        queryContext.setMemoryPool(reservedPool);
        if (useReservedPool) {
            reservedPool.free(secondQuery, "test", secondQueryMemory);
        }
        // Free memory
        userMemoryContext.close();
        revocableMemoryContext.close();
    }
}

Example 22 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestBasicQueryInfo method testConstructor.

@Test
public void testConstructor() {
    BasicQueryInfo basicInfo = new BasicQueryInfo(new QueryInfo(new QueryId("0"), TEST_SESSION.toSessionRepresentation(), RUNNING, new MemoryPoolId("reserved"), false, URI.create("1"), ImmutableList.of("2", "3"), "SELECT 4", Optional.empty(), new QueryStats(DateTime.parse("1991-09-06T05:00-05:30"), DateTime.parse("1991-09-06T05:01-05:30"), DateTime.parse("1991-09-06T05:02-05:30"), DateTime.parse("1991-09-06T06:00-05:30"), Duration.valueOf("8m"), Duration.valueOf("7m"), Duration.valueOf("34m"), Duration.valueOf("35m"), Duration.valueOf("44m"), Duration.valueOf("9m"), Duration.valueOf("10m"), Duration.valueOf("11m"), Duration.valueOf("12m"), Duration.valueOf("12m"), Duration.valueOf("12m"), 13, 14, 15, 16, 17, 18, 34, 19, 20.0, DataSize.valueOf("21GB"), DataSize.valueOf("22GB"), DataSize.valueOf("23GB"), DataSize.valueOf("24GB"), DataSize.valueOf("25GB"), DataSize.valueOf("26GB"), DataSize.valueOf("27GB"), DataSize.valueOf("28GB"), DataSize.valueOf("29GB"), true, Duration.valueOf("23m"), Duration.valueOf("24m"), Duration.valueOf("26m"), true, ImmutableSet.of(BlockedReason.WAITING_FOR_MEMORY), DataSize.valueOf("271GB"), 281, DataSize.valueOf("272GB"), 282, DataSize.valueOf("27GB"), 28, DataSize.valueOf("29GB"), 30, DataSize.valueOf("31GB"), 32, DataSize.valueOf("32GB"), ImmutableList.of(new StageGcStatistics(101, 102, 103, 104, 105, 106, 107)), ImmutableList.of()), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableMap.of(), ImmutableSet.of(), ImmutableMap.of(), ImmutableMap.of(), ImmutableSet.of(), Optional.empty(), false, "33", Optional.empty(), null, StandardErrorCode.ABANDONED_QUERY.toErrorCode(), ImmutableList.of(), ImmutableSet.of(), Optional.empty(), false, Optional.empty(), false));
    assertEquals(basicInfo.getQueryId().getId(), "0");
    assertEquals(basicInfo.getState(), RUNNING);
    assertEquals(basicInfo.getMemoryPool().getId(), "reserved");
    assertEquals(basicInfo.isScheduled(), false);
    assertEquals(basicInfo.getQuery(), "SELECT 4");
    assertEquals(basicInfo.getQueryStats().getCreateTime(), DateTime.parse("1991-09-06T05:00-05:30"));
    assertEquals(basicInfo.getQueryStats().getEndTime(), DateTime.parse("1991-09-06T06:00-05:30"));
    assertEquals(basicInfo.getQueryStats().getElapsedTime(), Duration.valueOf("8m"));
    assertEquals(basicInfo.getQueryStats().getExecutionTime(), Duration.valueOf("44m"));
    assertEquals(basicInfo.getQueryStats().getTotalDrivers(), 16);
    assertEquals(basicInfo.getQueryStats().getQueuedDrivers(), 17);
    assertEquals(basicInfo.getQueryStats().getRunningDrivers(), 18);
    assertEquals(basicInfo.getQueryStats().getCompletedDrivers(), 19);
    assertEquals(basicInfo.getQueryStats().getCumulativeUserMemory(), 20.0);
    assertEquals(basicInfo.getQueryStats().getUserMemoryReservation(), DataSize.valueOf("21GB"));
    assertEquals(basicInfo.getQueryStats().getTotalMemoryReservation(), DataSize.valueOf("23GB"));
    assertEquals(basicInfo.getQueryStats().getPeakUserMemoryReservation(), DataSize.valueOf("24GB"));
    assertEquals(basicInfo.getQueryStats().getTotalCpuTime(), Duration.valueOf("24m"));
    assertEquals(basicInfo.getQueryStats().isFullyBlocked(), true);
    assertEquals(basicInfo.getQueryStats().getBlockedReasons(), ImmutableSet.of(BlockedReason.WAITING_FOR_MEMORY));
    assertEquals(basicInfo.getQueryStats().getProgressPercentage(), OptionalDouble.of(100));
    assertEquals(basicInfo.getErrorCode(), StandardErrorCode.ABANDONED_QUERY.toErrorCode());
    assertEquals(basicInfo.getErrorType(), StandardErrorCode.ABANDONED_QUERY.toErrorCode().getType());
}

Example 23 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestSnapshotUtils method testDeleteSnapshot.

@Test
public void testDeleteSnapshot() throws IOException {
    QueryId queryId1 = new QueryId("query1");
    snapshotUtils.getOrCreateQuerySnapshotManager(queryId1, TEST_SNAPSHOT_SESSION);
    QueryId queryId2 = new QueryId("query2");
    snapshotUtils.getOrCreateQuerySnapshotManager(queryId2, TEST_SNAPSHOT_SESSION);
    when(fileSystemClient.deleteRecursively(anyObject())).thenReturn(// OK to remove query 1
    true).thenThrow(// Fail to remove query 2
    new IOException()).thenReturn(// OK to remove query 1-or-2 (in cleanup)
    true).thenThrow(// Fail to remove query 1-or-2 (in cleanup)
    new IOException()).thenReturn(// OK to remove query 1-or-2 (in cleanup)
    true);
    // OK
    snapshotUtils.removeQuerySnapshotManager(queryId1);
    verify(fileSystemClient).deleteRecursively(anyObject());
    // Fail
    snapshotUtils.removeQuerySnapshotManager(queryId2);
    verify(fileSystemClient, times(2)).deleteRecursively(anyObject());
    // 1 OK 1 Fail
    snapshotUtils.cleanupSnapshots();
    verify(fileSystemClient, times(4)).deleteRecursively(anyObject());
    // OK
    snapshotUtils.cleanupSnapshots();
    verify(fileSystemClient, times(6)).deleteRecursively(anyObject());
    // No-op
    snapshotUtils.cleanupSnapshots();
    verify(fileSystemClient, times(8)).deleteRecursively(anyObject());
}

Example 24 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestTaskSnapshotManager method testSpilledDeleted.

@Test
public void testSpilledDeleted() throws Exception {
    queryId = new QueryId("spilleddeletedquery");
    TaskId taskId1 = new TaskId(queryId.getId(), 1, 0);
    TaskSnapshotManager snapshotManager = new TaskSnapshotManager(taskId1, 0, snapshotUtils);
    snapshotManager.setTotalComponents(1);
    QuerySnapshotManager querySnapshotManager = new QuerySnapshotManager(queryId, snapshotUtils, TEST_SNAPSHOT_SESSION);
    querySnapshotManager.addNewTask(taskId1);
    snapshotUtils.getOrCreateQuerySnapshotManager(queryId, TEST_SNAPSHOT_SESSION);
    // first store
    long firstSnapshotId = 1L;
    File dirs = new File("/tmp/test_snapshot_manager/" + queryId + "/");
    File firstFile = new File("/tmp/test_snapshot_manager/" + queryId + "/firstFile");
    dirs.mkdirs();
    firstFile.createNewFile();
    try (FileWriter fw1 = new FileWriter(firstFile)) {
        String firstStr = "first string";
        fw1.write(firstStr);
        SnapshotStateId stateId = SnapshotStateId.forOperator(firstSnapshotId, taskId1, 3, 4, 5);
        snapshotManager.storeFile(stateId, firstFile.toPath());
        snapshotManager.succeededToCapture(stateId);
    }
    // second store, then deleted
    File secondFile = new File("/tmp/test_snapshot_manager/" + queryId + "/secondFile");
    secondFile.createNewFile();
    try (FileWriter fw2 = new FileWriter(secondFile)) {
        String secondStr = "second string";
        fw2.write(secondStr);
        SnapshotStateId secondId = SnapshotStateId.forOperator(2L, taskId1, 3, 4, 5);
        snapshotManager.storeFile(secondId, secondFile.toPath());
        snapshotManager.succeededToCapture(secondId);
        querySnapshotManager.snapshotInitiated(firstSnapshotId);
        querySnapshotManager.updateQueryCapture(taskId1, Collections.singletonMap(firstSnapshotId, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL)));
        File secondFileOperator = new File("/tmp/test_snapshot_manager/" + queryId + "/2/1/0/3/4/5/secondFile");
        assertTrue(secondFileOperator.exists());
        secondFileOperator.delete();
        assertFalse(secondFileOperator.exists());
        assertFalse(snapshotManager.loadFile(secondId, secondFile.toPath()));
    }
}

Example 25 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestTaskSnapshotManager method testStoreAndLoadFileBacktrack.

@Test
public void testStoreAndLoadFileBacktrack() throws Exception {
    queryId = new QueryId("filebacktrack");
    TaskId taskId = new TaskId(queryId.getId(), 2, 3);
    TaskSnapshotManager snapshotManager = new TaskSnapshotManager(taskId, 0, snapshotUtils);
    snapshotManager.setTotalComponents(1);
    // Create a file
    Path sourcePath = Paths.get(SNAPSHOT_FILE_SYSTEM_DIR + "/source/spill-test.txt");
    sourcePath.getParent().toFile().mkdirs();
    String fileContent = "Spill Contents";
    FileWriter fileWriter = new FileWriter(SNAPSHOT_FILE_SYSTEM_DIR + "/source/spill-test.txt");
    fileWriter.write(fileContent);
    fileWriter.close();
    QuerySnapshotManager querySnapshotManager = new QuerySnapshotManager(queryId, snapshotUtils, TEST_SNAPSHOT_SESSION);
    querySnapshotManager.addNewTask(taskId);
    SnapshotStateId id4save = new SnapshotStateId(2, taskId, "component1");
    snapshotManager.storeFile(id4save, sourcePath);
    snapshotManager.succeededToCapture(id4save);
    SnapshotStateId id4load = new SnapshotStateId(3, taskId, "component1");
    Path targetPath = Paths.get(SNAPSHOT_FILE_SYSTEM_DIR + "/target/spill-test.txt");
    // Try1: Previous snapshot ids not setup, so load fails
    assertNull(snapshotManager.loadFile(id4load, targetPath));
    // Try2: Previous snapshots are setup, so load should be successful
    querySnapshotManager.snapshotInitiated(2L);
    querySnapshotManager.updateQueryCapture(taskId, Collections.singletonMap(2L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL)));
    querySnapshotManager.getResumeSnapshotId();
    assertTrue(snapshotManager.loadFile(id4load, targetPath));
    String output = Files.readAllLines(targetPath).get(0);
    Assert.assertEquals(output, fileContent);
    // Try3: Previous snapshot failed
    querySnapshotManager.updateQueryRestore(taskId, Optional.of(new RestoreResult(2, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL))));
    querySnapshotManager.updateQueryCapture(taskId, Collections.singletonMap(2L, SnapshotInfo.withStatus(SnapshotResult.FAILED)));
    querySnapshotManager.snapshotInitiated(3L);
    querySnapshotManager.updateQueryCapture(taskId, Collections.singletonMap(3L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL)));
    querySnapshotManager.getResumeSnapshotId();
    assertFalse(snapshotManager.loadFile(id4load, targetPath));
}