Java Code Examples for org.apache.flink.runtime.io.network.partition.consumer.BufferOrEvent#isBuffer()

/**
 * Tests that each checkpoint is only aborted once in case of an interleaved cancellation
 * barrier arrival of two consecutive checkpoints.
 */
@Test
public void testInterleavedCancellationBarriers() throws Exception {
	BufferOrEvent[] sequence = {
		createBarrier(1L, 0),
		createCancellationBarrier(2L, 0),
		createCancellationBarrier(1L, 1),
		createCancellationBarrier(2L, 1),
		createCancellationBarrier(1L, 2),
		createCancellationBarrier(2L, 2),
		createBuffer(0)
	};
	CheckpointSequenceValidator validator =
		new CheckpointSequenceValidator(-1, -2);
	inputGate = createBarrierTracker(3, sequence, validator);

	for (BufferOrEvent boe : sequence) {
		if (boe.isBuffer() || boe.getEvent().getClass() != CancelCheckpointMarker.class) {
			assertEquals(boe, inputGate.pollNext().get());
		}
	}
}
 

@Test
public void testSingleChannelWithBarriers() throws Exception {
	BufferOrEvent[] sequence = {
			createBuffer(0), createBuffer(0), createBuffer(0),
			createBarrier(1, 0),
			createBuffer(0), createBuffer(0), createBuffer(0), createBuffer(0),
			createBarrier(2, 0), createBarrier(3, 0),
			createBuffer(0), createBuffer(0),
			createBarrier(4, 0), createBarrier(5, 0), createBarrier(6, 0),
			createBuffer(0)
	};
	CheckpointSequenceValidator validator =
		new CheckpointSequenceValidator(1, 2, 3, 4, 5, 6);
	inputGate = createBarrierTracker(1, sequence, validator);

	for (BufferOrEvent boe : sequence) {
		if (boe.isBuffer() || boe.getEvent().getClass() != CheckpointBarrier.class) {
			assertEquals(boe, inputGate.pollNext().get());
		}
	}
}
 

/**
 * Tests that each checkpoint is only aborted once in case of an interleaved cancellation
 * barrier arrival of two consecutive checkpoints.
 */
@Test
public void testInterleavedCancellationBarriers() throws Exception {
	BufferOrEvent[] sequence = {
		createBarrier(1L, 0),
		createCancellationBarrier(2L, 0),
		createCancellationBarrier(1L, 1),
		createCancellationBarrier(2L, 1),
		createCancellationBarrier(1L, 2),
		createCancellationBarrier(2L, 2),
		createBuffer(0)
	};
	CheckpointSequenceValidator validator =
		new CheckpointSequenceValidator(-1, -2);
	inputGate = createBarrierTracker(3, sequence, validator);

	for (BufferOrEvent boe : sequence) {
		if (boe.isBuffer() || (boe.getEvent().getClass() != CheckpointBarrier.class && boe.getEvent().getClass() != CancelCheckpointMarker.class)) {
			assertEquals(boe, inputGate.pollNext().get());
		}
	}
}
 

@Override
public BufferOrEvent getNextNonBlocked() throws Exception {
	while (true) {
		Optional<BufferOrEvent> next = inputGate.getNextBufferOrEvent();
		if (!next.isPresent()) {
			// buffer or input exhausted
			return null;
		}

		BufferOrEvent bufferOrEvent = next.get();
		if (bufferOrEvent.isBuffer()) {
			return bufferOrEvent;
		}
		else if (bufferOrEvent.getEvent().getClass() == CheckpointBarrier.class) {
			processBarrier((CheckpointBarrier) bufferOrEvent.getEvent(), bufferOrEvent.getChannelIndex());
		}
		else if (bufferOrEvent.getEvent().getClass() == CancelCheckpointMarker.class) {
			processCheckpointAbortBarrier((CancelCheckpointMarker) bufferOrEvent.getEvent(), bufferOrEvent.getChannelIndex());
		}
		else {
			// some other event
			return bufferOrEvent;
		}
	}
}
 

private static void check(BufferOrEvent expected, BufferOrEvent present) {
	assertNotNull(expected);
	assertNotNull(present);
	assertEquals(expected.isBuffer(), present.isBuffer());

	if (expected.isBuffer()) {
		assertEquals(expected.getBuffer().getMaxCapacity(), present.getBuffer().getMaxCapacity());
		assertEquals(expected.getBuffer().getSize(), present.getBuffer().getSize());
		MemorySegment expectedMem = expected.getBuffer().getMemorySegment();
		MemorySegment presentMem = present.getBuffer().getMemorySegment();
		assertTrue("memory contents differs", expectedMem.compare(presentMem, 0, 0, PAGE_SIZE) == 0);
	}
	else {
		assertEquals(expected.getEvent(), present.getEvent());
	}
}
 

