org.apache.flink.runtime.io.network.partition.consumer.BufferOrEvent Java Examples

@Test
public void testSingleChannelAbortCheckpoint() throws Exception {
	BufferOrEvent[] sequence = {
			createBuffer(0),
			createBarrier(1, 0),
			createBuffer(0),
			createBarrier(2, 0),
			createCancellationBarrier(4, 0),
			createBarrier(5, 0),
			createBuffer(0),
			createCancellationBarrier(6, 0),
			createBuffer(0)
	};
	// negative values mean an expected cancellation call!
	CheckpointSequenceValidator validator =
		new CheckpointSequenceValidator(1, 2, -4, 5, -6);
	inputGate = createBarrierTracker(1, sequence, validator);

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

@Test
public void testEarlyCleanup() throws Exception {
	ValidatingCheckpointHandler handler = new ValidatingCheckpointHandler(1);
	inputGate = createInputGate(3, handler);

	// checkpoint 1
	final BufferOrEvent[] sequence1 = addSequence(inputGate,
		createBuffer(0), createBuffer(1), createBuffer(2),
		createBarrier(1, 1), createBarrier(1, 2), createBarrier(1, 0));
	assertOutput(sequence1);
	assertEquals(1L, channelStateWriter.getLastStartedCheckpointId());
	assertInflightData();

	// checkpoint 2
	final BufferOrEvent[] sequence2 = addSequence(inputGate,
		createBuffer(2), createBuffer(1), createBuffer(0),
		createBarrier(2, 1),
		createBuffer(1), createBuffer(1), createEndOfPartition(1), createBuffer(0), createBuffer(2),
		createBarrier(2, 2),
		createBuffer(2), createEndOfPartition(2), createBuffer(0), createEndOfPartition(0));
	assertOutput(sequence2);
	assertEquals(2L, channelStateWriter.getLastStartedCheckpointId());
	assertInflightData();
}
 

@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;
		}
	}
}
 

/**
 * Validates that the buffer behaves correctly if no checkpoint barriers come,
 * for an input with multiple input channels.
 */
@Test
public void testMultiChannelNoBarriers() throws Exception {
	BufferOrEvent[] sequence = {
		createBuffer(2, PAGE_SIZE), createBuffer(2, PAGE_SIZE), createBuffer(0, PAGE_SIZE),
		createBuffer(1, PAGE_SIZE), createBuffer(0, PAGE_SIZE), createEndOfPartition(0),
		createBuffer(3, PAGE_SIZE), createBuffer(1, PAGE_SIZE), createEndOfPartition(3),
		createBuffer(1, PAGE_SIZE), createEndOfPartition(1), createBuffer(2, PAGE_SIZE), createEndOfPartition(2)
	};

	MockInputGate gate = new MockInputGate(PAGE_SIZE, 4, Arrays.asList(sequence));
	BarrierBuffer buffer = createBarrierHandler(gate);

	for (BufferOrEvent boe : sequence) {
		assertEquals(boe, buffer.getNextNonBlocked());
	}

	assertEquals(0L, buffer.getAlignmentDurationNanos());

	assertNull(buffer.getNextNonBlocked());
	assertNull(buffer.getNextNonBlocked());

	buffer.cleanup();
}
 

private static BufferOrEvent generateAndWriteEvent(FileChannel fileChannel, Random rnd, int numberOfChannels) throws IOException {
	long magicNumber = rnd.nextLong();
	byte[] data = new byte[rnd.nextInt(1000)];
	rnd.nextBytes(data);
	TestEvent evt = new TestEvent(magicNumber, data);

	int channelIndex = rnd.nextInt(numberOfChannels);

	ByteBuffer serializedEvent = EventSerializer.toSerializedEvent(evt);
	ByteBuffer header = ByteBuffer.allocate(9);
	header.order(ByteOrder.LITTLE_ENDIAN);

	header.putInt(channelIndex);
	header.putInt(serializedEvent.remaining());
	header.put((byte) 1);
	header.flip();

	FileUtils.writeCompletely(fileChannel, header);
	FileUtils.writeCompletely(fileChannel, serializedEvent);
	return new BufferOrEvent(evt, channelIndex);
}
 

/**
 * Validates that the buffer behaves correctly if no checkpoint barriers come,
 * for a single input channel.
 */
@Test
public void testSingleChannelNoBarriers() throws Exception {
	BufferOrEvent[] sequence = {
		createBuffer(0, PAGE_SIZE), createBuffer(0, PAGE_SIZE),
		createBuffer(0, PAGE_SIZE), createEndOfPartition(0)
	};

	MockInputGate gate = new MockInputGate(PAGE_SIZE, 1, Arrays.asList(sequence));
	BarrierBuffer buffer = createBarrierHandler(gate);

	for (BufferOrEvent boe : sequence) {
		assertEquals(boe, buffer.getNextNonBlocked());
	}

	assertEquals(0L, buffer.getAlignmentDurationNanos());

	assertNull(buffer.getNextNonBlocked());
	assertNull(buffer.getNextNonBlocked());

	buffer.cleanup();
}
 

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());
	}
}
 

