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

Example 1 with LocalMemoryContext

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

the class ParquetReader method allocateBlock.

private byte[] allocateBlock(int length) {
    byte[] buffer = new byte[length];
    LocalMemoryContext blockMemoryContext = currentRowGroupMemoryContext.newLocalMemoryContext(ParquetReader.class.getSimpleName());
    blockMemoryContext.setBytes(buffer.length);
    return buffer;
}

Example 2 with LocalMemoryContext

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

the class TestOperatorMemoryRevocation method testOperatorMemoryRevocation.

@Test
public void testOperatorMemoryRevocation() {
    AtomicInteger counter = new AtomicInteger();
    OperatorContext operatorContext = TestingOperatorContext.create(scheduledExecutor);
    LocalMemoryContext revocableMemoryContext = operatorContext.localRevocableMemoryContext();
    revocableMemoryContext.setBytes(1000);
    operatorContext.setMemoryRevocationRequestListener(() -> counter.incrementAndGet());
    operatorContext.requestMemoryRevoking();
    assertTrue(operatorContext.isMemoryRevokingRequested());
    assertEquals(counter.get(), 1);
    // calling resetMemoryRevokingRequested() should clear the memory revoking requested flag
    operatorContext.resetMemoryRevokingRequested();
    assertFalse(operatorContext.isMemoryRevokingRequested());
    operatorContext.requestMemoryRevoking();
    assertEquals(counter.get(), 2);
    assertTrue(operatorContext.isMemoryRevokingRequested());
}

Example 3 with LocalMemoryContext

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

the class TestOperatorMemoryRevocation method testRevocationAlreadyRequested.

@Test
public void testRevocationAlreadyRequested() {
    AtomicInteger counter = new AtomicInteger();
    OperatorContext operatorContext = TestingOperatorContext.create(scheduledExecutor);
    LocalMemoryContext revocableMemoryContext = operatorContext.localRevocableMemoryContext();
    revocableMemoryContext.setBytes(1000);
    // when memory revocation is already requested setting a listener should immediately execute it
    operatorContext.requestMemoryRevoking();
    operatorContext.setMemoryRevocationRequestListener(() -> counter.incrementAndGet());
    assertTrue(operatorContext.isMemoryRevokingRequested());
    assertEquals(counter.get(), 1);
}

Example 4 with LocalMemoryContext

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

the class TestMemoryTracking method testDestroy.

@Test
public void testDestroy() {
    LocalMemoryContext newLocalSystemMemoryContext = operatorContext.newLocalSystemMemoryContext("test");
    LocalMemoryContext newLocalUserMemoryContext = operatorContext.localUserMemoryContext();
    LocalMemoryContext newLocalRevocableMemoryContext = operatorContext.localRevocableMemoryContext();
    newLocalSystemMemoryContext.setBytes(100_000);
    newLocalRevocableMemoryContext.setBytes(200_000);
    newLocalUserMemoryContext.setBytes(400_000);
    assertEquals(operatorContext.getOperatorMemoryContext().getSystemMemory(), 100_000);
    assertEquals(operatorContext.getOperatorMemoryContext().getUserMemory(), 400_000);
    operatorContext.destroy();
    assertOperatorMemoryAllocations(operatorContext.getOperatorMemoryContext(), 0, 0, 0);
}

Example 5 with LocalMemoryContext

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

the class TestMemoryTracking method testTrySetBytes.

@Test
public void testTrySetBytes() {
    LocalMemoryContext localMemoryContext = operatorContext.localUserMemoryContext();
    assertTrue(localMemoryContext.trySetBytes(100_000_000));
    assertStats(operatorContext.getOperatorStats(), driverContext.getDriverStats(), pipelineContext.getPipelineStats(), taskContext.getTaskStats(), 100_000_000, 0, 0);
    assertTrue(localMemoryContext.trySetBytes(200_000_000));
    assertStats(operatorContext.getOperatorStats(), driverContext.getDriverStats(), pipelineContext.getPipelineStats(), taskContext.getTaskStats(), 200_000_000, 0, 0);
    assertTrue(localMemoryContext.trySetBytes(100_000_000));
    assertStats(operatorContext.getOperatorStats(), driverContext.getDriverStats(), pipelineContext.getPipelineStats(), taskContext.getTaskStats(), 100_000_000, 0, 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.getOperatorStats(), driverContext.getDriverStats(), pipelineContext.getPipelineStats(), taskContext.getTaskStats(), 100_000_000, 0, 0);
}