org.apache.flink.runtime.io.network.api.EndOfPartitionEvent Java Examples

/**
 * Tests {@link EventSerializer#isEvent(Buffer, Class)}
 * whether it peaks into the buffer only, i.e. after the call, the buffer
 * is still de-serializable.
 */
@Test
public void testIsEventPeakOnly() throws Exception {
	final Buffer serializedEvent =
		EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE);
	try {
		final ClassLoader cl = getClass().getClassLoader();
		assertTrue(
			EventSerializer.isEvent(serializedEvent, EndOfPartitionEvent.class));
		EndOfPartitionEvent event = (EndOfPartitionEvent) EventSerializer
			.fromBuffer(serializedEvent, cl);
		assertEquals(EndOfPartitionEvent.INSTANCE, event);
	} finally {
		serializedEvent.recycleBuffer();
	}
}
 

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 Optional<BufferOrEvent> getNext() {
	BufferOrEvent next = bufferOrEvents.poll();
	if (!finishAfterLastBuffer && bufferOrEvents.isEmpty()) {
		resetIsAvailable();
	}
	if (next == null) {
		return Optional.empty();
	}

	int channelIdx = next.getChannelIndex();
	if (closed[channelIdx]) {
		throw new RuntimeException("Inconsistent: Channel " + channelIdx
			+ " has data even though it is already closed.");
	}
	if (next.isEvent() && next.getEvent() instanceof EndOfPartitionEvent) {
		closed[channelIdx] = true;
	}
	return Optional.of(next);
}
 

/**
 * Tests {@link EventSerializer#isEvent(Buffer, Class)}
 * whether it peaks into the buffer only, i.e. after the call, the buffer
 * is still de-serializable.
 */
@Test
public void testIsEventPeakOnly() throws Exception {
	final Buffer serializedEvent =
		EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE);
	try {
		final ClassLoader cl = getClass().getClassLoader();
		assertTrue(
			EventSerializer.isEvent(serializedEvent, EndOfPartitionEvent.class));
		EndOfPartitionEvent event = (EndOfPartitionEvent) EventSerializer
			.fromBuffer(serializedEvent, cl);
		assertEquals(EndOfPartitionEvent.INSTANCE, event);
	} finally {
		serializedEvent.recycleBuffer();
	}
}
 

private void verifyBroadcastBufferOrEventIndependence(boolean broadcastEvent) throws Exception {
	@SuppressWarnings("unchecked")
	ArrayDeque<BufferConsumer>[] queues = new ArrayDeque[]{new ArrayDeque(), new ArrayDeque()};

	ResultPartitionWriter partition =
		new CollectingPartitionWriter(queues, new TestPooledBufferProvider(Integer.MAX_VALUE));
	RecordWriter<IntValue> writer = createRecordWriter(partition);

	if (broadcastEvent) {
		writer.broadcastEvent(EndOfPartitionEvent.INSTANCE);
	} else {
		writer.broadcastEmit(new IntValue(0));
	}

	// verify added to all queues
	assertEquals(1, queues[0].size());
	assertEquals(1, queues[1].size());

	// these two buffers may share the memory but not the indices!
	Buffer buffer1 = buildSingleBuffer(queues[0].remove());
	Buffer buffer2 = buildSingleBuffer(queues[1].remove());
	assertEquals(0, buffer1.getReaderIndex());
	assertEquals(0, buffer2.getReaderIndex());
	buffer1.setReaderIndex(1);
	assertEquals("Buffer 2 shares the same reader index as buffer 1", 0, buffer2.getReaderIndex());
}
 

/**
 * Tests that broadcasted events' buffers are independent (in their (reader) indices) once they
 * are put into the queue for Netty when broadcasting events to multiple channels.
 */
