org.apache.flink.runtime.io.network.partition.ResultPartitionID Java Examples

@Override
public void onPartitionConsumable(ExecutionVertexID executionVertexId, ResultPartitionID resultPartitionId) {
	final SchedulingResultPartition resultPartition = schedulingTopology
		.getResultPartitionOrThrow(resultPartitionId.getPartitionId());

	if (!resultPartition.getPartitionType().isPipelined()) {
		return;
	}

	final SchedulingExecutionVertex producerVertex = schedulingTopology.getVertexOrThrow(executionVertexId);
	if (!producerVertex.getProducedResultPartitions().contains(resultPartition)) {
		throw new IllegalStateException("partition " + resultPartitionId
				+ " is not the produced partition of " + executionVertexId);
	}

	allocateSlotsAndDeployExecutionVertices(resultPartition.getConsumers());
}
 

/**
 * Sends a task event backwards to an intermediate result partition producer.
 *
 * <p>Backwards task events flow between readers and writers and therefore
 * will only work when both are running at the same time, which is only
 * guaranteed to be the case when both the respective producer and
 * consumer task run pipelined.
 */
@Override
public void sendTaskEvent(ResultPartitionID partitionId, TaskEvent event, final RemoteInputChannel inputChannel) throws IOException {
	checkNotClosed();

	tcpChannel.writeAndFlush(new TaskEventRequest(event, partitionId, inputChannel.getInputChannelId()))
			.addListener(
					new ChannelFutureListener() {
						@Override
						public void operationComplete(ChannelFuture future) throws Exception {
							if (!future.isSuccess()) {
								SocketAddress remoteAddr = future.channel().remoteAddress();
								inputChannel.onError(new LocalTransportException(
									String.format("Sending the task event to '%s' failed.", remoteAddr),
									future.channel().localAddress(), future.cause()
								));
							}
						}
					});
}
 

@Override
public CompletableFuture<ExecutionState> requestPartitionProducerState(
		JobID jobId,
		IntermediateDataSetID intermediateDataSetId,
		ResultPartitionID resultPartitionId) {

	JobManagerMessages.RequestPartitionProducerState msg = new JobManagerMessages.RequestPartitionProducerState(
		jobId,
		intermediateDataSetId, resultPartitionId
	);

	scala.concurrent.Future<ExecutionState> futureResponse = jobManager
		.ask(msg, timeout)
		.mapTo(ClassTag$.MODULE$.<ExecutionState>apply(ExecutionState.class));

	return FutureUtils.toJava(futureResponse);
}
 

/**
 * Creates and returns a remote input channel for the specific input gate with specific partition request client.
 *
 * @param inputGate The input gate owns the created input channel.
 * @param client The client is used to send partition request.
 * @param initialBackoff initial back off (in ms) for retriggering subpartition requests (must be <tt>&gt; 0</tt> to activate)
 * @param maxBackoff after which delay (in ms) to stop retriggering subpartition requests
 * @return The new created remote input channel.
 */
static RemoteInputChannel createRemoteInputChannel(
		SingleInputGate inputGate,
		PartitionRequestClient client,
		int initialBackoff,
		int maxBackoff) throws Exception {
	final ConnectionManager connectionManager = mock(ConnectionManager.class);
	when(connectionManager.createPartitionRequestClient(any(ConnectionID.class)))
		.thenReturn(client);

	ResultPartitionID partitionId = new ResultPartitionID();
	RemoteInputChannel inputChannel = new RemoteInputChannel(
		inputGate,
		0,
		partitionId,
		mock(ConnectionID.class),
		connectionManager,
		initialBackoff,
		maxBackoff,
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());

	inputGate.setInputChannel(partitionId.getPartitionId(), inputChannel);
	return inputChannel;
}
 

/**
 * Tests that {@link PartitionRequestServerHandler} responds {@link ErrorResponse} with wrapped
 * {@link PartitionNotFoundException} after receiving invalid {@link PartitionRequest}.
 */