private void processBufferOrEvent(BufferOrEvent bufferOrEvent) throws Exception {
	if (bufferOrEvent.isBuffer()) {
		currentChannel = bufferOrEvent.getChannelIndex();
		currentRecordDeserializer = recordDeserializers[currentChannel];
		currentRecordDeserializer.setNextBuffer(bufferOrEvent.getBuffer());
	}
	else {
		// Event received
		final AbstractEvent event = bufferOrEvent.getEvent();
		// TODO: with barrierHandler.isFinished() we might not need to support any events on this level.
		if (event.getClass() != EndOfPartitionEvent.class) {
			throw new IOException("Unexpected event: " + event);
		}

		handleEndOfPartitionEvent(bufferOrEvent.getChannelIndex());
	}
}
 

@Override
public void close() {
	BufferOrEvent boe;
	while ((boe = currentBuffers.poll()) != null) {
		if (boe.isBuffer()) {
			boe.getBuffer().recycleBuffer();
		}
	}
}
 

private static void check(BufferOrEvent expected, BufferOrEvent present, int pageSize) {
	assertNotNull(expected);
	assertNotNull(present);
	assertEquals(expected.isBuffer(), present.isBuffer());

	if (expected.isBuffer()) {
		assertEquals(expected.getBuffer().getMaxCapacity(), present.getBuffer().getMaxCapacity());
		assertEquals(expected.getBuffer().getSize(), present.getBuffer().getSize());
		MemorySegment expectedMem = expected.getBuffer().getMemorySegment();
		MemorySegment presentMem = present.getBuffer().getMemorySegment();
		assertTrue("memory contents differs", expectedMem.compare(presentMem, 0, 0, pageSize) == 0);
	} else {
		assertEquals(expected.getEvent(), present.getEvent());
	}
}
 

@Override
public Optional<BufferOrEvent> pollNext() throws Exception {
	while (true) {
		Optional<BufferOrEvent> next = inputGate.pollNext();

		if (!next.isPresent()) {
			return handleEmptyBuffer();
		}

		BufferOrEvent bufferOrEvent = next.get();
		checkState(!barrierHandler.isBlocked(bufferOrEvent.getChannelInfo()));

		if (bufferOrEvent.isBuffer()) {
			return next;
		}
		else if (bufferOrEvent.getEvent().getClass() == CheckpointBarrier.class) {
			CheckpointBarrier checkpointBarrier = (CheckpointBarrier) bufferOrEvent.getEvent();
			barrierHandler.processBarrier(checkpointBarrier, bufferOrEvent.getChannelInfo());
			return next;
		}
		else if (bufferOrEvent.getEvent().getClass() == CancelCheckpointMarker.class) {
			barrierHandler.processCancellationBarrier((CancelCheckpointMarker) bufferOrEvent.getEvent());
		}
		else {
			if (bufferOrEvent.getEvent().getClass() == EndOfPartitionEvent.class) {
				barrierHandler.processEndOfPartition();
			}
			return next;
		}
	}
}
 

@Override
public void validateAlignmentBuffered(long actualBytesBuffered, BufferOrEvent... sequence) {
	long expectedBuffered = 0;
	for (BufferOrEvent boe : sequence) {
		if (boe.isBuffer()) {
			expectedBuffered += BufferSpiller.HEADER_SIZE + boe.getBuffer().getSize();
		}
	}

	assertEquals("Wrong alignment buffered bytes", actualBytesBuffered, expectedBuffered);
}
 

@Test
public void testMultiChannelSkippingCheckpoints() throws Exception {
	BufferOrEvent[] sequence = {
			createBuffer(0), createBuffer(2), createBuffer(0),
			createBarrier(1, 1), createBarrier(1, 2),
			createBuffer(2), createBuffer(1),
			createBarrier(1, 0),

			createBuffer(0), createBuffer(0), createBuffer(1), createBuffer(1), createBuffer(2),
			createBarrier(2, 0), createBarrier(2, 1), createBarrier(2, 2),

			createBuffer(2), createBuffer(2),
			createBarrier(3, 2),
			createBuffer(2), createBuffer(2),

			// jump to checkpoint 4
			createBarrier(4, 0),
			createBuffer(0), createBuffer(1), createBuffer(2),
			createBarrier(4, 1),
			createBuffer(1),
			createBarrier(4, 2),

			createBuffer(0)
	};
	CheckpointSequenceValidator validator =
		new CheckpointSequenceValidator(1, 2, 4);
	inputGate = createBarrierTracker(3, sequence, validator);

	for (BufferOrEvent boe : sequence) {
		if (boe.isBuffer() || boe.getEvent().getClass() != CheckpointBarrier.class) {
			assertEquals(boe, inputGate.pollNext().get());
		}
	}
}
 

