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

Example 1 with MemoryTrackingContext

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

the class TestTableWriterOperator method assertMemoryIsReleased.

private void assertMemoryIsReleased(TableWriterOperator tableWriterOperator) {
    OperatorContext tableWriterOperatorOperatorContext = tableWriterOperator.getOperatorContext();
    MemoryTrackingContext tableWriterMemoryContext = tableWriterOperatorOperatorContext.getOperatorMemoryContext();
    assertEquals(tableWriterMemoryContext.getUserMemory(), 0);
    assertEquals(tableWriterMemoryContext.getRevocableMemory(), 0);
    Operator statisticAggregationOperator = tableWriterOperator.getStatisticAggregationOperator();
    assertTrue(statisticAggregationOperator instanceof AggregationOperator);
    AggregationOperator aggregationOperator = (AggregationOperator) statisticAggregationOperator;
    OperatorContext aggregationOperatorOperatorContext = aggregationOperator.getOperatorContext();
    MemoryTrackingContext aggregationOperatorMemoryContext = aggregationOperatorOperatorContext.getOperatorMemoryContext();
    assertEquals(aggregationOperatorMemoryContext.getUserMemory(), 0);
    assertEquals(aggregationOperatorMemoryContext.getRevocableMemory(), 0);
}

Example 2 with MemoryTrackingContext

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

the class TestMemoryTracking method testOperatorAllocations.

@Test
public void testOperatorAllocations() {
    MemoryTrackingContext operatorMemoryContext = operatorContext.getOperatorMemoryContext();
    LocalMemoryContext userMemory = operatorContext.localUserMemoryContext();
    LocalMemoryContext revocableMemory = operatorContext.localRevocableMemoryContext();
    userMemory.setBytes(100);
    assertOperatorMemoryAllocations(operatorMemoryContext, 100, 0);
    assertOperatorMemoryAllocations(operatorMemoryContext, 100, 0);
    userMemory.setBytes(500);
    assertOperatorMemoryAllocations(operatorMemoryContext, 500, 0);
    userMemory.setBytes(userMemory.getBytes() - 500);
    assertOperatorMemoryAllocations(operatorMemoryContext, 0, 0);
    revocableMemory.setBytes(300);
    assertOperatorMemoryAllocations(operatorMemoryContext, 0, 300);
    assertThatThrownBy(() -> userMemory.setBytes(userMemory.getBytes() - 500)).isInstanceOf(IllegalArgumentException.class).hasMessage("bytes cannot be negative");
    operatorContext.destroy();
    assertOperatorMemoryAllocations(operatorMemoryContext, 0, 0);
}

Example 3 with MemoryTrackingContext

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

the class TestingOperatorContext method create.

public static OperatorContext create(ScheduledExecutorService scheduledExecutor) {
    Executor executor = MoreExecutors.directExecutor();
    TaskContext taskContext = TestingTaskContext.createTaskContext(executor, scheduledExecutor, TestingSession.testSessionBuilder().build());
    MemoryTrackingContext pipelineMemoryContext = new MemoryTrackingContext(newSimpleAggregatedMemoryContext(), newSimpleAggregatedMemoryContext());
    PipelineContext pipelineContext = new PipelineContext(1, taskContext, executor, scheduledExecutor, pipelineMemoryContext, false, false, false);
    DriverContext driverContext = new DriverContext(pipelineContext, executor, scheduledExecutor, pipelineMemoryContext, Lifespan.taskWide(), 0L);
    OperatorContext operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("test"), "operator type");
    return operatorContext;
}