@Test
public void testResponsePartitionNotFoundException() {
	final PartitionRequestServerHandler serverHandler = new PartitionRequestServerHandler(
		new ResultPartitionManager(),
		new TaskEventDispatcher(),
		new PartitionRequestQueue());
	final EmbeddedChannel channel = new EmbeddedChannel(serverHandler);
	final ResultPartitionID partitionId = new ResultPartitionID();

	// Write the message of partition request to server
	channel.writeInbound(new PartitionRequest(partitionId, 0, new InputChannelID(), 2));
	channel.runPendingTasks();

	// Read the response message after handling partition request
	final Object msg = channel.readOutbound();
	assertThat(msg, instanceOf(ErrorResponse.class));

	final ErrorResponse err = (ErrorResponse) msg;
	assertThat(err.cause, instanceOf(PartitionNotFoundException.class));

	final ResultPartitionID actualPartitionId = ((PartitionNotFoundException) err.cause).getPartitionId();
	assertThat(partitionId, is(actualPartitionId));
}
 

private static TaskExecutorPartitionTracker createPartitionTrackerWithFixedPartitionReport(ShuffleEnvironment<?, ?> shuffleEnvironment) {
	final ClusterPartitionReport.ClusterPartitionReportEntry clusterPartitionReportEntry =
		new ClusterPartitionReport.ClusterPartitionReportEntry(
			new IntermediateDataSetID(),
			Collections.singleton(new ResultPartitionID()),
			4);

	final ClusterPartitionReport clusterPartitionReport = new ClusterPartitionReport(Collections.singletonList(clusterPartitionReportEntry));

	return new TaskExecutorPartitionTrackerImpl(shuffleEnvironment) {
		@Override
		public ClusterPartitionReport createClusterPartitionReport() {
			return clusterPartitionReport;
		}
	};
}
 

protected ResultPartitionWriter createResultPartition(
		JobID jobId,
		ResultPartitionID partitionId,
		NetworkEnvironment environment,
		int channels) throws Exception {

	ResultPartition resultPartition = new ResultPartition(
		"sender task",
		new NoOpTaskActions(),
		jobId,
		partitionId,
		ResultPartitionType.PIPELINED_BOUNDED,
		channels,
		1,
		environment.getResultPartitionManager(),
		new NoOpResultPartitionConsumableNotifier(),
		ioManager,
		false);

	environment.setupPartition(resultPartition);

	return resultPartition;
}
 

/**
 * Tests that an update channel does not trigger a partition request before the UDF has
 * requested any partitions. Otherwise, this can lead to races when registering a listener at
 * the gate (e.g. in UnionInputGate), which can result in missed buffer notifications at the
 * listener.
 */
@Test
public void testUpdateChannelBeforeRequest() throws Exception {
	SingleInputGate inputGate = createInputGate(1);

	ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);

	InputChannel unknown = InputChannelBuilder.newBuilder()
		.setPartitionManager(partitionManager)
		.buildUnknownAndSetToGate(inputGate);

	// Update to a local channel and verify that no request is triggered
	ResultPartitionID resultPartitionID = unknown.getPartitionId();
	ResourceID location = ResourceID.generate();
	inputGate.updateInputChannel(location, createRemoteWithIdAndLocation(resultPartitionID.getPartitionId(), location));

	verify(partitionManager, never()).createSubpartitionView(
		any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class));
}
 

/**
 * Helper to create simple {@link ResultPartition} instance for use by a {@link Task} inside
 * {@link NetworkEnvironment#registerTask(Task)}.
 *
 * @param partitionType
 * 		the produced partition type
 * @param channels
 * 		the number of output channels
 *
 * @return instance with minimal data set and some mocks so that it is useful for {@link
 * NetworkEnvironment#registerTask(Task)}
 */
private static ResultPartition createResultPartition(
		final ResultPartitionType partitionType, final int channels) {
	return new ResultPartition(
		"TestTask-" + partitionType + ":" + channels,
		mock(TaskActions.class),
		new JobID(),
		new ResultPartitionID(),
		partitionType,
		channels,
		channels,
		mock(ResultPartitionManager.class),
		new NoOpResultPartitionConsumableNotifier(),
		mock(IOManager.class),
		false);
}
 

void triggerPartitionStateCheck(ResultPartitionID partitionId) {
	partitionProducerStateProvider.requestPartitionProducerState(
		consumedResultId,
		partitionId,
		((PartitionProducerStateProvider.ResponseHandle responseHandle) -> {
			boolean isProducingState = new RemoteChannelStateChecker(partitionId, owningTaskName)
				.isProducerReadyOrAbortConsumption(responseHandle);
			if (isProducingState) {
				try {
					retriggerPartitionRequest(partitionId.getPartitionId());
				} catch (IOException t) {
					responseHandle.failConsumption(t);
				}
			}
		}));
}
 