@Override
public void close() throws IOException {
	BufferOrEvent boe;
	while ((boe = currentBuffers.poll()) != null) {
		if (boe.isBuffer()) {
			boe.getBuffer().recycleBuffer();
		}
	}
	super.close();
}
 

private static void check(BufferOrEvent expected, BufferOrEvent present, int pageSize) {
	assertNotNull(expected);
	assertNotNull(present);
	assertEquals(expected.isBuffer(), present.isBuffer());

	if (expected.isBuffer()) {
		assertEquals(expected.getBuffer().getMaxCapacity(), present.getBuffer().getMaxCapacity());
		assertEquals(expected.getBuffer().getSize(), present.getBuffer().getSize());
		MemorySegment expectedMem = expected.getBuffer().getMemorySegment();
		MemorySegment presentMem = present.getBuffer().getMemorySegment();
		assertTrue("memory contents differs", expectedMem.compare(presentMem, 0, 0, pageSize) == 0);
	} else {
		assertEquals(expected.getEvent(), present.getEvent());
	}
}
 

@Override
public void add(BufferOrEvent boe) throws IOException {
	try {
		ByteBuffer contents;
		if (boe.isBuffer()) {
			Buffer buf = boe.getBuffer();
			contents = buf.getNioBufferReadable();
		}
		else {
			contents = EventSerializer.toSerializedEvent(boe.getEvent());
		}

		headBuffer.clear();
		headBuffer.putInt(boe.getChannelIndex());
		headBuffer.putInt(contents.remaining());
		headBuffer.put((byte) (boe.isBuffer() ? 0 : 1));
		headBuffer.flip();

		bytesWritten += (headBuffer.remaining() + contents.remaining());

		FileUtils.writeCompletely(currentChannel, headBuffer);
		FileUtils.writeCompletely(currentChannel, contents);
	}
	finally {
		if (boe.isBuffer()) {
			boe.getBuffer().recycleBuffer();
		}
	}
}
 

/**
 * Adds a buffer or event to the sequence of spilled buffers and events.
 *
 * @param boe The buffer or event to add and spill.
 * @throws IOException Thrown, if the buffer of event could not be spilled.
 */
@Override
public void add(BufferOrEvent boe) throws IOException {
	try {
		ByteBuffer contents;
		if (boe.isBuffer()) {
			Buffer buf = boe.getBuffer();
			contents = buf.getNioBufferReadable();
		}
		else {
			contents = EventSerializer.toSerializedEvent(boe.getEvent());
		}

		headBuffer.clear();
		headBuffer.putInt(boe.getChannelIndex());
		headBuffer.putInt(contents.remaining());
		headBuffer.put((byte) (boe.isBuffer() ? 0 : 1));
		headBuffer.flip();

		bytesWritten += (headBuffer.remaining() + contents.remaining());

		FileUtils.writeCompletely(currentChannel, headBuffer);
		FileUtils.writeCompletely(currentChannel, contents);
	}
	finally {
		if (boe.isBuffer()) {
			boe.getBuffer().recycleBuffer();
		}
	}
}
 

/**
 * This test validates that the barrier tracker does not immediately
 * discard a pending checkpoint as soon as it sees a barrier from a
 * later checkpoint from some channel.
 *
 * <p>This behavior is crucial, otherwise topologies where different inputs
 * have different latency (and that latency is close to or higher than the
 * checkpoint interval) may skip many checkpoints, or fail to complete a
 * checkpoint all together.
 */