@Test
public void testBroadcastEventBufferIndependence() throws Exception {
	@SuppressWarnings("unchecked")
	ArrayDeque<BufferConsumer>[] queues =
		new ArrayDeque[]{new ArrayDeque(), new ArrayDeque()};

	ResultPartitionWriter partition =
		new CollectingPartitionWriter(queues, new TestPooledBufferProvider(Integer.MAX_VALUE));
	RecordWriter<?> writer = new RecordWriter<>(partition);

	writer.broadcastEvent(EndOfPartitionEvent.INSTANCE);

	// Verify added to all queues
	assertEquals(1, queues[0].size());
	assertEquals(1, queues[1].size());

	// these two buffers may share the memory but not the indices!
	Buffer buffer1 = buildSingleBuffer(queues[0].remove());
	Buffer buffer2 = buildSingleBuffer(queues[1].remove());
	assertEquals(0, buffer1.getReaderIndex());
	assertEquals(0, buffer2.getReaderIndex());
	buffer1.setReaderIndex(1);
	assertEquals("Buffer 2 shares the same reader index as buffer 1", 0, buffer2.getReaderIndex());
}
 

@Override
public Optional<BufferOrEvent> getNextBufferOrEvent() {
	BufferOrEvent next = bufferOrEvents.poll();
	if (next == null) {
		return Optional.empty();
	}

	int channelIdx = next.getChannelIndex();
	if (closed[channelIdx]) {
		throw new RuntimeException("Inconsistent: Channel " + channelIdx
			+ " has data even though it is already closed.");
	}
	if (next.isEvent() && next.getEvent() instanceof EndOfPartitionEvent) {
		closed[channelIdx] = true;
		closedChannels++;
	}
	return Optional.of(next);
}
 

private void verifyBroadcastBufferOrEventIndependence(boolean broadcastEvent) throws Exception {
	@SuppressWarnings("unchecked")
	ArrayDeque<BufferConsumer>[] queues = new ArrayDeque[]{new ArrayDeque(), new ArrayDeque()};

	ResultPartitionWriter partition =
		new CollectingPartitionWriter(queues, new TestPooledBufferProvider(Integer.MAX_VALUE));
	RecordWriter<IntValue> writer = new RecordWriterBuilder().build(partition);

	if (broadcastEvent) {
		writer.broadcastEvent(EndOfPartitionEvent.INSTANCE);
	} else {
		writer.broadcastEmit(new IntValue(0));
	}

	// verify added to all queues
	assertEquals(1, queues[0].size());
	assertEquals(1, queues[1].size());

	// these two buffers may share the memory but not the indices!
	Buffer buffer1 = buildSingleBuffer(queues[0].remove());
	Buffer buffer2 = buildSingleBuffer(queues[1].remove());
	assertEquals(0, buffer1.getReaderIndex());
	assertEquals(0, buffer2.getReaderIndex());
	buffer1.setReaderIndex(1);
	assertEquals("Buffer 2 shares the same reader index as buffer 1", 0, buffer2.getReaderIndex());
}
 

@Test
public void testSerializeDeserializeEvent() throws Exception {
	AbstractEvent[] events = {
			EndOfPartitionEvent.INSTANCE,
			EndOfSuperstepEvent.INSTANCE,
			new CheckpointBarrier(1678L, 4623784L, CheckpointOptions.forCheckpointWithDefaultLocation()),
			new TestTaskEvent(Math.random(), 12361231273L),
			new CancelCheckpointMarker(287087987329842L)
	};

	for (AbstractEvent evt : events) {
		ByteBuffer serializedEvent = EventSerializer.toSerializedEvent(evt);
		assertTrue(serializedEvent.hasRemaining());

		AbstractEvent deserialized =
				EventSerializer.fromSerializedEvent(serializedEvent, getClass().getClassLoader());
		assertNotNull(deserialized);
		assertEquals(evt, deserialized);
	}
}
 

@Test
public void testSerializeDeserializeEvent() throws Exception {
	AbstractEvent[] events = {
			EndOfPartitionEvent.INSTANCE,
			EndOfSuperstepEvent.INSTANCE,
			new CheckpointBarrier(1678L, 4623784L, CheckpointOptions.forCheckpointWithDefaultLocation()),
			new TestTaskEvent(Math.random(), 12361231273L),
			new CancelCheckpointMarker(287087987329842L)
	};

	for (AbstractEvent evt : events) {
		ByteBuffer serializedEvent = EventSerializer.toSerializedEvent(evt);
		assertTrue(serializedEvent.hasRemaining());

		AbstractEvent deserialized =
				EventSerializer.fromSerializedEvent(serializedEvent, getClass().getClassLoader());
		assertNotNull(deserialized);
		assertEquals(evt, deserialized);
	}
}
 