private void testPartitionReleaseOnStateTransitionsAfterRunning(Consumer<Execution> stateTransition1, Consumer<Execution> stateTransition2) throws Exception {
	final SimpleAckingTaskManagerGateway taskManagerGateway = new SimpleAckingTaskManagerGateway();
	final CompletableFuture<Tuple2<JobID, Collection<ResultPartitionID>>> releasePartitionsCallFuture = new CompletableFuture<>();
	taskManagerGateway.setReleasePartitionsConsumer(((jobID, partitionIds) -> releasePartitionsCallFuture.complete(Tuple2.of(jobID, partitionIds))));

	final TestingShuffleMaster testingShuffleMaster = new TestingShuffleMaster();

	setupExecutionGraphAndStartRunningJob(ResultPartitionType.PIPELINED, NoOpPartitionTracker.INSTANCE, taskManagerGateway, testingShuffleMaster);

	stateTransition1.accept(execution);
	assertFalse(releasePartitionsCallFuture.isDone());

	stateTransition2.accept(execution);
	assertTrue(releasePartitionsCallFuture.isDone());

	final Tuple2<JobID, Collection<ResultPartitionID>> releasePartitionsCall = releasePartitionsCallFuture.get();
	assertEquals(jobId, releasePartitionsCall.f0);
	assertEquals(Collections.singletonList(descriptor.getShuffleDescriptor().getResultPartitionID()), releasePartitionsCall.f1);

	assertEquals(1, testingShuffleMaster.externallyReleasedPartitions.size());
	assertEquals(descriptor.getShuffleDescriptor(), testingShuffleMaster.externallyReleasedPartitions.poll());
}
 

static TaskEventRequest readFrom(ByteBuf buffer, ClassLoader classLoader) throws IOException {
	// directly deserialize fromNetty's buffer
	int length = buffer.readInt();
	ByteBuffer serializedEvent = buffer.nioBuffer(buffer.readerIndex(), length);
	// assume this event's content is read from the ByteBuf (positions are not shared!)
	buffer.readerIndex(buffer.readerIndex() + length);

	TaskEvent event =
		(TaskEvent) EventSerializer.fromSerializedEvent(serializedEvent, classLoader);

	ResultPartitionID partitionId =
		new ResultPartitionID(
			IntermediateResultPartitionID.fromByteBuf(buffer),
			ExecutionAttemptID.fromByteBuf(buffer));

	InputChannelID receiverId = InputChannelID.fromByteBuf(buffer);

	return new TaskEventRequest(event, partitionId, receiverId);
}
 

@Test
public void testOnFailedPartitionRequest() throws Exception {
	final ConnectionManager connectionManager = mock(ConnectionManager.class);
	when(connectionManager.createPartitionRequestClient(any(ConnectionID.class)))
			.thenReturn(mock(PartitionRequestClient.class));

	final ResultPartitionID partitionId = new ResultPartitionID();

	final SingleInputGate inputGate = mock(SingleInputGate.class);

	final RemoteInputChannel ch = new RemoteInputChannel(
			inputGate,
			0,
			partitionId,
			mock(ConnectionID.class),
			connectionManager,
			UnregisteredMetricGroups.createUnregisteredTaskMetricGroup().getIOMetricGroup());

	ch.onFailedPartitionRequest();

	verify(inputGate).triggerPartitionStateCheck(eq(partitionId));
}
 

@Test
public void testPartitionRequestNoBackoff() throws Exception {
	// Setup
	SingleInputGate inputGate = createSingleInputGate(1);
	ResultPartitionID partitionId = new ResultPartitionID();
	TestVerifyPartitionRequestClient client = new TestVerifyPartitionRequestClient();
	ConnectionManager connectionManager = new TestVerifyConnectionManager(client);
	RemoteInputChannel ch = createRemoteInputChannel(inputGate, connectionManager, partitionId, 0, 0);

	// No delay for first request
	ch.requestSubpartition(0);
	client.verifyResult(partitionId, 0, 0);

	// Exception, because backoff is disabled.
	try {
		ch.retriggerSubpartitionRequest(0);
		ch.getNextBuffer();
		fail("Did not throw expected exception.");
	}
	catch (Exception expected) {
	}
}
 

