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

Example 21 with LocalMemoryContext

use of io.prestosql.memory.context.LocalMemoryContext in project hetu-core by openlookeng.

the class TestMemoryTracking method testLocalSystemAllocations.

@Test
public void testLocalSystemAllocations() {
    long pipelineLocalAllocation = 1_000_000;
    long taskLocalAllocation = 10_000_000;
    LocalMemoryContext pipelineLocalSystemMemoryContext = pipelineContext.localSystemMemoryContext();
    pipelineLocalSystemMemoryContext.setBytes(pipelineLocalAllocation);
    assertLocalMemoryAllocations(pipelineContext.getPipelineMemoryContext(), pipelineLocalAllocation, 0, pipelineLocalAllocation);
    LocalMemoryContext taskLocalSystemMemoryContext = taskContext.localSystemMemoryContext();
    taskLocalSystemMemoryContext.setBytes(taskLocalAllocation);
    assertLocalMemoryAllocations(taskContext.getTaskMemoryContext(), pipelineLocalAllocation + taskLocalAllocation, 0, taskLocalAllocation);
    assertEquals(pipelineContext.getPipelineStats().getSystemMemoryReservation().toBytes(), pipelineLocalAllocation, "task level allocations should not be visible at the pipeline level");
    pipelineLocalSystemMemoryContext.setBytes(pipelineLocalSystemMemoryContext.getBytes() - pipelineLocalAllocation);
    assertLocalMemoryAllocations(pipelineContext.getPipelineMemoryContext(), taskLocalAllocation, 0, 0);
    taskLocalSystemMemoryContext.setBytes(taskLocalSystemMemoryContext.getBytes() - taskLocalAllocation);
    assertLocalMemoryAllocations(taskContext.getTaskMemoryContext(), 0, 0, 0);
}

Example 22 with LocalMemoryContext

use of io.prestosql.memory.context.LocalMemoryContext in project hetu-core by openlookeng.

the class TestQueryContext method testMoveTaggedAllocations.

@Test
public void testMoveTaggedAllocations() {
    MemoryPool generalPool = new MemoryPool(GENERAL_POOL, new DataSize(10_000, BYTE));
    MemoryPool reservedPool = new MemoryPool(RESERVED_POOL, new DataSize(10_000, BYTE));
    QueryId queryId = new QueryId("query");
    QueryContext queryContext = createQueryContext(queryId, generalPool);
    TaskStateMachine taskStateMachine = new TaskStateMachine(TaskId.valueOf("task-id"), TEST_EXECUTOR);
    TaskContext taskContext = queryContext.addTaskContext(taskStateMachine, TEST_SESSION, false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
    DriverContext driverContext = taskContext.addPipelineContext(0, false, false, false).addDriverContext();
    OperatorContext operatorContext = driverContext.addOperatorContext(0, new PlanNodeId("test"), "test");
    // allocate some memory in the general pool
    LocalMemoryContext memoryContext = operatorContext.aggregateUserMemoryContext().newLocalMemoryContext("test_context");
    memoryContext.setBytes(1_000);
    Map<String, Long> allocations = generalPool.getTaggedMemoryAllocations().get(queryId);
    assertEquals(allocations, ImmutableMap.of("test_context", 1_000L));
    queryContext.setMemoryPool(reservedPool);
    assertNull(generalPool.getTaggedMemoryAllocations().get(queryId));
    allocations = reservedPool.getTaggedMemoryAllocations().get(queryId);
    assertEquals(allocations, ImmutableMap.of("test_context", 1_000L));
    assertEquals(generalPool.getFreeBytes(), 10_000);
    assertEquals(reservedPool.getFreeBytes(), 9_000);
    memoryContext.close();
    assertEquals(generalPool.getFreeBytes(), 10_000);
    assertEquals(reservedPool.getFreeBytes(), 10_000);
}

Example 23 with LocalMemoryContext

use of io.prestosql.memory.context.LocalMemoryContext in project hetu-core by openlookeng.

the class TestMemoryTracking method testTrySetZeroBytesFullPool.

@Test
public void testTrySetZeroBytesFullPool() {
    LocalMemoryContext localMemoryContext = operatorContext.localUserMemoryContext();
    // fill up the pool
    memoryPool.reserve(new QueryId("test_query"), "test", memoryPool.getFreeBytes());
    // try to reserve 0 bytes in the full pool
    assertTrue(localMemoryContext.trySetBytes(localMemoryContext.getBytes()));
}