@Override
public synchronized void finish() throws IOException {
	synchronized (buffers) {
		if (add(EventSerializer.toBufferConsumer(EndOfPartitionEvent.INSTANCE), true)) {
			isFinished = true;
		}

		flush();
	}

	// If we are spilling/have spilled, wait for the writer to finish
	if (spillWriter != null) {
		spillWriter.close();
	}
	LOG.debug("{}: Finished {}.", parent.getOwningTaskName(), this);
}
 

/**
 * Tests {@link EventSerializer#isEvent(Buffer, Class)}
 * whether it peaks into the buffer only, i.e. after the call, the buffer
 * is still de-serializable.
 */
@Test
public void testIsEventPeakOnly() throws Exception {
	final Buffer serializedEvent =
		EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE);
	try {
		final ClassLoader cl = getClass().getClassLoader();
		assertTrue(
			EventSerializer.isEvent(serializedEvent, EndOfPartitionEvent.class));
		EndOfPartitionEvent event = (EndOfPartitionEvent) EventSerializer
			.fromBuffer(serializedEvent, cl);
		assertEquals(EndOfPartitionEvent.INSTANCE, event);
	} finally {
		serializedEvent.recycleBuffer();
	}
}
 

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

@Override
public Optional<BufferOrEvent> getNext() {
	BufferOrEvent next = bufferOrEvents.poll();
	if (!finishAfterLastBuffer && bufferOrEvents.isEmpty()) {
		availabilityHelper.resetUnavailable();
	}
	if (next == null) {
		return Optional.empty();
	}

	int channelIdx = next.getChannelInfo().getInputChannelIdx();
	if (closed[channelIdx]) {
		throw new RuntimeException("Inconsistent: Channel " + channelIdx
			+ " has data even though it is already closed.");
	}
	if (next.isEvent() && next.getEvent() instanceof EndOfPartitionEvent) {
		closed[channelIdx] = true;
	}
	return Optional.of(next);
}
 

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

@Test
public void testForwardsIsBuffer() throws IOException {
	assertEquals(buffer.isBuffer(), buffer.readOnlySlice().isBuffer());
	assertEquals(buffer.isBuffer(), buffer.readOnlySlice(1, 2).isBuffer());
	Buffer eventBuffer = EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE);
	assertEquals(eventBuffer.isBuffer(), eventBuffer.readOnlySlice().isBuffer());
	assertEquals(eventBuffer.isBuffer(), eventBuffer.readOnlySlice(1, 2).isBuffer());
}
 

@Test
public void testForwardsIsBuffer() throws IOException {
	assertEquals(buffer.isBuffer(), buffer.readOnlySlice().isBuffer());
	assertEquals(buffer.isBuffer(), buffer.readOnlySlice(1, 2).isBuffer());
	Buffer eventBuffer = EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE);
	assertEquals(eventBuffer.isBuffer(), eventBuffer.readOnlySlice().isBuffer());
	assertEquals(eventBuffer.isBuffer(), eventBuffer.readOnlySlice(1, 2).isBuffer());
}
 

private IteratorWrappingTestSingleInputGate<T> wrapIterator(MutableObjectIterator<T> iterator) throws IOException, InterruptedException {
	inputIterator = iterator;
	serializer = new SpanningRecordSerializer<T>();

	// The input iterator can produce an infinite stream. That's why we have to serialize each
	// record on demand and cannot do it upfront.
	final BufferAndAvailabilityProvider answer = new BufferAndAvailabilityProvider() {

		private boolean hasData = inputIterator.next(reuse) != null;

		@Override
		public Optional<BufferAndAvailability> getBufferAvailability() throws IOException {
			if (hasData) {
				serializer.serializeRecord(reuse);
				BufferBuilder bufferBuilder = createBufferBuilder(bufferSize);
				BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
				serializer.copyToBufferBuilder(bufferBuilder);

				hasData = inputIterator.next(reuse) != null;

				// Call getCurrentBuffer to ensure size is set
				return Optional.of(new BufferAndAvailability(bufferConsumer.build(), true, 0));
			} else {
				inputChannel.setReleased();

				return Optional.of(new BufferAndAvailability(EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE),
					false,
					0));
			}
		}
	};

	inputChannel.addBufferAndAvailability(answer);

	inputGate.setInputChannels(inputChannel);

	return this;
}
 