@Override
public void requestPartitionProducerState(
		final IntermediateDataSetID intermediateDataSetId,
		final ResultPartitionID resultPartitionId,
		Consumer<? super ResponseHandle> responseConsumer) {

	final CompletableFuture<ExecutionState> futurePartitionState =
		partitionProducerStateChecker.requestPartitionProducerState(
			jobId,
			intermediateDataSetId,
			resultPartitionId);

	FutureUtils.assertNoException(
		futurePartitionState
			.handle(PartitionProducerStateResponseHandle::new)
			.thenAcceptAsync(responseConsumer, executor));
}
 

@Override
public void requestSubpartitionView(
	ResultPartitionProvider partitionProvider,
	ResultPartitionID resultPartitionId,
	int subPartitionIndex) throws IOException {

	synchronized (requestLock) {
		if (subpartitionView == null) {
			// This this call can trigger a notification we have to
			// schedule a separate task at the event loop that will
			// start consuming this. Otherwise the reference to the
			// view cannot be available in getNextBuffer().
			this.subpartitionView = partitionProvider.createSubpartitionView(
				resultPartitionId,
				subPartitionIndex,
				this);
		} else {
			throw new IllegalStateException("Subpartition already requested");
		}
	}
}
 

@Test
public void testProducerFailedException() throws Exception {
	PartitionRequestQueue queue = new PartitionRequestQueue();

	ResultSubpartitionView view = new ReleasedResultSubpartitionView();

	ResultPartitionProvider partitionProvider =
		(partitionId, index, availabilityListener) -> view;

	EmbeddedChannel ch = new EmbeddedChannel(queue);

	CreditBasedSequenceNumberingViewReader seqView = new CreditBasedSequenceNumberingViewReader(new InputChannelID(), 2, queue);
	seqView.requestSubpartitionView(partitionProvider, new ResultPartitionID(), 0);
	// Add available buffer to trigger enqueue the erroneous view
	seqView.notifyDataAvailable();

	ch.runPendingTasks();

	// Read the enqueued msg
	Object msg = ch.readOutbound();

	assertEquals(msg.getClass(), NettyMessage.ErrorResponse.class);

	NettyMessage.ErrorResponse err = (NettyMessage.ErrorResponse) msg;
	assertTrue(err.cause instanceof CancelTaskException);
}
 

@Test
public void testProducerFailedException() throws Exception {
	PartitionRequestQueue queue = new PartitionRequestQueue();

	ResultSubpartitionView view = new ReleasedResultSubpartitionView();

	ResultPartitionProvider partitionProvider =
		(partitionId, index, availabilityListener) -> view;

	EmbeddedChannel ch = new EmbeddedChannel(queue);

	CreditBasedSequenceNumberingViewReader seqView = new CreditBasedSequenceNumberingViewReader(new InputChannelID(), 2, queue);
	seqView.requestSubpartitionView(partitionProvider, new ResultPartitionID(), 0);
	// Add available buffer to trigger enqueue the erroneous view
	seqView.notifyDataAvailable();

	ch.runPendingTasks();

	// Read the enqueued msg
	Object msg = ch.readOutbound();

	assertEquals(msg.getClass(), NettyMessage.ErrorResponse.class);

	NettyMessage.ErrorResponse err = (NettyMessage.ErrorResponse) msg;
	assertTrue(err.cause instanceof CancelTaskException);
}
 

@Test
public void testTaskExecutorShutdownHandling() throws Exception {
	final CompletableFuture<Collection<IntermediateDataSetID>> clusterPartitionReleaseFuture = new CompletableFuture<>();
	runTest(
		builder -> builder.setReleaseClusterPartitionsConsumer(clusterPartitionReleaseFuture::complete),
		(resourceManagerGateway, taskManagerId1, taskManagerId2) -> {
			IntermediateDataSetID dataSetID = new IntermediateDataSetID();

			resourceManagerGateway.heartbeatFromTaskManager(
				taskManagerId1,
				createTaskExecutorHeartbeatPayload(dataSetID, 2, new ResultPartitionID()));

			// we need a partition on another task executor so that there's something to release when one task executor goes down
			resourceManagerGateway.heartbeatFromTaskManager(
				taskManagerId2,
				createTaskExecutorHeartbeatPayload(dataSetID, 2, new ResultPartitionID()));

			resourceManagerGateway.disconnectTaskManager(taskManagerId2, new RuntimeException("test exception"));
			Collection<IntermediateDataSetID> intermediateDataSetIDS = clusterPartitionReleaseFuture.get(TIMEOUT.toMilliseconds(), TimeUnit.MILLISECONDS);
			assertThat(intermediateDataSetIDS, contains(dataSetID));
		});
}
 