@Test
public void testCompleteCheckpointsOnLateBarriers() {
	try {
		BufferOrEvent[] sequence = {
				// checkpoint 2
				createBuffer(1), createBuffer(1), createBuffer(0), createBuffer(2),
				createBarrier(2, 1), createBarrier(2, 0), createBarrier(2, 2),

				// incomplete checkpoint 3
				createBuffer(1), createBuffer(0),
				createBarrier(3, 1), createBarrier(3, 2),

				// some barriers from checkpoint 4
				createBuffer(1), createBuffer(0),
				createBarrier(4, 2), createBarrier(4, 1),
				createBuffer(1), createBuffer(2),

				// last barrier from checkpoint 3
				createBarrier(3, 0),

				// complete checkpoint 4
				createBuffer(0), createBarrier(4, 0),

				// regular checkpoint 5
				createBuffer(1), createBuffer(2), createBarrier(5, 1),
				createBuffer(0), createBarrier(5, 0),
				createBuffer(1), createBarrier(5, 2),

				// checkpoint 6 (incomplete),
				createBuffer(1), createBarrier(6, 1),
				createBuffer(0), createBarrier(6, 0),

				// checkpoint 7, with early barriers for checkpoints 8 and 9
				createBuffer(1), createBarrier(7, 1),
				createBuffer(0), createBarrier(7, 2),
				createBuffer(2), createBarrier(8, 2),
				createBuffer(0), createBarrier(8, 1),
				createBuffer(1), createBarrier(9, 1),

				// complete checkpoint 7, first barriers from checkpoint 10
				createBarrier(7, 0),
				createBuffer(0), createBarrier(9, 2),
				createBuffer(2), createBarrier(10, 2),

				// complete checkpoint 8 and 9
				createBarrier(8, 0),
				createBuffer(1), createBuffer(2), createBarrier(9, 0),

				// trailing data
				createBuffer(1), createBuffer(0), createBuffer(2)
		};

		MockInputGate gate = new MockInputGate(PAGE_SIZE, 3, Arrays.asList(sequence));
		BarrierTracker tracker = new BarrierTracker(gate);

		CheckpointSequenceValidator validator =
				new CheckpointSequenceValidator(2, 3, 4, 5, 7, 8, 9);
		tracker.registerCheckpointEventHandler(validator);

		for (BufferOrEvent boe : sequence) {
			if (boe.isBuffer() || boe.getEvent().getClass() != CheckpointBarrier.class) {
				assertEquals(boe, tracker.getNextNonBlocked());
			}
		}

		assertNull(tracker.getNextNonBlocked());
		assertNull(tracker.getNextNonBlocked());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

@Override
public BufferOrEvent getNextNonBlocked() throws Exception {
	while (true) {
		// process buffered BufferOrEvents before grabbing new ones
		Optional<BufferOrEvent> next;
		if (currentBuffered == null) {
			next = inputGate.getNextBufferOrEvent();
		}
		else {
			next = Optional.ofNullable(currentBuffered.getNext());
			if (!next.isPresent()) {
				completeBufferedSequence();
				return getNextNonBlocked();
			}
		}

		if (!next.isPresent()) {
			if (!endOfStream) {
				// end of input stream. stream continues with the buffered data
				endOfStream = true;
				releaseBlocksAndResetBarriers();
				return getNextNonBlocked();
			}
			else {
				// final end of both input and buffered data
				return null;
			}
		}

		BufferOrEvent bufferOrEvent = next.get();
		if (isBlocked(bufferOrEvent.getChannelIndex())) {
			// if the channel is blocked, we just store the BufferOrEvent
			bufferBlocker.add(bufferOrEvent);
			checkSizeLimit();
		}
		else if (bufferOrEvent.isBuffer()) {
			return bufferOrEvent;
		}
		else if (bufferOrEvent.getEvent().getClass() == CheckpointBarrier.class) {
			if (!endOfStream) {
				// process barriers only if there is a chance of the checkpoint completing
				processBarrier((CheckpointBarrier) bufferOrEvent.getEvent(), bufferOrEvent.getChannelIndex());
			}
		}
		else if (bufferOrEvent.getEvent().getClass() == CancelCheckpointMarker.class) {
			processCancellationBarrier((CancelCheckpointMarker) bufferOrEvent.getEvent());
		}
		else {
			if (bufferOrEvent.getEvent().getClass() == EndOfPartitionEvent.class) {
				processEndOfPartition();
			}
			return bufferOrEvent;
		}
	}
}
 

protected boolean getNextRecord(T target) throws IOException, InterruptedException {
	// The action of partition request was removed from InputGate#setup since FLINK-16536, and this is the only
	// unified way for launching partition request for batch jobs. In order to avoid potential performance concern,
	// we might consider migrating this action back to the setup based on some condition judgement future.
	if (!requestedPartitions) {
		inputGate.requestPartitions();
		requestedPartitions = true;
	}

	if (isFinished) {
		return false;
	}

	while (true) {
		if (currentRecordDeserializer != null) {
			DeserializationResult result = currentRecordDeserializer.getNextRecord(target);

			if (result.isBufferConsumed()) {
				final Buffer currentBuffer = currentRecordDeserializer.getCurrentBuffer();

				currentBuffer.recycleBuffer();
				currentRecordDeserializer = null;
			}

			if (result.isFullRecord()) {
				return true;
			}
		}

		final BufferOrEvent bufferOrEvent = inputGate.getNext().orElseThrow(IllegalStateException::new);

		if (bufferOrEvent.isBuffer()) {
			currentRecordDeserializer = recordDeserializers.get(bufferOrEvent.getChannelInfo());
			currentRecordDeserializer.setNextBuffer(bufferOrEvent.getBuffer());
		}
		else {
			// sanity check for leftover data in deserializers. events should only come between
			// records, not in the middle of a fragment
			if (recordDeserializers.get(bufferOrEvent.getChannelInfo()).hasUnfinishedData()) {
				throw new IOException(
						"Received an event in channel " + bufferOrEvent.getChannelInfo() + " while still having "
						+ "data from a record. This indicates broken serialization logic. "
						+ "If you are using custom serialization code (Writable or Value types), check their "
						+ "serialization routines. In the case of Kryo, check the respective Kryo serializer.");
			}

			if (handleEvent(bufferOrEvent.getEvent())) {
				if (inputGate.isFinished()) {
					isFinished = true;
					return false;
				}
				else if (hasReachedEndOfSuperstep()) {
					return false;
				}
				// else: More data is coming...
			}
		}
	}
}
 

protected boolean getNextRecord(T target) throws IOException, InterruptedException {
	if (isFinished) {
		return false;
	}

	while (true) {
		if (currentRecordDeserializer != null) {
			DeserializationResult result = currentRecordDeserializer.getNextRecord(target);

			if (result.isBufferConsumed()) {
				final Buffer currentBuffer = currentRecordDeserializer.getCurrentBuffer();

				currentBuffer.recycleBuffer();
				currentRecordDeserializer = null;
			}

			if (result.isFullRecord()) {
				return true;
			}
		}

		final BufferOrEvent bufferOrEvent = inputGate.getNextBufferOrEvent().orElseThrow(IllegalStateException::new);

		if (bufferOrEvent.isBuffer()) {
			currentRecordDeserializer = recordDeserializers[bufferOrEvent.getChannelIndex()];
			currentRecordDeserializer.setNextBuffer(bufferOrEvent.getBuffer());
		}
		else {
			// sanity check for leftover data in deserializers. events should only come between
			// records, not in the middle of a fragment
			if (recordDeserializers[bufferOrEvent.getChannelIndex()].hasUnfinishedData()) {
				throw new IOException(
						"Received an event in channel " + bufferOrEvent.getChannelIndex() + " while still having "
						+ "data from a record. This indicates broken serialization logic. "
						+ "If you are using custom serialization code (Writable or Value types), check their "
						+ "serialization routines. In the case of Kryo, check the respective Kryo serializer.");
			}

			if (handleEvent(bufferOrEvent.getEvent())) {
				if (inputGate.isFinished()) {
					isFinished = true;
					return false;
				}
				else if (hasReachedEndOfSuperstep()) {
					return false;
				}
				// else: More data is coming...
			}
		}
	}
}
 

protected boolean getNextRecord(T target) throws IOException, InterruptedException {
	if (isFinished) {
		return false;
	}

	while (true) {
		if (currentRecordDeserializer != null) {
			DeserializationResult result = currentRecordDeserializer.getNextRecord(target);

			if (result.isBufferConsumed()) {
				final Buffer currentBuffer = currentRecordDeserializer.getCurrentBuffer();

				currentBuffer.recycleBuffer();
				currentRecordDeserializer = null;
			}

			if (result.isFullRecord()) {
				return true;
			}
		}

		final BufferOrEvent bufferOrEvent = inputGate.getNext().orElseThrow(IllegalStateException::new);

		if (bufferOrEvent.isBuffer()) {
			currentRecordDeserializer = recordDeserializers[bufferOrEvent.getChannelIndex()];
			currentRecordDeserializer.setNextBuffer(bufferOrEvent.getBuffer());
		}
		else {
			// sanity check for leftover data in deserializers. events should only come between
			// records, not in the middle of a fragment
			if (recordDeserializers[bufferOrEvent.getChannelIndex()].hasUnfinishedData()) {
				throw new IOException(
						"Received an event in channel " + bufferOrEvent.getChannelIndex() + " while still having "
						+ "data from a record. This indicates broken serialization logic. "
						+ "If you are using custom serialization code (Writable or Value types), check their "
						+ "serialization routines. In the case of Kryo, check the respective Kryo serializer.");
			}

			if (handleEvent(bufferOrEvent.getEvent())) {
				if (inputGate.isFinished()) {
					isFinished = true;
					return false;
				}
				else if (hasReachedEndOfSuperstep()) {
					return false;
				}
				// else: More data is coming...
			}
		}
	}
}