Examples with MemorySegment org.apache.flink.core.memory.MemorySegment used on opensource projects

Example 1 with MemorySegment

use of org.apache.flink.core.memory.MemorySegment in project flink by apache.

the class MemoryManager method allocatePages.

/**
	 * Allocates a set of memory segments from this memory manager. If the memory manager pre-allocated the
	 * segments, they will be taken from the pool of memory segments. Otherwise, they will be allocated
	 * as part of this call.
	 * 
	 * @param owner The owner to associate with the memory segment, for the fallback release.
	 * @param target The list into which to put the allocated memory pages.
	 * @param numPages The number of pages to allocate.
	 * @throws MemoryAllocationException Thrown, if this memory manager does not have the requested amount
	 *                                   of memory pages any more.
	 */
public void allocatePages(Object owner, List<MemorySegment> target, int numPages) throws MemoryAllocationException {
    // sanity check
    if (owner == null) {
        throw new IllegalArgumentException("The memory owner must not be null.");
    }
    // reserve array space, if applicable
    if (target instanceof ArrayList) {
        ((ArrayList<MemorySegment>) target).ensureCapacity(numPages);
    }
    // -------------------- BEGIN CRITICAL SECTION -------------------
    synchronized (lock) {
        if (isShutDown) {
            throw new IllegalStateException("Memory manager has been shut down.");
        }
        // lazy case, the 'freeSegments.size()' is zero.
        if (numPages > (memoryPool.getNumberOfAvailableMemorySegments() + numNonAllocatedPages)) {
            throw new MemoryAllocationException("Could not allocate " + numPages + " pages. Only " + (memoryPool.getNumberOfAvailableMemorySegments() + numNonAllocatedPages) + " pages are remaining.");
        }
        Set<MemorySegment> segmentsForOwner = allocatedSegments.get(owner);
        if (segmentsForOwner == null) {
            segmentsForOwner = new HashSet<MemorySegment>(numPages);
            allocatedSegments.put(owner, segmentsForOwner);
        }
        if (isPreAllocated) {
            for (int i = numPages; i > 0; i--) {
                MemorySegment segment = memoryPool.requestSegmentFromPool(owner);
                target.add(segment);
                segmentsForOwner.add(segment);
            }
        } else {
            for (int i = numPages; i > 0; i--) {
                MemorySegment segment = memoryPool.allocateNewSegment(owner);
                target.add(segment);
                segmentsForOwner.add(segment);
            }
            numNonAllocatedPages -= numPages;
        }
    }
// -------------------- END CRITICAL SECTION -------------------
}

Example 2 with MemorySegment

use of org.apache.flink.core.memory.MemorySegment in project flink by apache.

the class MemoryManagerLazyAllocationTest method allocateAllSingle.

@Test
public void allocateAllSingle() {
    try {
        final AbstractInvokable mockInvoke = new DummyInvokable();
        List<MemorySegment> segments = new ArrayList<MemorySegment>();
        try {
            for (int i = 0; i < NUM_PAGES; i++) {
                segments.add(this.memoryManager.allocatePages(mockInvoke, 1).get(0));
            }
        } catch (MemoryAllocationException e) {
            fail("Unable to allocate memory");
        }
        for (MemorySegment seg : segments) {
            this.memoryManager.release(seg);
        }
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}

Example 3 with MemorySegment

use of org.apache.flink.core.memory.MemorySegment in project flink by apache.

the class BufferTest method testgetNioBufferThreadSafe.

@Test
public void testgetNioBufferThreadSafe() {
    final MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(1024);
    final BufferRecycler recycler = Mockito.mock(BufferRecycler.class);
    Buffer buffer = new Buffer(segment, recycler);
    ByteBuffer buf1 = buffer.getNioBuffer();
    ByteBuffer buf2 = buffer.getNioBuffer();
    assertNotNull(buf1);
    assertNotNull(buf2);
    assertTrue("Repeated call to getNioBuffer() returns the same nio buffer", buf1 != buf2);
}

Example 4 with MemorySegment

use of org.apache.flink.core.memory.MemorySegment in project flink by apache.

the class RecordWriterTest method createBufferProvider.

private BufferProvider createBufferProvider(final int bufferSize) throws IOException, InterruptedException {
    BufferProvider bufferProvider = mock(BufferProvider.class);
    when(bufferProvider.requestBufferBlocking()).thenAnswer(new Answer<Buffer>() {

        @Override
        public Buffer answer(InvocationOnMock invocationOnMock) throws Throwable {
            MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(bufferSize);
            Buffer buffer = new Buffer(segment, DiscardingRecycler.INSTANCE);
            return buffer;
        }
    });
    return bufferProvider;
}

Example 5 with MemorySegment

use of org.apache.flink.core.memory.MemorySegment in project flink by apache.

the class BufferFileWriterFileSegmentReaderTest method fillBufferWithAscendingNumbers.

public static int fillBufferWithAscendingNumbers(Buffer buffer, int currentNumber) {
    MemorySegment segment = buffer.getMemorySegment();
    final int size = buffer.getSize();
    for (int i = 0; i < size; i += 4) {
        segment.putInt(i, currentNumber++);
    }
    return currentNumber;
}