@Test
public void testInputConstraintANY() throws Exception {
	List<JobVertex> vertices = createOrderedVertices();
	ExecutionGraph eg = createExecutionGraph(vertices, InputDependencyConstraint.ANY);
	ExecutionVertex ev11 = eg.getJobVertex(vertices.get(0).getID()).getTaskVertices()[0];
	ExecutionVertex ev21 = eg.getJobVertex(vertices.get(1).getID()).getTaskVertices()[0];
	ExecutionVertex ev22 = eg.getJobVertex(vertices.get(1).getID()).getTaskVertices()[1];
	ExecutionVertex ev31 = eg.getJobVertex(vertices.get(2).getID()).getTaskVertices()[0];

	eg.start(mainThreadExecutor);
	eg.scheduleForExecution();


	// Inputs constraint not satisfied on init
	assertFalse(ev31.checkInputDependencyConstraints());

	// Input1 consumable satisfies the constraint
	IntermediateResultPartition partition11 = ev11.getProducedPartitions().values().iterator().next();
	ev11.scheduleOrUpdateConsumers(new ResultPartitionID(partition11.getPartitionId(),
		ev11.getCurrentExecutionAttempt().getAttemptId()));
	assertTrue(ev31.checkInputDependencyConstraints());

	// Inputs constraint not satisfied after failover
	ev11.fail(new Exception());


	waitUntilJobRestarted(eg);


	assertFalse(ev31.checkInputDependencyConstraints());

	// Input2 consumable satisfies the constraint
	waitUntilExecutionVertexState(ev21, ExecutionState.DEPLOYING, 2000L);
	waitUntilExecutionVertexState(ev22, ExecutionState.DEPLOYING, 2000L);
	ev21.getCurrentExecutionAttempt().markFinished();
	ev22.getCurrentExecutionAttempt().markFinished();
	assertTrue(ev31.checkInputDependencyConstraints());

}
 

@Override
public CompletableFuture<ExecutionState> requestPartitionState(
		final IntermediateDataSetID intermediateResultId,
		final ResultPartitionID resultPartitionId) {

	try {
		return CompletableFuture.completedFuture(schedulerNG.requestPartitionState(intermediateResultId, resultPartitionId));
	} catch (PartitionProducerDisposedException e) {
		log.info("Error while requesting partition state", e);
		return FutureUtils.completedExceptionally(e);
	}
}
 

private void scheduleResultPartitionCleanup(JobID jobId) {
	final Collection<CompletableFuture<ExecutionState>> taskTerminationFutures = taskResultPartitionCleanupFuturesPerJob.remove(jobId);
	if (taskTerminationFutures != null) {
		FutureUtils.waitForAll(taskTerminationFutures)
			.thenRunAsync(() -> {
				Collection<ResultPartitionID> partitionsForJob = partitionTable.stopTrackingPartitions(jobId);
				shuffleEnvironment.releasePartitionsLocally(partitionsForJob);
			}, getMainThreadExecutor());
	}
}
 

