Search in sources :

Example 1 with BufferAndChannel

use of org.apache.flink.runtime.io.network.util.TestProducerSource.BufferAndChannel in project flink by apache.

the class TestPartitionProducer method call.

@Override
public Boolean call() throws Exception {
    boolean success = false;
    try {
        BufferAndChannel bufferAndChannel;
        while ((bufferAndChannel = source.getNextBuffer()) != null) {
            ByteBuffer record = ByteBuffer.wrap(bufferAndChannel.getBuffer());
            partition.emitRecord(record, bufferAndChannel.getTargetChannel());
            // Check for interrupted flag after adding data to prevent resource leaks
            if (Thread.interrupted()) {
                throw new InterruptedException();
            }
            if (isSlowProducer) {
                Thread.sleep(random.nextInt(MAX_SLEEP_TIME_MS + 1));
            }
        }
        partition.finish();
        success = true;
        return true;
    } finally {
        if (!success) {
            partition.release();
        }
    }
}
Also used : BufferAndChannel(org.apache.flink.runtime.io.network.util.TestProducerSource.BufferAndChannel) ByteBuffer(java.nio.ByteBuffer)

Example 2 with BufferAndChannel

use of org.apache.flink.runtime.io.network.util.TestProducerSource.BufferAndChannel in project flink by apache.

the class TestSubpartitionProducer method call.

@Override
public Boolean call() throws Exception {
    boolean success = false;
    try {
        BufferAndChannel bufferAndChannel;
        while ((bufferAndChannel = source.getNextBuffer()) != null) {
            MemorySegment segment = MemorySegmentFactory.wrap(bufferAndChannel.getBuffer());
            subpartition.add(new BufferConsumer(new NetworkBuffer(segment, MemorySegment::free, Buffer.DataType.DATA_BUFFER), segment.size()));
            // Check for interrupted flag after adding data to prevent resource leaks
            if (Thread.interrupted()) {
                throw new InterruptedException();
            }
            if (isSlowProducer) {
                Thread.sleep(random.nextInt(MAX_SLEEP_TIME_MS + 1));
            }
        }
        subpartition.finish();
        success = true;
        return true;
    } finally {
        if (!success) {
            subpartition.release();
        }
    }
}
Also used : BufferAndChannel(org.apache.flink.runtime.io.network.util.TestProducerSource.BufferAndChannel) NetworkBuffer(org.apache.flink.runtime.io.network.buffer.NetworkBuffer) BufferConsumer(org.apache.flink.runtime.io.network.buffer.BufferConsumer) MemorySegment(org.apache.flink.core.memory.MemorySegment)

Aggregations

BufferAndChannel (org.apache.flink.runtime.io.network.util.TestProducerSource.BufferAndChannel)2 ByteBuffer (java.nio.ByteBuffer)1 MemorySegment (org.apache.flink.core.memory.MemorySegment)1 BufferConsumer (org.apache.flink.runtime.io.network.buffer.BufferConsumer)1 NetworkBuffer (org.apache.flink.runtime.io.network.buffer.NetworkBuffer)1