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

Example 51 with MemorySegment

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

the class NetworkBufferAllocator method allocateUnPooledNetworkBuffer.

/**
 * Allocates an un-pooled network buffer with the specific size.
 *
 * @param size The requested buffer size.
 * @param dataType The data type this buffer represents.
 * @return The un-pooled network buffer.
 */
Buffer allocateUnPooledNetworkBuffer(int size, Buffer.DataType dataType) {
    checkArgument(size > 0, "Illegal buffer size, must be positive.");
    byte[] byteArray = new byte[size];
    MemorySegment memSeg = MemorySegmentFactory.wrap(byteArray);
    return new NetworkBuffer(memSeg, FreeingBufferRecycler.INSTANCE, dataType);
}

Example 52 with MemorySegment

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

the class LocalBufferPool method requestMemorySegmentFromGlobal.

private boolean requestMemorySegmentFromGlobal() {
    assert Thread.holdsLock(availableMemorySegments);
    if (isRequestedSizeReached()) {
        return false;
    }
    checkState(!isDestroyed, "Destroyed buffer pools should never acquire segments - this will lead to buffer leaks.");
    MemorySegment segment = networkBufferPool.requestPooledMemorySegment();
    if (segment != null) {
        availableMemorySegments.add(segment);
        numberOfRequestedMemorySegments++;
        return true;
    }
    return false;
}

Example 53 with MemorySegment

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

the class LocalBufferPool method returnExcessMemorySegments.

private void returnExcessMemorySegments() {
    assert Thread.holdsLock(availableMemorySegments);
    while (hasExcessBuffers()) {
        MemorySegment segment = availableMemorySegments.poll();
        if (segment == null) {
            return;
        }
        returnMemorySegment(segment);
    }
}

Example 54 with MemorySegment

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

the class ReadOnlySlicedBufferTest method setUp.

@Before
public void setUp() throws Exception {
    final MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(BUFFER_SIZE);
    buffer = new NetworkBuffer(segment, FreeingBufferRecycler.INSTANCE, Buffer.DataType.DATA_BUFFER, 0);
    for (int i = 0; i < DATA_SIZE; ++i) {
        buffer.writeByte(i);
    }
}

Example 55 with MemorySegment

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

the class BatchShuffleReadBufferPoolTest method testBufferFulfilledByRecycledBuffers.

@Test
public void testBufferFulfilledByRecycledBuffers() throws Exception {
    int numRequestThreads = 2;
    AtomicReference<Throwable> exception = new AtomicReference<>();
    BatchShuffleReadBufferPool bufferPool = createBufferPool();
    Map<Object, List<MemorySegment>> buffers = new ConcurrentHashMap<>();
    try {
        Object[] owners = new Object[] { new Object(), new Object(), new Object(), new Object() };
        for (int i = 0; i < 4; ++i) {
            buffers.put(owners[i], bufferPool.requestBuffers());
        }
        assertEquals(0, bufferPool.getAvailableBuffers());
        Thread[] requestThreads = new Thread[numRequestThreads];
        for (int i = 0; i < numRequestThreads; ++i) {
            requestThreads[i] = new Thread(() -> {
                try {
                    Object owner = new Object();
                    List<MemorySegment> allocated = null;
                    while (allocated == null || allocated.isEmpty()) {
                        allocated = bufferPool.requestBuffers();
                    }
                    buffers.put(owner, allocated);
                } catch (Throwable throwable) {
                    exception.set(throwable);
                }
            });
            requestThreads[i].start();
        }
        // recycle one by one
        for (MemorySegment segment : buffers.remove(owners[0])) {
            bufferPool.recycle(segment);
        }
        // bulk recycle
        bufferPool.recycle(buffers.remove(owners[1]));
        for (Thread requestThread : requestThreads) {
            requestThread.join();
        }
        assertNull(exception.get());
        assertEquals(0, bufferPool.getAvailableBuffers());
        assertEquals(4, buffers.size());
    } finally {
        for (Object owner : buffers.keySet()) {
            bufferPool.recycle(buffers.remove(owner));
        }
        assertEquals(bufferPool.getNumTotalBuffers(), bufferPool.getAvailableBuffers());
        bufferPool.destroy();
    }
}