private void testPartitionTrackingForStateTransition(final Consumer<Execution> stateTransition, final PartitionReleaseResult partitionReleaseResult) throws Exception {
	CompletableFuture<Tuple2<ResourceID, ResultPartitionDeploymentDescriptor>> partitionStartTrackingFuture = new CompletableFuture<>();
	CompletableFuture<Collection<ResultPartitionID>> partitionStopTrackingFuture = new CompletableFuture<>();
	CompletableFuture<Collection<ResultPartitionID>> partitionStopTrackingAndReleaseFuture = new CompletableFuture<>();
	final TestingPartitionTracker partitionTracker = new TestingPartitionTracker();
	partitionTracker.setStartTrackingPartitionsConsumer(
		(resourceID, resultPartitionDeploymentDescriptor) ->
			partitionStartTrackingFuture.complete(Tuple2.of(resourceID, resultPartitionDeploymentDescriptor))
	);
	partitionTracker.setStopTrackingPartitionsConsumer(partitionStopTrackingFuture::complete);
	partitionTracker.setStopTrackingAndReleasePartitionsConsumer(partitionStopTrackingAndReleaseFuture::complete);

	setupExecutionGraphAndStartRunningJob(ResultPartitionType.BLOCKING, partitionTracker, new SimpleAckingTaskManagerGateway(), NettyShuffleMaster.INSTANCE);

	Tuple2<ResourceID, ResultPartitionDeploymentDescriptor> startTrackingCall = partitionStartTrackingFuture.get();
	assertThat(startTrackingCall.f0, equalTo(taskExecutorResourceId));
	assertThat(startTrackingCall.f1, equalTo(descriptor));

	stateTransition.accept(execution);

	switch (partitionReleaseResult) {
		case NONE:
			assertFalse(partitionStopTrackingFuture.isDone());
			assertFalse(partitionStopTrackingAndReleaseFuture.isDone());
			break;
		case STOP_TRACKING:
			assertTrue(partitionStopTrackingFuture.isDone());
			assertFalse(partitionStopTrackingAndReleaseFuture.isDone());
			final Collection<ResultPartitionID> stopTrackingCall = partitionStopTrackingFuture.get();
			assertEquals(Collections.singletonList(descriptor.getShuffleDescriptor().getResultPartitionID()), stopTrackingCall);
			break;
		case STOP_TRACKING_AND_RELEASE:
			assertFalse(partitionStopTrackingFuture.isDone());
			assertTrue(partitionStopTrackingAndReleaseFuture.isDone());
			final Collection<ResultPartitionID> stopTrackingAndReleaseCall = partitionStopTrackingAndReleaseFuture.get();
			assertEquals(Collections.singletonList(descriptor.getShuffleDescriptor().getResultPartitionID()), stopTrackingAndReleaseCall);
			break;
	}
}
 

/**
 * Tests {@link TaskEventDispatcher#publish(ResultPartitionID, TaskEvent)} and {@link TaskEventDispatcher#subscribeToEvent(ResultPartitionID, EventListener, Class)} methods.
 */
@Test
public void publishSubscribe() throws Exception {
	ResultPartitionID partitionId1 = new ResultPartitionID();
	ResultPartitionID partitionId2 = new ResultPartitionID();
	TaskEventDispatcher ted = new TaskEventDispatcher();

	AllWorkersDoneEvent event1 = new AllWorkersDoneEvent();
	TerminationEvent event2 = new TerminationEvent();
	assertFalse(ted.publish(partitionId1, event1));

	ted.registerPartition(partitionId1);
	ted.registerPartition(partitionId2);

	// no event listener subscribed yet, but the event is forwarded to a TaskEventHandler
	assertTrue(ted.publish(partitionId1, event1));

	OneShotEventListener eventListener1a = new OneShotEventListener(event1);
	ZeroShotEventListener eventListener1b = new ZeroShotEventListener();
	ZeroShotEventListener eventListener2 = new ZeroShotEventListener();
	OneShotEventListener eventListener3 = new OneShotEventListener(event2);
	ted.subscribeToEvent(partitionId1, eventListener1a, AllWorkersDoneEvent.class);
	ted.subscribeToEvent(partitionId2, eventListener1b, AllWorkersDoneEvent.class);
	ted.subscribeToEvent(partitionId1, eventListener2, TaskEvent.class);
	ted.subscribeToEvent(partitionId1, eventListener3, TerminationEvent.class);

	assertTrue(ted.publish(partitionId1, event1));
	assertTrue("listener should have fired for AllWorkersDoneEvent", eventListener1a.fired);
	assertFalse("listener should not have fired for AllWorkersDoneEvent", eventListener3.fired);

	// publish another event, verify that only the right subscriber is called
	assertTrue(ted.publish(partitionId1, event2));
	assertTrue("listener should have fired for TerminationEvent", eventListener3.fired);
}
 

/**
 * Verifies that {@link RemoteInputChannel#onFailedPartitionRequest()} is called when a
 * {@link PartitionNotFoundException} is received.
 */