private void processBufferOrEvent(BufferOrEvent bufferOrEvent) throws IOException {
	if (bufferOrEvent.isBuffer()) {
		lastChannel = channelIndexes.get(bufferOrEvent.getChannelInfo());
		checkState(lastChannel != StreamTaskInput.UNSPECIFIED);
		currentRecordDeserializer = recordDeserializers[lastChannel];
		checkState(currentRecordDeserializer != null,
			"currentRecordDeserializer has already been released");

		currentRecordDeserializer.setNextBuffer(bufferOrEvent.getBuffer());
	}
	else {
		// Event received
		final AbstractEvent event = bufferOrEvent.getEvent();
		// TODO: with checkpointedInputGate.isFinished() we might not need to support any events on this level.
		if (event.getClass() != EndOfPartitionEvent.class) {
			throw new IOException("Unexpected event: " + event);
		}

		// release the record deserializer immediately,
		// which is very valuable in case of bounded stream
		releaseDeserializer(channelIndexes.get(bufferOrEvent.getChannelInfo()));
	}
}
 

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

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

@Test
public void testMultiChannelNoBarriers() {
	try {
		BufferOrEvent[] sequence = { createBuffer(2), createBuffer(2), createBuffer(0),
				createBuffer(1), createBuffer(0), createBuffer(3),
				createBuffer(1), createBuffer(1), createBuffer(2)
		};

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

		for (BufferOrEvent boe : sequence) {
			assertEquals(boe, tracker.getNextNonBlocked());
		}

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

@Test
public void testMissingCancellationBarriers() throws Exception {
	BufferOrEvent[] sequence = {
		createBarrier(1L, 0),
		createCancellationBarrier(2L, 0),
		createCancellationBarrier(3L, 0),
		createCancellationBarrier(3L, 1),
		createBuffer(0)
	};
	AbstractInvokable validator = new CheckpointSequenceValidator(-3);
	inputGate = createBarrierBuffer(2, sequence, validator);

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

@Test
public void testSingleChannelNoBarriers() {
	try {
		BufferOrEvent[] sequence = { createBuffer(0), createBuffer(0), createBuffer(0) };

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

		for (BufferOrEvent boe : sequence) {
			assertEquals(boe, tracker.getNextNonBlocked());
		}

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

private static void validateBuffer(BufferOrEvent boe, int expectedSize, int expectedChannelIndex) {
	assertEquals("wrong channel index", expectedChannelIndex, boe.getChannelIndex());
	assertTrue("is not buffer", boe.isBuffer());

	Buffer buf = boe.getBuffer();
	assertEquals("wrong buffer size", expectedSize, buf.getSize());

	MemorySegment seg = buf.getMemorySegment();
	for (int i = 0; i < expectedSize; i++) {
		byte expected = (byte) i;
		if (expected != seg.get(i)) {
			fail(String.format(
				"wrong buffer contents at position %s : expected=%d , found=%d", i, expected, seg.get(i)));
		}
	}
}
 

@Override
public Optional<BufferOrEvent> getNext() throws IOException, InterruptedException {
	currentChannel = (currentChannel + 1) % numberOfChannels;

	if (barrierGens[currentChannel].isNextBarrier()) {
		return Optional.of(
			new BufferOrEvent(
				new CheckpointBarrier(
					++currentBarriers[currentChannel],
					System.currentTimeMillis(),
					CheckpointOptions.forCheckpointWithDefaultLocation()),
				currentChannel));
	} else {
		Buffer buffer = bufferPools[currentChannel].requestBuffer();
		buffer.getMemorySegment().putLong(0, c++);
		return Optional.of(new BufferOrEvent(buffer, currentChannel));
	}
}
 

private static BufferOrEvent generateAndWriteEvent(FileChannel fileChannel, Random rnd, int numberOfChannels) throws IOException {
	long magicNumber = rnd.nextLong();
	byte[] data = new byte[rnd.nextInt(1000)];
	rnd.nextBytes(data);
	TestEvent evt = new TestEvent(magicNumber, data);

	int channelIndex = rnd.nextInt(numberOfChannels);

	ByteBuffer serializedEvent = EventSerializer.toSerializedEvent(evt);
	ByteBuffer header = ByteBuffer.allocate(9);
	header.order(ByteOrder.LITTLE_ENDIAN);

	header.putInt(channelIndex);
	header.putInt(serializedEvent.remaining());
	header.put((byte) 1);
	header.flip();

	FileUtils.writeCompletely(fileChannel, header);
	FileUtils.writeCompletely(fileChannel, serializedEvent);
	return new BufferOrEvent(evt, channelIndex);
}
 

@Test
public void testEndOfStreamWhileCheckpoint() throws Exception {
	ValidatingCheckpointHandler handler = new ValidatingCheckpointHandler(1);
	inputGate = createInputGate(3, handler);

	// one checkpoint
	final BufferOrEvent[] sequence1 = addSequence(inputGate,
		createBarrier(1, 0), createBarrier(1, 1), createBarrier(1, 2));
	assertOutput(sequence1);
	assertEquals(1L, channelStateWriter.getLastStartedCheckpointId());
	assertInflightData();

	final BufferOrEvent[] sequence2 = addSequence(inputGate,
		// some buffers
		createBuffer(0), createBuffer(0), createBuffer(2),

		// start the checkpoint that will be incomplete
		createBarrier(2, 2), createBarrier(2, 0),
		createBuffer(0), createBuffer(2), createBuffer(1),

		// close one after the barrier one before the barrier
		createEndOfPartition(2), createEndOfPartition(1),
		createBuffer(0),

		// final end of stream
		createEndOfPartition(0));
	assertOutput(sequence2);
	assertEquals(2L, channelStateWriter.getLastStartedCheckpointId());
	assertInflightData(sequence2[7]);
}
 

@Test
public void testSingleChannelAbortCheckpoint() throws Exception {
	BufferOrEvent[] sequence = {
			createBuffer(0),
			createBarrier(1, 0),
			createBuffer(0),
			createBarrier(2, 0),
			createCancellationBarrier(4, 0),
			createBarrier(5, 0),
			createBuffer(0),
			createCancellationBarrier(6, 0),
			createBuffer(0)
	};

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

	// negative values mean an expected cancellation call!
	CheckpointSequenceValidator validator =
			new CheckpointSequenceValidator(1, 2, -4, 5, -6);
	tracker.registerCheckpointEventHandler(validator);

	for (BufferOrEvent boe : sequence) {
		if (boe.isBuffer()) {
			assertEquals(boe, tracker.getNextNonBlocked());
		}
		assertTrue(tracker.isEmpty());
	}

	assertNull(tracker.getNextNonBlocked());
	assertNull(tracker.getNextNonBlocked());
}
 

public static BufferOrEvent generateRandomBuffer(int size, int channelIndex) {
	MemorySegment seg = MemorySegmentFactory.allocateUnpooledSegment(PAGE_SIZE);
	for (int i = 0; i < size; i++) {
		seg.put(i, (byte) i);
	}

	Buffer buf = new NetworkBuffer(seg, FreeingBufferRecycler.INSTANCE);
	buf.setSize(size);
	return new BufferOrEvent(buf, channelIndex);
}
 

private static BufferOrEvent generateRandomEvent(Random rnd, int numberOfChannels) {
	long magicNumber = rnd.nextLong();
	byte[] data = new byte[rnd.nextInt(1000)];
	rnd.nextBytes(data);
	TestEvent evt = new TestEvent(magicNumber, data);

	int channelIndex = rnd.nextInt(numberOfChannels);

	return new BufferOrEvent(evt, channelIndex);
}
 

@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);
}
 

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

@Test
public void testSingleChannelAbortCheckpoint() throws Exception {
	ValidatingCheckpointHandler handler = new ValidatingCheckpointHandler(1);
	inputGate = createInputGate(1, handler);
	final BufferOrEvent[] sequence1 = addSequence(
		inputGate,
		createBuffer(0),
		createBarrier(1, 0),
		createBuffer(0),
		createBarrier(2, 0),
		createCancellationBarrier(4, 0));

	assertOutput(sequence1);
	assertEquals(2L, channelStateWriter.getLastStartedCheckpointId());
	assertEquals(4L, handler.getLastCanceledCheckpointId());
	assertInflightData();

	final BufferOrEvent[] sequence2 = addSequence(
		inputGate,
		createBarrier(5, 0),
		createBuffer(0),
		createCancellationBarrier(6, 0),
		createBuffer(0),
		createEndOfPartition(0));

	assertOutput(sequence2);
	assertEquals(5L, channelStateWriter.getLastStartedCheckpointId());
	assertEquals(6L, handler.getLastCanceledCheckpointId());
	assertInflightData();
}
 

public static BufferOrEvent generateRandomBuffer(int size, int channelIndex) {
	MemorySegment seg = MemorySegmentFactory.allocateUnpooledSegment(PAGE_SIZE);
	for (int i = 0; i < size; i++) {
		seg.put(i, (byte) i);
	}

	Buffer buf = new NetworkBuffer(seg, FreeingBufferRecycler.INSTANCE);
	buf.setSize(size);
	return new BufferOrEvent(buf, channelIndex);
}
 

@Test
public void testMultiChannelWithBarriers() 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),
			createBarrier(3, 0), createBarrier(3, 1),

			createBarrier(4, 1), createBarrier(4, 2), createBarrier(4, 0),

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

	for (BufferOrEvent boe : sequence) {
		assertEquals(boe, inputGate.pollNext().get());
	}
}
 

private static CheckpointedInputGate createBarrierTracker(
		int numberOfChannels,
		BufferOrEvent[] sequence,
		@Nullable AbstractInvokable toNotifyOnCheckpoint) {
	MockInputGate gate = new MockInputGate(numberOfChannels, Arrays.asList(sequence));
	return new CheckpointedInputGate(
		gate,
		new CheckpointBarrierTracker(gate.getNumberOfInputChannels(), toNotifyOnCheckpoint));
}
 

@Test
public void testSingleChannelAbortCheckpoint() throws Exception {
	BufferOrEvent[] sequence = {
		createBuffer(0, PAGE_SIZE),
		createBarrier(1, 0),
		createBuffer(0, PAGE_SIZE),
		createBarrier(2, 0),
		createCancellationBarrier(4, 0),
		createBarrier(5, 0),
		createBuffer(0, PAGE_SIZE),
		createCancellationBarrier(6, 0),
		createBuffer(0, PAGE_SIZE)
	};

	MockInputGate gate = new MockInputGate(PAGE_SIZE, 1, Arrays.asList(sequence));
	BarrierBuffer buffer = createBarrierHandler(gate);

	AbstractInvokable toNotify = mock(AbstractInvokable.class);
	buffer.registerCheckpointEventHandler(toNotify);

	check(sequence[0], buffer.getNextNonBlocked(), PAGE_SIZE);
	check(sequence[2], buffer.getNextNonBlocked(), PAGE_SIZE);
	verify(toNotify, times(1)).triggerCheckpointOnBarrier(argThat(new CheckpointMatcher(1L)), any(CheckpointOptions.class), any(CheckpointMetrics.class));
	assertEquals(0L, buffer.getAlignmentDurationNanos());

	check(sequence[6], buffer.getNextNonBlocked(), PAGE_SIZE);
	assertEquals(5L, buffer.getCurrentCheckpointId());
	verify(toNotify, times(1)).triggerCheckpointOnBarrier(argThat(new CheckpointMatcher(2L)), any(CheckpointOptions.class), any(CheckpointMetrics.class));
	verify(toNotify, times(1)).abortCheckpointOnBarrier(eq(4L), any(CheckpointDeclineOnCancellationBarrierException.class));
	verify(toNotify, times(1)).triggerCheckpointOnBarrier(argThat(new CheckpointMatcher(5L)), any(CheckpointOptions.class), any(CheckpointMetrics.class));
	assertEquals(0L, buffer.getAlignmentDurationNanos());

	check(sequence[8], buffer.getNextNonBlocked(), PAGE_SIZE);
	assertEquals(6L, buffer.getCurrentCheckpointId());
	verify(toNotify, times(1)).abortCheckpointOnBarrier(eq(6L), any(CheckpointDeclineOnCancellationBarrierException.class));
	assertEquals(0L, buffer.getAlignmentDurationNanos());

	buffer.cleanup();
}
 

private static BufferOrEvent createBuffer(int channel, int size) {
	byte[] bytes = new byte[size];
	RND.nextBytes(bytes);

	MemorySegment memory = MemorySegmentFactory.allocateUnpooledSegment(PAGE_SIZE);
	memory.put(0, bytes);

	Buffer buf = new NetworkBuffer(memory, FreeingBufferRecycler.INSTANCE);
	buf.setSize(size);

	// retain an additional time so it does not get disposed after being read by the input gate
	buf.retainBuffer();

	return new BufferOrEvent(buf, channel);
}
 

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

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

@Test
public void testIsAvailableWithBufferedDataInDeserializer() throws Exception {
	List<BufferOrEvent> buffers = Collections.singletonList(createDataBuffer());

	VerifyRecordsDataOutput output = new VerifyRecordsDataOutput<>();
	StreamTaskNetworkInput input = createStreamTaskNetworkInput(buffers, output);

	assertHasNextElement(input, output);
	assertHasNextElement(input, output);
	assertEquals(2, output.getNumberOfEmittedRecords());
}
 

public MockInputGate(
		int numberOfChannels,
		List<BufferOrEvent> bufferOrEvents,
		boolean finishAfterLastBuffer) {
	this.numberOfChannels = numberOfChannels;
	this.bufferOrEvents = new ArrayDeque<BufferOrEvent>(bufferOrEvents);
	this.closed = new boolean[numberOfChannels];
	this.finishAfterLastBuffer = finishAfterLastBuffer;

	availabilityHelper.resetAvailable();
}