TestInputChannel readEndOfPartitionEvent() {
	addBufferAndAvailability(
		() -> {
			setReleased();
			return Optional.of(new BufferAndAvailability(EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE),
				false,
				0));
		}
	);
	return this;
}
 

private IteratorWrappingTestSingleInputGate<T> wrapIterator(MutableObjectIterator<T> iterator) throws IOException, InterruptedException {
	inputIterator = iterator;
	serializer = new SpanningRecordSerializer<T>();

	// The input iterator can produce an infinite stream. That's why we have to serialize each
	// record on demand and cannot do it upfront.
	final BufferAndAvailabilityProvider answer = new BufferAndAvailabilityProvider() {

		private boolean hasData = inputIterator.next(reuse) != null;

		@Override
		public Optional<BufferAndAvailability> getBufferAvailability() throws IOException {
			if (hasData) {
				serializer.serializeRecord(reuse);
				BufferBuilder bufferBuilder = createBufferBuilder(bufferSize);
				serializer.copyToBufferBuilder(bufferBuilder);

				hasData = inputIterator.next(reuse) != null;

				// Call getCurrentBuffer to ensure size is set
				return Optional.of(new BufferAndAvailability(buildSingleBuffer(bufferBuilder), true, 0));
			} else {
				inputChannel.setReleased();

				return Optional.of(new BufferAndAvailability(EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE),
					false,
					0));
			}
		}
	};

	inputChannel.addBufferAndAvailability(answer);

	inputGate.setInputChannel(new IntermediateResultPartitionID(), inputChannel);

	return this;
}
 

/**
 * Handles the event and returns whether the reader reached an end-of-stream event (either the
 * end of the whole stream or the end of an superstep).
 */
protected boolean handleEvent(AbstractEvent event) throws IOException {
	final Class<?> eventType = event.getClass();

	try {
		// ------------------------------------------------------------
		// Runtime events
		// ------------------------------------------------------------

		// This event is also checked at the (single) input gate to release the respective
		// channel, at which it was received.
		if (eventType == EndOfPartitionEvent.class) {
			return true;
		}
		else if (eventType == EndOfSuperstepEvent.class) {
			return incrementEndOfSuperstepEventAndCheck();
		}

		// ------------------------------------------------------------
		// Task events (user)
		// ------------------------------------------------------------
		else if (event instanceof TaskEvent) {
			taskEventHandler.publish((TaskEvent) event);

			return false;
		}
		else {
			throw new IllegalStateException("Received unexpected event of type " + eventType + " at reader.");
		}
	}
	catch (Throwable t) {
		throw new IOException("Error while handling event of type " + eventType + ": " + t.getMessage(), t);
	}
}
 

private void handleEndOfPartitionEvent(BufferOrEvent bufferOrEvent, InputGate inputGate) {
	if (bufferOrEvent.isEvent()
		&& bufferOrEvent.getEvent().getClass() == EndOfPartitionEvent.class
		&& inputGate.isFinished()) {

		checkState(!bufferOrEvent.moreAvailable());
		if (!inputGatesWithRemainingData.remove(inputGate)) {
			throw new IllegalStateException("Couldn't find input gate in set of remaining " +
				"input gates.");
		}
		if (isFinished()) {
			markAvailable();
		}
	}
}
 

