Examples with LocalMemoryContext io.trino.memory.context.LocalMemoryContext used on opensource projects

Example 6 with LocalMemoryContext

use of io.trino.memory.context.LocalMemoryContext in project trino by trinodb.

the class TestMemoryRevokingScheduler method testScheduleMemoryRevoking.

@Test
public void testScheduleMemoryRevoking() throws Exception {
    QueryContext q1 = getOrCreateQueryContext(new QueryId("q1"));
    QueryContext q2 = getOrCreateQueryContext(new QueryId("q2"));
    SqlTask sqlTask1 = newSqlTask(q1.getQueryId());
    SqlTask sqlTask2 = newSqlTask(q2.getQueryId());
    TaskContext taskContext1 = getOrCreateTaskContext(sqlTask1);
    PipelineContext pipelineContext11 = taskContext1.addPipelineContext(0, false, false, false);
    DriverContext driverContext111 = pipelineContext11.addDriverContext();
    OperatorContext operatorContext1 = driverContext111.addOperatorContext(1, new PlanNodeId("na"), "na");
    OperatorContext operatorContext2 = driverContext111.addOperatorContext(2, new PlanNodeId("na"), "na");
    DriverContext driverContext112 = pipelineContext11.addDriverContext();
    OperatorContext operatorContext3 = driverContext112.addOperatorContext(3, new PlanNodeId("na"), "na");
    TaskContext taskContext2 = getOrCreateTaskContext(sqlTask2);
    PipelineContext pipelineContext21 = taskContext2.addPipelineContext(1, false, false, false);
    DriverContext driverContext211 = pipelineContext21.addDriverContext();
    OperatorContext operatorContext4 = driverContext211.addOperatorContext(4, new PlanNodeId("na"), "na");
    OperatorContext operatorContext5 = driverContext211.addOperatorContext(5, new PlanNodeId("na"), "na");
    Collection<SqlTask> tasks = ImmutableList.of(sqlTask1, sqlTask2);
    MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(memoryPool, () -> tasks, executor, 1.0, 1.0);
    allOperatorContexts = ImmutableSet.of(operatorContext1, operatorContext2, operatorContext3, operatorContext4, operatorContext5);
    assertMemoryRevokingNotRequested();
    requestMemoryRevoking(scheduler);
    assertEquals(10, memoryPool.getFreeBytes());
    assertMemoryRevokingNotRequested();
    LocalMemoryContext revocableMemory1 = operatorContext1.localRevocableMemoryContext();
    LocalMemoryContext revocableMemory3 = operatorContext3.localRevocableMemoryContext();
    LocalMemoryContext revocableMemory4 = operatorContext4.localRevocableMemoryContext();
    LocalMemoryContext revocableMemory5 = operatorContext5.localRevocableMemoryContext();
    revocableMemory1.setBytes(3);
    revocableMemory3.setBytes(6);
    assertEquals(1, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // we are still good - no revoking needed
    assertMemoryRevokingNotRequested();
    revocableMemory4.setBytes(7);
    assertEquals(-6, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // we need to revoke 3 and 6
    assertMemoryRevokingRequestedFor(operatorContext1, operatorContext3);
    // yet another revoking request should not change anything
    requestMemoryRevoking(scheduler);
    assertMemoryRevokingRequestedFor(operatorContext1, operatorContext3);
    // lets revoke some bytes
    revocableMemory1.setBytes(0);
    operatorContext1.resetMemoryRevokingRequested();
    requestMemoryRevoking(scheduler);
    assertMemoryRevokingRequestedFor(operatorContext3);
    assertEquals(-3, memoryPool.getFreeBytes());
    // and allocate some more
    revocableMemory5.setBytes(3);
    assertEquals(-6, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // we are still good with just OC3 in process of revoking
    assertMemoryRevokingRequestedFor(operatorContext3);
    // and allocate some more
    revocableMemory5.setBytes(4);
    assertEquals(-7, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // no we have to trigger revoking for OC4
    assertMemoryRevokingRequestedFor(operatorContext3, operatorContext4);
}

Example 7 with LocalMemoryContext

use of io.trino.memory.context.LocalMemoryContext in project trino by trinodb.

the class TestMemoryTracking method testLocalTotalMemoryLimitExceeded.

@Test
public void testLocalTotalMemoryLimitExceeded() {
    LocalMemoryContext memoryContext = operatorContext.newLocalUserMemoryContext("test");
    memoryContext.setBytes(100);
    assertOperatorMemoryAllocations(operatorContext.getOperatorMemoryContext(), 100, 0);
    memoryContext.setBytes(queryMaxMemory.toBytes());
    assertOperatorMemoryAllocations(operatorContext.getOperatorMemoryContext(), queryMaxMemory.toBytes(), 0);
    assertThatThrownBy(() -> memoryContext.setBytes(queryMaxMemory.toBytes() + 1)).isInstanceOf(ExceededMemoryLimitException.class).hasMessage("Query exceeded per-node memory limit of %1$s [Allocated: %1$s, Delta: 1B, Top Consumers: {test=%1$s}]", queryMaxMemory);
}

Example 8 with LocalMemoryContext

use of io.trino.memory.context.LocalMemoryContext in project trino by trinodb.

the class TestMemoryTracking method testTrySetBytes.

@Test
public void testTrySetBytes() {
    LocalMemoryContext localMemoryContext = operatorContext.localUserMemoryContext();
    assertTrue(localMemoryContext.trySetBytes(100_000_000));
    assertStats(operatorContext.getNestedOperatorStats(), driverContext.getDriverStats(), pipelineContext.getPipelineStats(), taskContext.getTaskStats(), 100_000_000, 0);
    assertTrue(localMemoryContext.trySetBytes(200_000_000));
    assertStats(operatorContext.getNestedOperatorStats(), driverContext.getDriverStats(), pipelineContext.getPipelineStats(), taskContext.getTaskStats(), 200_000_000, 0);
    assertTrue(localMemoryContext.trySetBytes(100_000_000));
    assertStats(operatorContext.getNestedOperatorStats(), driverContext.getDriverStats(), pipelineContext.getPipelineStats(), taskContext.getTaskStats(), 100_000_000, 0);
    // allocating more than the pool size should fail and we should have the same stats as before
    assertFalse(localMemoryContext.trySetBytes(memoryPool.getMaxBytes() + 1));
    assertStats(operatorContext.getNestedOperatorStats(), driverContext.getDriverStats(), pipelineContext.getPipelineStats(), taskContext.getTaskStats(), 100_000_000, 0);
}

Example 9 with LocalMemoryContext

use of io.trino.memory.context.LocalMemoryContext in project trino by trinodb.

the class TestMemoryTracking method testRevocableMemoryAllocations.

@Test
public void testRevocableMemoryAllocations() {
    LocalMemoryContext userMemory = operatorContext.localUserMemoryContext();
    LocalMemoryContext revocableMemory = operatorContext.localRevocableMemoryContext();
    revocableMemory.setBytes(100_000_000);
    assertStats(operatorContext.getNestedOperatorStats(), driverContext.getDriverStats(), pipelineContext.getPipelineStats(), taskContext.getTaskStats(), 0, 100_000_000);
    userMemory.setBytes(100_000_000);
    revocableMemory.setBytes(200_000_000);
    assertStats(operatorContext.getNestedOperatorStats(), driverContext.getDriverStats(), pipelineContext.getPipelineStats(), taskContext.getTaskStats(), 100_000_000, 200_000_000);
}

Example 10 with LocalMemoryContext

use of io.trino.memory.context.LocalMemoryContext in project trino by trinodb.

the class TestPageProcessor method testProjectLazyLoad.

@Test
public void testProjectLazyLoad() {
    PageProcessor pageProcessor = new PageProcessor(Optional.of(new SelectAllFilter()), ImmutableList.of(new LazyPagePageProjection()), OptionalInt.of(MAX_BATCH_SIZE));
    // if channel 1 is loaded, test will fail
    Page inputPage = new Page(createLongSequenceBlock(0, 100), new LazyBlock(100, () -> {
        throw new AssertionError("Lazy block should not be loaded");
    }));
    LocalMemoryContext memoryContext = newSimpleAggregatedMemoryContext().newLocalMemoryContext(PageProcessor.class.getSimpleName());
    Iterator<Optional<Page>> output = pageProcessor.process(SESSION, new DriverYieldSignal(), memoryContext, inputPage);
    List<Optional<Page>> outputPages = ImmutableList.copyOf(output);
    assertEquals(outputPages.size(), 1);
    assertPageEquals(ImmutableList.of(BIGINT), outputPages.get(0).orElse(null), new Page(createLongSequenceBlock(0, 100)));
}