@Test
public void testReceivePartitionNotFoundException() throws Exception {
	// Minimal mock of a remote input channel
	final BufferProvider bufferProvider = mock(BufferProvider.class);
	when(bufferProvider.requestBuffer()).thenReturn(TestBufferFactory.createBuffer(0));

	final RemoteInputChannel inputChannel = mock(RemoteInputChannel.class);
	when(inputChannel.getInputChannelId()).thenReturn(new InputChannelID());
	when(inputChannel.getBufferProvider()).thenReturn(bufferProvider);

	final ErrorResponse partitionNotFound = new ErrorResponse(
		new PartitionNotFoundException(new ResultPartitionID()),
		inputChannel.getInputChannelId());

	final PartitionRequestClientHandler client = new PartitionRequestClientHandler();
	client.addInputChannel(inputChannel);

	// Mock channel context
	ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
	when(ctx.channel()).thenReturn(mock(Channel.class));

	client.channelActive(ctx);

	client.channelRead(ctx, partitionNotFound);

	verify(inputChannel, times(1)).onFailedPartitionRequest();
}
 

/**
 * Stops the tracking of all partition for the given key.
 */
public Collection<ResultPartitionID> stopTrackingPartitions(K key) {
	Preconditions.checkNotNull(key);

	Set<ResultPartitionID> storedPartitions = trackedPartitionsPerKey.remove(key);
	return storedPartitions == null
		? Collections.emptyList()
		: storedPartitions;
}
 

LocalRecoveredInputChannel(
		SingleInputGate inputGate,
		int channelIndex,
		ResultPartitionID partitionId,
		ResultPartitionManager partitionManager,
		TaskEventPublisher taskEventPublisher,
		int initialBackOff,
		int maxBackoff,
		InputChannelMetrics metrics) {
	super(inputGate, channelIndex, partitionId, initialBackOff, maxBackoff, metrics.getNumBytesInLocalCounter(), metrics.getNumBuffersInLocalCounter());

	this.partitionManager = checkNotNull(partitionManager);
	this.taskEventPublisher = checkNotNull(taskEventPublisher);
}
 

public RemoteInputChannel(
	SingleInputGate inputGate,
	int channelIndex,
	ResultPartitionID partitionId,
	ConnectionID connectionId,
	ConnectionManager connectionManager,
	TaskIOMetricGroup metrics) {

	this(inputGate, channelIndex, partitionId, connectionId, connectionManager, 0, 0, metrics);
}
 

@Test
public void unregisterPartition() throws Exception {
	ResultPartitionID partitionId1 = new ResultPartitionID();
	ResultPartitionID partitionId2 = new ResultPartitionID();
	TaskEventDispatcher ted = new TaskEventDispatcher();

	AllWorkersDoneEvent event = new AllWorkersDoneEvent();
	assertFalse(ted.publish(partitionId1, event));

	ted.registerPartition(partitionId1);
	ted.registerPartition(partitionId2);

	OneShotEventListener eventListener1a = new OneShotEventListener(event);
	ZeroShotEventListener eventListener1b = new ZeroShotEventListener();
	OneShotEventListener eventListener2 = new OneShotEventListener(event);
	ted.subscribeToEvent(partitionId1, eventListener1a, AllWorkersDoneEvent.class);
	ted.subscribeToEvent(partitionId2, eventListener1b, AllWorkersDoneEvent.class);
	ted.subscribeToEvent(partitionId1, eventListener2, AllWorkersDoneEvent.class);

	ted.unregisterPartition(partitionId2);

	// publish something for partitionId1 triggering all according listeners
	assertTrue(ted.publish(partitionId1, event));
	assertTrue("listener should have fired for AllWorkersDoneEvent", eventListener1a.fired);
	assertTrue("listener should have fired for AllWorkersDoneEvent", eventListener2.fired);

	// now publish something for partitionId2 which should not trigger any listeners
	assertFalse(ted.publish(partitionId2, event));
}
 

/**
 * Registers the given partition for incoming task events allowing calls to {@link
 * #subscribeToEvent(ResultPartitionID, EventListener, Class)}.
 *
 * @param partitionId
 * 		the partition ID
 */
public void registerPartition(ResultPartitionID partitionId) {
	checkNotNull(partitionId);

	synchronized (registeredHandlers) {
		LOG.debug("registering {}", partitionId);
		if (registeredHandlers.put(partitionId, new TaskEventHandler()) != null) {
			throw new IllegalStateException(
				"Partition " + partitionId + " already registered at task event dispatcher.");
		}
	}
}