@Test
public void testReleasingDeserializerTimely()
	throws Exception {

	int numInputChannels = 2;
	LongSerializer inSerializer = LongSerializer.INSTANCE;
	StreamTestSingleInputGate inputGate = new StreamTestSingleInputGate<>(numInputChannels, 0, inSerializer, 1024);

	TestRecordDeserializer[] deserializers = new TestRecordDeserializer[numInputChannels];
	for (int i = 0; i < deserializers.length; i++) {
		deserializers[i] = new TestRecordDeserializer(ioManager.getSpillingDirectoriesPaths());
	}

	TestRecordDeserializer[] copiedDeserializers = Arrays.copyOf(deserializers, deserializers.length);
	DataOutput output = new NoOpDataOutput<>();
	StreamTaskNetworkInput input = new StreamTaskNetworkInput<>(
		new CheckpointedInputGate(
			inputGate.getInputGate(),
			new CheckpointBarrierTracker(1, new DummyCheckpointInvokable())),
		inSerializer,
		new StatusWatermarkValve(1, output),
		0,
		deserializers);

	for (int i = 0; i < numInputChannels; i++) {
		assertNotNull(deserializers[i]);
		inputGate.sendEvent(EndOfPartitionEvent.INSTANCE, i);
		input.emitNext(output);
		assertNull(deserializers[i]);
		assertTrue(copiedDeserializers[i].isCleared());
	}
}
 

/**
 * Tests that event buffers are properly recycled when broadcasting events
 * to multiple channels.
 */
@Test
public void testBroadcastEventBufferReferenceCounting() throws Exception {

	@SuppressWarnings("unchecked")
	ArrayDeque<BufferConsumer>[] queues = new ArrayDeque[] { new ArrayDeque(), new ArrayDeque() };

	ResultPartitionWriter partition =
		new CollectingPartitionWriter(queues, new TestPooledBufferProvider(Integer.MAX_VALUE));
	RecordWriter<?> writer = createRecordWriter(partition);

	writer.broadcastEvent(EndOfPartitionEvent.INSTANCE);

	// Verify added to all queues
	assertEquals(1, queues[0].size());
	assertEquals(1, queues[1].size());

	// get references to buffer consumers (copies from the original event buffer consumer)
	BufferConsumer bufferConsumer1 = queues[0].getFirst();
	BufferConsumer bufferConsumer2 = queues[1].getFirst();

	// process all collected events (recycles the buffer)
	for (int i = 0; i < queues.length; i++) {
		assertTrue(parseBuffer(queues[i].remove(), i).isEvent());
	}

	assertTrue(bufferConsumer1.isRecycled());
	assertTrue(bufferConsumer2.isRecycled());
}
 

public static ByteBuffer toSerializedEvent(AbstractEvent event) throws IOException {
	final Class<?> eventClass = event.getClass();
	if (eventClass == EndOfPartitionEvent.class) {
		return ByteBuffer.wrap(new byte[] { 0, 0, 0, END_OF_PARTITION_EVENT });
	}
	else if (eventClass == CheckpointBarrier.class) {
		return serializeCheckpointBarrier((CheckpointBarrier) event);
	}
	else if (eventClass == EndOfSuperstepEvent.class) {
		return ByteBuffer.wrap(new byte[] { 0, 0, 0, END_OF_SUPERSTEP_EVENT });
	}
	else if (eventClass == EndOfChannelStateEvent.class) {
		return ByteBuffer.wrap(new byte[] { 0, 0, 0, END_OF_CHANNEL_STATE_EVENT });
	}
	else if (eventClass == CancelCheckpointMarker.class) {
		CancelCheckpointMarker marker = (CancelCheckpointMarker) event;

		ByteBuffer buf = ByteBuffer.allocate(12);
		buf.putInt(0, CANCEL_CHECKPOINT_MARKER_EVENT);
		buf.putLong(4, marker.getCheckpointId());
		return buf;
	}
	else {
		try {
			final DataOutputSerializer serializer = new DataOutputSerializer(128);
			serializer.writeInt(OTHER_EVENT);
			serializer.writeUTF(event.getClass().getName());
			event.write(serializer);
			return serializer.wrapAsByteBuffer();
		}
		catch (IOException e) {
			throw new IOException("Error while serializing event.", e);
		}
	}
}
 

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

private BufferOrEvent transformEvent(
		Buffer buffer,
		boolean moreAvailable,
		InputChannel currentChannel) throws IOException, InterruptedException {
	final AbstractEvent event;
	try {
		event = EventSerializer.fromBuffer(buffer, getClass().getClassLoader());
	} finally {
		buffer.recycleBuffer();
	}

	if (event.getClass() == EndOfPartitionEvent.class) {
		channelsWithEndOfPartitionEvents.set(currentChannel.getChannelIndex());

		if (channelsWithEndOfPartitionEvents.cardinality() == numberOfInputChannels) {
			// Because of race condition between:
			// 1. releasing inputChannelsWithData lock in this method and reaching this place
			// 2. empty data notification that re-enqueues a channel
			// we can end up with moreAvailable flag set to true, while we expect no more data.
			checkState(!moreAvailable || !pollNext().isPresent());
			moreAvailable = false;
			hasReceivedAllEndOfPartitionEvents = true;
			markAvailable();
		}

		currentChannel.releaseAllResources();
	}

	return new BufferOrEvent(event, currentChannel.getChannelInfo(), moreAvailable, buffer.getSize());
}
 

public static ByteBuffer toSerializedEvent(AbstractEvent event) throws IOException {
	final Class<?> eventClass = event.getClass();
	if (eventClass == EndOfPartitionEvent.class) {
		return ByteBuffer.wrap(new byte[] { 0, 0, 0, END_OF_PARTITION_EVENT });
	}
	else if (eventClass == CheckpointBarrier.class) {
		return serializeCheckpointBarrier((CheckpointBarrier) event);
	}
	else if (eventClass == EndOfSuperstepEvent.class) {
		return ByteBuffer.wrap(new byte[] { 0, 0, 0, END_OF_SUPERSTEP_EVENT });
	}
	else if (eventClass == CancelCheckpointMarker.class) {
		CancelCheckpointMarker marker = (CancelCheckpointMarker) event;

		ByteBuffer buf = ByteBuffer.allocate(12);
		buf.putInt(0, CANCEL_CHECKPOINT_MARKER_EVENT);
		buf.putLong(4, marker.getCheckpointId());
		return buf;
	}
	else {
		try {
			final DataOutputSerializer serializer = new DataOutputSerializer(128);
			serializer.writeInt(OTHER_EVENT);
			serializer.writeUTF(event.getClass().getName());
			event.write(serializer);
			return serializer.wrapAsByteBuffer();
		}
		catch (IOException e) {
			throw new IOException("Error while serializing event.", e);
		}
	}
}
 

TestInputChannel readEndOfPartitionEvent() {
	addBufferAndAvailability(
		() -> {
			setReleased();
			return Optional.of(new BufferAndAvailability(EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE),
				false,
				0));
		}
	);
	return this;
}
 

/**
 * Tests that event buffers are properly recycled when broadcasting events
 * to multiple channels.
 */
@Test
public void testBroadcastEventBufferReferenceCounting() throws Exception {

	@SuppressWarnings("unchecked")
	ArrayDeque<BufferConsumer>[] queues = new ArrayDeque[] { new ArrayDeque(), new ArrayDeque() };

	ResultPartitionWriter partition =
		new CollectingPartitionWriter(queues, new TestPooledBufferProvider(Integer.MAX_VALUE));
	RecordWriter<?> writer = new RecordWriterBuilder().build(partition);

	writer.broadcastEvent(EndOfPartitionEvent.INSTANCE);

	// Verify added to all queues
	assertEquals(1, queues[0].size());
	assertEquals(1, queues[1].size());

	// get references to buffer consumers (copies from the original event buffer consumer)
	BufferConsumer bufferConsumer1 = queues[0].getFirst();
	BufferConsumer bufferConsumer2 = queues[1].getFirst();

	// process all collected events (recycles the buffer)
	for (int i = 0; i < queues.length; i++) {
		assertTrue(parseBuffer(queues[i].remove(), i).isEvent());
	}

	assertTrue(bufferConsumer1.isRecycled());
	assertTrue(bufferConsumer2.isRecycled());
}