You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by tr...@apache.org on 2018/05/24 14:49:45 UTC

[1/4] flink git commit: [FLINK-9416] Make all RestClusterClient calls retriable

Repository: flink
Updated Branches:
  refs/heads/master a5966fd6a -> 88987ebb9


[FLINK-9416] Make all RestClusterClient calls retriable

This commit changes the RestClusterClient calls such that they are all retriable wrt
to connection errors and if the service is currently unavailable (return code 503).

Moreover, it changes the retry behaviour for polling the JobResult such that it fails
now if the cluster returns a NOT_FOUND code.

This closes #6069.


Project: http://git-wip-us.apache.org/repos/asf/flink/repo
Commit: http://git-wip-us.apache.org/repos/asf/flink/commit/88987ebb
Tree: http://git-wip-us.apache.org/repos/asf/flink/tree/88987ebb
Diff: http://git-wip-us.apache.org/repos/asf/flink/diff/88987ebb

Branch: refs/heads/master
Commit: 88987ebb93b4d56cc5a2efd7ef47dedab9b9d3d9
Parents: 5ad0ca2
Author: Till Rohrmann <tr...@apache.org>
Authored: Thu May 24 10:01:23 2018 +0200
Committer: Till Rohrmann <tr...@apache.org>
Committed: Thu May 24 16:45:15 2018 +0200

----------------------------------------------------------------------
 .../client/program/rest/RestClusterClient.java  |  89 +++++++---------
 .../program/rest/RestClusterClientTest.java     | 101 ++++++++++++++++++-
 2 files changed, 139 insertions(+), 51 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/flink/blob/88987ebb/flink-clients/src/main/java/org/apache/flink/client/program/rest/RestClusterClient.java
----------------------------------------------------------------------
diff --git a/flink-clients/src/main/java/org/apache/flink/client/program/rest/RestClusterClient.java b/flink-clients/src/main/java/org/apache/flink/client/program/rest/RestClusterClient.java
index 0921642..7466e32 100644
--- a/flink-clients/src/main/java/org/apache/flink/client/program/rest/RestClusterClient.java
+++ b/flink-clients/src/main/java/org/apache/flink/client/program/rest/RestClusterClient.java
@@ -100,6 +100,7 @@ import org.apache.flink.util.Preconditions;
 import org.apache.flink.util.function.CheckedSupplier;
 
 import org.apache.flink.shaded.netty4.io.netty.channel.ConnectTimeoutException;
+import org.apache.flink.shaded.netty4.io.netty.handler.codec.http.HttpResponseStatus;
 
 import akka.actor.AddressFromURIString;
 
@@ -274,7 +275,9 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 		final JobMessageParameters  params = new JobMessageParameters();
 		params.jobPathParameter.resolve(jobId);
 
-		CompletableFuture<JobDetailsInfo> responseFuture = sendRequest(detailsHeaders, params);
+		CompletableFuture<JobDetailsInfo> responseFuture = sendRequest(
+			detailsHeaders,
+			params);
 
 		return responseFuture.thenApply(JobDetailsInfo::getJobStatus);
 	}
@@ -293,11 +296,9 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 			() -> {
 				final JobMessageParameters messageParameters = new JobMessageParameters();
 				messageParameters.jobPathParameter.resolve(jobId);
-				return sendRetryableRequest(
+				return sendRequest(
 					JobExecutionResultHeaders.getInstance(),
-					messageParameters,
-					EmptyRequestBody.getInstance(),
-					isConnectionProblemException().or(isHttpStatusUnsuccessfulException()));
+					messageParameters);
 			});
 	}
 
@@ -313,8 +314,7 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 		jobGraph.setAllowQueuedScheduling(true);
 
 		log.info("Requesting blob server port.");
-		CompletableFuture<BlobServerPortResponseBody> portFuture = sendRequest(
-			BlobServerPortHeaders.getInstance());
+		CompletableFuture<BlobServerPortResponseBody> portFuture = sendRequest(BlobServerPortHeaders.getInstance());
 
 		CompletableFuture<JobGraph> jobUploadFuture = portFuture.thenCombine(
 			getDispatcherAddress(),
@@ -366,8 +366,7 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 		params.terminationModeQueryParameter.resolve(Collections.singletonList(TerminationModeQueryParameter.TerminationMode.STOP));
 		CompletableFuture<EmptyResponseBody> responseFuture = sendRequest(
 			JobTerminationHeaders.getInstance(),
-			params
-		);
+			params);
 		responseFuture.get(timeout.toMillis(), TimeUnit.MILLISECONDS);
 	}
 
@@ -378,8 +377,7 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 		params.terminationModeQueryParameter.resolve(Collections.singletonList(TerminationModeQueryParameter.TerminationMode.CANCEL));
 		CompletableFuture<EmptyResponseBody> responseFuture = sendRequest(
 			JobTerminationHeaders.getInstance(),
-			params
-		);
+			params);
 		responseFuture.get(timeout.toMillis(), TimeUnit.MILLISECONDS);
 	}
 
@@ -391,7 +389,7 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 	@Override
 	public CompletableFuture<String> triggerSavepoint(
 			final JobID jobId,
-			final @Nullable String savepointDirectory) throws FlinkException {
+			final @Nullable String savepointDirectory) {
 		return triggerSavepoint(jobId, savepointDirectory, false);
 	}
 
@@ -429,8 +427,7 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 
 		CompletableFuture<JobAccumulatorsInfo> responseFuture = sendRequest(
 			accumulatorsHeaders,
-			accMsgParams
-		);
+			accMsgParams);
 
 		Map<String, OptionalFailure<Object>> result = Collections.emptyMap();
 
@@ -463,16 +460,14 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 				savepointStatusHeaders.getUnresolvedMessageParameters();
 			savepointStatusMessageParameters.jobIdPathParameter.resolve(jobId);
 			savepointStatusMessageParameters.triggerIdPathParameter.resolve(triggerID);
-			return sendRetryableRequest(
+			return sendRequest(
 				savepointStatusHeaders,
-				savepointStatusMessageParameters,
-				EmptyRequestBody.getInstance(),
-				isConnectionProblemException());
+				savepointStatusMessageParameters);
 		});
 	}
 
 	@Override
-	public CompletableFuture<Collection<JobStatusMessage>> listJobs() throws Exception {
+	public CompletableFuture<Collection<JobStatusMessage>> listJobs() {
 		return sendRequest(JobsOverviewHeaders.getInstance())
 			.thenApply(
 				(multipleJobsDetails) -> multipleJobsDetails
@@ -508,8 +503,7 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 
 		final CompletableFuture<TriggerResponse> rescalingTriggerResponseFuture = sendRequest(
 			rescalingTriggerHeaders,
-			rescalingTriggerMessageParameters,
-			EmptyRequestBody.getInstance());
+			rescalingTriggerMessageParameters);
 
 		final CompletableFuture<AsynchronousOperationInfo> rescalingOperationFuture = rescalingTriggerResponseFuture.thenCompose(
 			(TriggerResponse triggerResponse) -> {
@@ -521,11 +515,9 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 				rescalingStatusMessageParameters.triggerIdPathParameter.resolve(triggerId);
 
 				return pollResourceAsync(
-					() -> sendRetryableRequest(
+					() -> sendRequest(
 						rescalingStatusHeaders,
-						rescalingStatusMessageParameters,
-						EmptyRequestBody.getInstance(),
-						isConnectionProblemException()));
+						rescalingStatusMessageParameters));
 			});
 
 		return rescalingOperationFuture.thenApply(
@@ -544,7 +536,6 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 
 		final CompletableFuture<TriggerResponse> savepointDisposalTriggerFuture = sendRequest(
 			SavepointDisposalTriggerHeaders.getInstance(),
-			EmptyMessageParameters.getInstance(),
 			savepointDisposalRequest);
 
 		final CompletableFuture<AsynchronousOperationInfo> savepointDisposalFuture = savepointDisposalTriggerFuture.thenCompose(
@@ -555,11 +546,9 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 				savepointDisposalStatusMessageParameters.triggerIdPathParameter.resolve(triggerId);
 
 				return pollResourceAsync(
-					() -> sendRetryableRequest(
+					() -> sendRequest(
 						savepointDisposalStatusHeaders,
-						savepointDisposalStatusMessageParameters,
-						EmptyRequestBody.getInstance(),
-						isConnectionProblemException()));
+						savepointDisposalStatusMessageParameters));
 			});
 
 		return savepointDisposalFuture.thenApply(
@@ -575,11 +564,7 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 	@Override
 	public void shutDownCluster() {
 		try {
-			sendRetryableRequest(
-				ShutdownHeaders.getInstance(),
-				EmptyMessageParameters.getInstance(),
-				EmptyRequestBody.getInstance(),
-				isConnectionProblemException()).get();
+			sendRequest(ShutdownHeaders.getInstance()).get();
 		} catch (InterruptedException e) {
 			Thread.currentThread().interrupt();
 		} catch (ExecutionException e) {
@@ -681,24 +666,20 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 		return sendRequest(messageHeaders, EmptyMessageParameters.getInstance(), request);
 	}
 
-	private <M extends MessageHeaders<EmptyRequestBody, P, EmptyMessageParameters>, P extends ResponseBody> CompletableFuture<P>
+	@VisibleForTesting
+	<M extends MessageHeaders<EmptyRequestBody, P, EmptyMessageParameters>, P extends ResponseBody> CompletableFuture<P>
 			sendRequest(M messageHeaders) {
 		return sendRequest(messageHeaders, EmptyMessageParameters.getInstance(), EmptyRequestBody.getInstance());
 	}
 
 	private <M extends MessageHeaders<R, P, U>, U extends MessageParameters, R extends RequestBody, P extends ResponseBody> CompletableFuture<P>
 			sendRequest(M messageHeaders, U messageParameters, R request) {
-		return getWebMonitorBaseUrl().thenCompose(webMonitorBaseUrl -> {
-			try {
-				return restClient.sendRequest(webMonitorBaseUrl.getHost(), webMonitorBaseUrl.getPort(), messageHeaders, messageParameters, request);
-			} catch (IOException e) {
-				throw new CompletionException(e);
-			}
-		});
+		return sendRetriableRequest(
+			messageHeaders, messageParameters, request, isConnectionProblemOrServiceUnavailable());
 	}
 
 	private <M extends MessageHeaders<R, P, U>, U extends MessageParameters, R extends RequestBody, P extends ResponseBody> CompletableFuture<P>
-			sendRetryableRequest(M messageHeaders, U messageParameters, R request, Predicate<Throwable> retryPredicate) {
+			sendRetriableRequest(M messageHeaders, U messageParameters, R request, Predicate<Throwable> retryPredicate) {
 		return retry(() -> getWebMonitorBaseUrl().thenCompose(webMonitorBaseUrl -> {
 			try {
 				return restClient.sendRequest(webMonitorBaseUrl.getHost(), webMonitorBaseUrl.getPort(), messageHeaders, messageParameters, request);
@@ -719,6 +700,10 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 			new ScheduledExecutorServiceAdapter(retryExecutorService));
 	}
 
+	private static Predicate<Throwable> isConnectionProblemOrServiceUnavailable() {
+		return isConnectionProblemException().or(isServiceUnavailable());
+	}
+
 	private static Predicate<Throwable> isConnectionProblemException() {
 		return (throwable) ->
 			ExceptionUtils.findThrowable(throwable, java.net.ConnectException.class).isPresent() ||
@@ -727,13 +712,17 @@ public class RestClusterClient<T> extends ClusterClient<T> implements NewCluster
 				ExceptionUtils.findThrowable(throwable, IOException.class).isPresent();
 	}
 
-	private static Predicate<Throwable> isHttpStatusUnsuccessfulException() {
+	private static Predicate<Throwable> isServiceUnavailable() {
+		return httpExceptionCodePredicate(code -> code == HttpResponseStatus.SERVICE_UNAVAILABLE.code());
+	}
+
+	private static Predicate<Throwable> httpExceptionCodePredicate(Predicate<Integer> statusCodePredicate) {
 		return (throwable) -> ExceptionUtils.findThrowable(throwable, RestClientException.class)
-				.map(restClientException -> {
-					final int code = restClientException.getHttpResponseStatus().code();
-					return code < 200 || code > 299;
-				})
-				.orElse(false);
+			.map(restClientException -> {
+				final int code = restClientException.getHttpResponseStatus().code();
+				return statusCodePredicate.test(code);
+			})
+			.orElse(false);
 	}
 
 	@VisibleForTesting

http://git-wip-us.apache.org/repos/asf/flink/blob/88987ebb/flink-clients/src/test/java/org/apache/flink/client/program/rest/RestClusterClientTest.java
----------------------------------------------------------------------
diff --git a/flink-clients/src/test/java/org/apache/flink/client/program/rest/RestClusterClientTest.java b/flink-clients/src/test/java/org/apache/flink/client/program/rest/RestClusterClientTest.java
index df0ce23..f025d67 100644
--- a/flink-clients/src/test/java/org/apache/flink/client/program/rest/RestClusterClientTest.java
+++ b/flink-clients/src/test/java/org/apache/flink/client/program/rest/RestClusterClientTest.java
@@ -40,6 +40,7 @@ import org.apache.flink.runtime.jobmaster.JobResult;
 import org.apache.flink.runtime.messages.Acknowledge;
 import org.apache.flink.runtime.messages.webmonitor.JobDetails;
 import org.apache.flink.runtime.messages.webmonitor.MultipleJobsDetails;
+import org.apache.flink.runtime.rest.HttpMethodWrapper;
 import org.apache.flink.runtime.rest.RestClient;
 import org.apache.flink.runtime.rest.RestClientConfiguration;
 import org.apache.flink.runtime.rest.RestServerEndpoint;
@@ -369,7 +370,7 @@ public class RestClusterClientTest extends TestLogger {
 	public void testSubmitJobAndWaitForExecutionResult() throws Exception {
 		final TestJobExecutionResultHandler testJobExecutionResultHandler =
 			new TestJobExecutionResultHandler(
-				new RestHandlerException("should trigger retry", HttpResponseStatus.NOT_FOUND),
+				new RestHandlerException("should trigger retry", HttpResponseStatus.SERVICE_UNAVAILABLE),
 				JobExecutionResultResponseBody.inProgress(),
 				JobExecutionResultResponseBody.created(new JobResult.Builder()
 					.jobId(jobId)
@@ -709,6 +710,104 @@ public class RestClusterClientTest extends TestLogger {
 		assertThat(webMonitorBaseUrl.getPort(), equalTo(manualPort));
 	}
 
+	/**
+	 * Tests that the send operation is being retried.
+	 */
+	@Test
+	public void testRetriableSendOperationIfConnectionErrorOrServiceUnavailable() throws Exception {
+		final PingRestHandler pingRestHandler = new PingRestHandler(
+			FutureUtils.completedExceptionally(new RestHandlerException("test exception", HttpResponseStatus.SERVICE_UNAVAILABLE)),
+			CompletableFuture.completedFuture(EmptyResponseBody.getInstance()));
+
+		try (final TestRestServerEndpoint restServerEndpoint = createRestServerEndpoint(pingRestHandler)) {
+			final AtomicBoolean firstPollFailed = new AtomicBoolean();
+			failHttpRequest = (messageHeaders, messageParameters, requestBody) ->
+				messageHeaders instanceof PingRestHandlerHeaders && !firstPollFailed.getAndSet(true);
+
+			restClusterClient.sendRequest(PingRestHandlerHeaders.INSTANCE).get();
+		}
+	}
+
+	/**
+	 * Tests that the send operation is not being retried when receiving a NOT_FOUND return code.
+	 */
+	@Test
+	public void testSendIsNotRetriableIfHttpNotFound() throws Exception {
+		final String exceptionMessage = "test exception";
+		final PingRestHandler pingRestHandler = new PingRestHandler(
+			FutureUtils.completedExceptionally(new RestHandlerException(exceptionMessage, HttpResponseStatus.NOT_FOUND)));
+
+		try (final TestRestServerEndpoint restServerEndpoint = createRestServerEndpoint(pingRestHandler)) {
+			try {
+				restClusterClient.sendRequest(PingRestHandlerHeaders.INSTANCE).get();
+				fail("The rest request should have failed.");
+			}  catch (Exception e) {
+				assertThat(ExceptionUtils.findThrowableWithMessage(e, exceptionMessage).isPresent(), is(true));
+			}
+		}
+	}
+
+	private class PingRestHandler extends TestHandler<EmptyRequestBody, EmptyResponseBody, EmptyMessageParameters> {
+
+		private final Queue<CompletableFuture<EmptyResponseBody>> responseQueue;
+
+		private PingRestHandler(CompletableFuture<EmptyResponseBody> ... responses) {
+			super(PingRestHandlerHeaders.INSTANCE);
+			responseQueue = new ArrayDeque<>(Arrays.asList(responses));
+		}
+
+		@Override
+		protected CompletableFuture<EmptyResponseBody> handleRequest(@Nonnull HandlerRequest<EmptyRequestBody, EmptyMessageParameters> request, @Nonnull DispatcherGateway gateway) throws RestHandlerException {
+			final CompletableFuture<EmptyResponseBody> result = responseQueue.poll();
+
+			if (result != null) {
+				return result;
+			} else {
+				return CompletableFuture.completedFuture(EmptyResponseBody.getInstance());
+			}
+		}
+	}
+
+	private static final class PingRestHandlerHeaders implements MessageHeaders<EmptyRequestBody, EmptyResponseBody, EmptyMessageParameters> {
+
+		static final PingRestHandlerHeaders INSTANCE = new PingRestHandlerHeaders();
+
+		@Override
+		public Class<EmptyResponseBody> getResponseClass() {
+			return EmptyResponseBody.class;
+		}
+
+		@Override
+		public HttpResponseStatus getResponseStatusCode() {
+			return HttpResponseStatus.OK;
+		}
+
+		@Override
+		public String getDescription() {
+			return "foobar";
+		}
+
+		@Override
+		public Class<EmptyRequestBody> getRequestClass() {
+			return EmptyRequestBody.class;
+		}
+
+		@Override
+		public EmptyMessageParameters getUnresolvedMessageParameters() {
+			return EmptyMessageParameters.getInstance();
+		}
+
+		@Override
+		public HttpMethodWrapper getHttpMethod() {
+			return HttpMethodWrapper.GET;
+		}
+
+		@Override
+		public String getTargetRestEndpointURL() {
+			return "/foobar";
+		}
+	}
+
 	private class TestAccumulatorHandler extends TestHandler<EmptyRequestBody, JobAccumulatorsInfo, JobAccumulatorsMessageParameters> {
 
 		public TestAccumulatorHandler() {


[4/4] flink git commit: [FLINK-9427] Fix registration and request slot race condition in TaskExecutor

Posted by tr...@apache.org.
[FLINK-9427] Fix registration and request slot race condition in TaskExecutor

This commit fixes a race condition between the TaskExecutor and the ResourceManager. Before,
it could happen that the ResourceManager sends requestSlots message before the TaskExecutor
registration was completed. Due to this, the TaskExecutor did not have all information it needed
to accept task submissions.

The problem was that the TaskExecutor sent the SlotReport at registration time. Due to this, t
he SlotManager could already assign these slots to pending slot requests. With this commit, the
registration protocol changes such that the TaskExecutor first registers at the ResourceManager
and only after completing this step, it will announce the available slots to the SlotManager.

This closes #6067.


Project: http://git-wip-us.apache.org/repos/asf/flink/repo
Commit: http://git-wip-us.apache.org/repos/asf/flink/commit/47dc6997
Tree: http://git-wip-us.apache.org/repos/asf/flink/tree/47dc6997
Diff: http://git-wip-us.apache.org/repos/asf/flink/diff/47dc6997

Branch: refs/heads/master
Commit: 47dc69970996aa2a57c0dbfde866d3dac6d53001
Parents: a5966fd
Author: Till Rohrmann <tr...@apache.org>
Authored: Wed May 23 18:50:27 2018 +0200
Committer: Till Rohrmann <tr...@apache.org>
Committed: Thu May 24 16:45:15 2018 +0200

----------------------------------------------------------------------
 .../MesosResourceManagerTest.java               |   9 +-
 .../flink/runtime/jobmaster/JobMaster.java      |   3 +-
 .../RegistrationConnectionListener.java         |   5 +-
 .../resourcemanager/ResourceManager.java        |  18 +-
 .../resourcemanager/ResourceManagerGateway.java |  16 +-
 .../EstablishedResourceManagerConnection.java   |  62 +++
 .../runtime/taskexecutor/TaskExecutor.java      | 187 +++++----
 ...TaskExecutorToResourceManagerConnection.java |  56 +--
 .../clusterframework/ResourceManagerTest.java   |   3 -
 .../ResourceManagerTaskExecutorTest.java        |  14 +-
 .../resourcemanager/ResourceManagerTest.java    |   2 -
 .../slotmanager/SlotManagerTest.java            |  63 +++
 .../slotmanager/TestingResourceActions.java     |  44 +++
 .../utils/TestingResourceManagerGateway.java    |  50 ++-
 .../runtime/taskexecutor/TaskExecutorTest.java  | 385 ++++++++++++++-----
 .../TestingTaskExecutorGateway.java             |  10 +-
 .../TestingTaskExecutorGatewayBuilder.java      |  14 +-
 .../flink/yarn/YarnConfigurationITCase.java     |  16 +-
 .../flink/yarn/YarnResourceManagerTest.java     |  17 +-
 19 files changed, 730 insertions(+), 244 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-mesos/src/test/java/org/apache/flink/mesos/runtime/clusterframework/MesosResourceManagerTest.java
----------------------------------------------------------------------
diff --git a/flink-mesos/src/test/java/org/apache/flink/mesos/runtime/clusterframework/MesosResourceManagerTest.java b/flink-mesos/src/test/java/org/apache/flink/mesos/runtime/clusterframework/MesosResourceManagerTest.java
index 91647c0..56d2fbd 100644
--- a/flink-mesos/src/test/java/org/apache/flink/mesos/runtime/clusterframework/MesosResourceManagerTest.java
+++ b/flink-mesos/src/test/java/org/apache/flink/mesos/runtime/clusterframework/MesosResourceManagerTest.java
@@ -54,6 +54,7 @@ import org.apache.flink.runtime.jobmaster.JobMasterId;
 import org.apache.flink.runtime.jobmaster.JobMasterRegistrationSuccess;
 import org.apache.flink.runtime.leaderelection.TestingLeaderElectionService;
 import org.apache.flink.runtime.leaderretrieval.SettableLeaderRetrievalService;
+import org.apache.flink.runtime.messages.Acknowledge;
 import org.apache.flink.runtime.metrics.MetricRegistry;
 import org.apache.flink.runtime.metrics.MetricRegistryImpl;
 import org.apache.flink.runtime.registration.RegistrationResponse;
@@ -665,9 +666,15 @@ public class MesosResourceManagerTest extends TestLogger {
 			final HardwareDescription hardwareDescription = new HardwareDescription(1, 2L, 3L, 4L);
 			// send registration message
 			CompletableFuture<RegistrationResponse> successfulFuture =
-				resourceManager.registerTaskExecutor(task1Executor.address, task1Executor.resourceID, slotReport, dataPort, hardwareDescription, timeout);
+				resourceManager.registerTaskExecutor(task1Executor.address, task1Executor.resourceID, dataPort, hardwareDescription, timeout);
 			RegistrationResponse response = successfulFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);
 			assertTrue(response instanceof TaskExecutorRegistrationSuccess);
+			final TaskExecutorRegistrationSuccess registrationResponse = (TaskExecutorRegistrationSuccess) response;
+
+			final CompletableFuture<Acknowledge> initialSlotReportFuture = resourceManager.sendSlotReport(task1Executor.resourceID, registrationResponse.getRegistrationId(), slotReport, timeout);
+
+			// check for errors
+			initialSlotReportFuture.get();
 
 			// verify the internal state
 			assertThat(resourceManager.workersInLaunch, hasEntry(extractResourceID(task1), worker1launched));

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobMaster.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobMaster.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobMaster.java
index aad81a1..f96e0ae 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobMaster.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobMaster.java
@@ -1557,7 +1557,8 @@ public class JobMaster extends FencedRpcEndpoint<JobMasterId> implements JobMast
 		@Override
 		protected void onRegistrationSuccess(final JobMasterRegistrationSuccess success) {
 			runAsync(() -> {
-				// filter out replace connections
+				// filter out outdated connections
+				//noinspection ObjectEquality
 				if (this == resourceManagerConnection) {
 					establishResourceManagerConnection(success);
 				}

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/main/java/org/apache/flink/runtime/registration/RegistrationConnectionListener.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/registration/RegistrationConnectionListener.java b/flink-runtime/src/main/java/org/apache/flink/runtime/registration/RegistrationConnectionListener.java
index 360f982..3eb7243 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/registration/RegistrationConnectionListener.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/registration/RegistrationConnectionListener.java
@@ -22,14 +22,15 @@ package org.apache.flink.runtime.registration;
  * Classes which want to be notified about the registration result by the {@link RegisteredRpcConnection}
  * have to implement this interface.
  */
-public interface RegistrationConnectionListener<Success extends RegistrationResponse.Success> {
+public interface RegistrationConnectionListener<T extends RegisteredRpcConnection<?, ?, S>, S extends RegistrationResponse.Success> {
 
 	/**
 	 * This method is called by the {@link RegisteredRpcConnection} when the registration is success.
 	 *
 	 * @param success The concrete response information for successful registration.
+	 * @param connection The instance which established the connection
 	 */
-	void onRegistrationSuccess(Success success);
+	void onRegistrationSuccess(T connection, S success);
 
 	/**
 	 * This method is called by the {@link RegisteredRpcConnection} when the registration fails.

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/ResourceManager.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/ResourceManager.java b/flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/ResourceManager.java
index af736b2..6e5c824 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/ResourceManager.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/ResourceManager.java
@@ -338,7 +338,6 @@ public abstract class ResourceManager<WorkerType extends ResourceIDRetrievable>
 	public CompletableFuture<RegistrationResponse> registerTaskExecutor(
 			final String taskExecutorAddress,
 			final ResourceID taskExecutorResourceId,
-			final SlotReport slotReport,
 			final int dataPort,
 			final HardwareDescription hardwareDescription,
 			final Time timeout) {
@@ -354,7 +353,6 @@ public abstract class ResourceManager<WorkerType extends ResourceIDRetrievable>
 						taskExecutorGateway,
 						taskExecutorAddress,
 						taskExecutorResourceId,
-						slotReport,
 						dataPort,
 						hardwareDescription);
 				}
@@ -363,6 +361,18 @@ public abstract class ResourceManager<WorkerType extends ResourceIDRetrievable>
 	}
 
 	@Override
+	public CompletableFuture<Acknowledge> sendSlotReport(ResourceID taskManagerResourceId, InstanceID taskManagerRegistrationId, SlotReport slotReport, Time timeout) {
+		final WorkerRegistration<WorkerType> workerTypeWorkerRegistration = taskExecutors.get(taskManagerResourceId);
+
+		if (workerTypeWorkerRegistration.getInstanceID().equals(taskManagerRegistrationId)) {
+			slotManager.registerTaskManager(workerTypeWorkerRegistration, slotReport);
+			return CompletableFuture.completedFuture(Acknowledge.get());
+		} else {
+			return FutureUtils.completedExceptionally(new ResourceManagerException(String.format("Unknown TaskManager registration id %s.", taskManagerRegistrationId)));
+		}
+	}
+
+	@Override
 	public void heartbeatFromTaskManager(final ResourceID resourceID, final SlotReport slotReport) {
 		taskManagerHeartbeatManager.receiveHeartbeat(resourceID, slotReport);
 	}
@@ -669,7 +679,6 @@ public abstract class ResourceManager<WorkerType extends ResourceIDRetrievable>
 	 * @param taskExecutorGateway to communicate with the registering TaskExecutor
 	 * @param taskExecutorAddress address of the TaskExecutor
 	 * @param taskExecutorResourceId ResourceID of the TaskExecutor
-	 * @param slotReport initial slot report from the TaskExecutor
 	 * @param dataPort port used for data transfer
 	 * @param hardwareDescription of the registering TaskExecutor
 	 * @return RegistrationResponse
@@ -678,7 +687,6 @@ public abstract class ResourceManager<WorkerType extends ResourceIDRetrievable>
 			TaskExecutorGateway taskExecutorGateway,
 			String taskExecutorAddress,
 			ResourceID taskExecutorResourceId,
-			SlotReport slotReport,
 			int dataPort,
 			HardwareDescription hardwareDescription) {
 		WorkerRegistration<WorkerType> oldRegistration = taskExecutors.remove(taskExecutorResourceId);
@@ -702,8 +710,6 @@ public abstract class ResourceManager<WorkerType extends ResourceIDRetrievable>
 
 			taskExecutors.put(taskExecutorResourceId, registration);
 
-			slotManager.registerTaskManager(registration, slotReport);
-
 			taskManagerHeartbeatManager.monitorTarget(taskExecutorResourceId, new HeartbeatTarget<Void>() {
 				@Override
 				public void receiveHeartbeat(ResourceID resourceID, Void payload) {

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/ResourceManagerGateway.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/ResourceManagerGateway.java b/flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/ResourceManagerGateway.java
index bd282d6..ea79650 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/ResourceManagerGateway.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/ResourceManagerGateway.java
@@ -92,7 +92,6 @@ public interface ResourceManagerGateway extends FencedRpcGateway<ResourceManager
 	 *
 	 * @param taskExecutorAddress The address of the TaskExecutor that registers
 	 * @param resourceId The resource ID of the TaskExecutor that registers
-	 * @param slotReport The slot report containing free and allocated task slots
 	 * @param dataPort port used for data communication between TaskExecutors
 	 * @param hardwareDescription of the registering TaskExecutor
 	 * @param timeout The timeout for the response.
@@ -102,12 +101,25 @@ public interface ResourceManagerGateway extends FencedRpcGateway<ResourceManager
 	CompletableFuture<RegistrationResponse> registerTaskExecutor(
 		String taskExecutorAddress,
 		ResourceID resourceId,
-		SlotReport slotReport,
 		int dataPort,
 		HardwareDescription hardwareDescription,
 		@RpcTimeout Time timeout);
 
 	/**
+	 * Sends the given {@link SlotReport} to the ResourceManager.
+	 *
+	 * @param taskManagerRegistrationId id identifying the sending TaskManager
+	 * @param slotReport which is sent to the ResourceManager
+	 * @param timeout for the operation
+	 * @return Future which is completed with {@link Acknowledge} once the slot report has been received.
+	 */
+	CompletableFuture<Acknowledge> sendSlotReport(
+		ResourceID taskManagerResourceId,
+		InstanceID taskManagerRegistrationId,
+		SlotReport slotReport,
+		@RpcTimeout Time timeout);
+
+	/**
 	 * Sent by the TaskExecutor to notify the ResourceManager that a slot has become available.
 	 *
 	 * @param instanceId TaskExecutor's instance id

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/EstablishedResourceManagerConnection.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/EstablishedResourceManagerConnection.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/EstablishedResourceManagerConnection.java
new file mode 100644
index 0000000..5fb2c1c
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/EstablishedResourceManagerConnection.java
@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.taskexecutor;
+
+import org.apache.flink.runtime.clusterframework.types.ResourceID;
+import org.apache.flink.runtime.instance.InstanceID;
+import org.apache.flink.runtime.resourcemanager.ResourceManagerGateway;
+
+import javax.annotation.Nonnull;
+
+/**
+ * Container for the resource manager connection instances used by the
+ * {@link TaskExecutor}.
+ */
+class EstablishedResourceManagerConnection {
+
+	@Nonnull
+	private final ResourceManagerGateway resourceManagerGateway;
+
+	@Nonnull
+	private final ResourceID resourceManagerResourceId;
+
+	@Nonnull
+	private final InstanceID taskExecutorRegistrationId;
+
+	EstablishedResourceManagerConnection(@Nonnull ResourceManagerGateway resourceManagerGateway, @Nonnull ResourceID resourceManagerResourceId, @Nonnull InstanceID taskExecutorRegistrationId) {
+		this.resourceManagerGateway = resourceManagerGateway;
+		this.resourceManagerResourceId = resourceManagerResourceId;
+		this.taskExecutorRegistrationId = taskExecutorRegistrationId;
+	}
+
+	@Nonnull
+	public ResourceManagerGateway getResourceManagerGateway() {
+		return resourceManagerGateway;
+	}
+
+	@Nonnull
+	public ResourceID getResourceManagerResourceId() {
+		return resourceManagerResourceId;
+	}
+
+	@Nonnull
+	public InstanceID getTaskExecutorRegistrationId() {
+		return taskExecutorRegistrationId;
+	}
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutor.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutor.java
index e87c44d..dda2688 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutor.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutor.java
@@ -47,6 +47,7 @@ import org.apache.flink.runtime.heartbeat.HeartbeatServices;
 import org.apache.flink.runtime.heartbeat.HeartbeatTarget;
 import org.apache.flink.runtime.highavailability.HighAvailabilityServices;
 import org.apache.flink.runtime.instance.HardwareDescription;
+import org.apache.flink.runtime.instance.InstanceID;
 import org.apache.flink.runtime.io.network.NetworkEnvironment;
 import org.apache.flink.runtime.io.network.netty.PartitionProducerStateChecker;
 import org.apache.flink.runtime.io.network.partition.ResultPartitionConsumableNotifier;
@@ -56,6 +57,7 @@ import org.apache.flink.runtime.jobgraph.tasks.InputSplitProvider;
 import org.apache.flink.runtime.jobmaster.JMTMRegistrationSuccess;
 import org.apache.flink.runtime.jobmaster.JobMasterGateway;
 import org.apache.flink.runtime.jobmaster.JobMasterId;
+import org.apache.flink.runtime.jobmaster.ResourceManagerAddress;
 import org.apache.flink.runtime.leaderretrieval.LeaderRetrievalListener;
 import org.apache.flink.runtime.leaderretrieval.LeaderRetrievalService;
 import org.apache.flink.runtime.messages.Acknowledge;
@@ -82,6 +84,7 @@ import org.apache.flink.runtime.taskexecutor.exceptions.RegistrationTimeoutExcep
 import org.apache.flink.runtime.taskexecutor.exceptions.SlotAllocationException;
 import org.apache.flink.runtime.taskexecutor.exceptions.SlotOccupiedException;
 import org.apache.flink.runtime.taskexecutor.exceptions.TaskException;
+import org.apache.flink.runtime.taskexecutor.exceptions.TaskManagerException;
 import org.apache.flink.runtime.taskexecutor.exceptions.TaskSubmissionException;
 import org.apache.flink.runtime.taskexecutor.rpc.RpcCheckpointResponder;
 import org.apache.flink.runtime.taskexecutor.rpc.RpcInputSplitProvider;
@@ -167,10 +170,6 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway {
 	/** The network component in the task manager. */
 	private final NetworkEnvironment networkEnvironment;
 
-	// --------- resource manager --------
-
-	private TaskExecutorToResourceManagerConnection resourceManagerConnection;
-
 	// --------- job manager connections -----------
 
 	private final Map<ResourceID, JobManagerConnection> jobManagerConnections;
@@ -191,6 +190,17 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway {
 
 	private FileCache fileCache;
 
+	// --------- resource manager --------
+
+	@Nullable
+	private ResourceManagerAddress resourceManagerAddress;
+
+	@Nullable
+	private EstablishedResourceManagerConnection establishedResourceManagerConnection;
+
+	@Nullable
+	private TaskExecutorToResourceManagerConnection resourceManagerConnection;
+
 	@Nullable
 	private UUID currentRegistrationTimeoutId;
 
@@ -242,6 +252,8 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway {
 		this.hardwareDescription = HardwareDescription.extractFromSystem(
 			taskExecutorServices.getMemoryManager().getMemorySize());
 
+		this.resourceManagerAddress = null;
+		this.resourceManagerConnection = null;
 		this.currentRegistrationTimeoutId = null;
 	}
 
@@ -732,19 +744,10 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway {
 			allocationId, jobId, resourceManagerId);
 
 		try {
-			if (resourceManagerConnection == null) {
-				final String message = "TaskManager is not connected to a resource manager.";
+			if (!isConnectedToResourceManager(resourceManagerId)) {
+				final String message = String.format("TaskManager is not connected to the resource manager %s.", resourceManagerId);
 				log.debug(message);
-				throw new SlotAllocationException(message);
-			}
-
-			if (!Objects.equals(resourceManagerConnection.getTargetLeaderId(), resourceManagerId)) {
-				final String message = "The leader id " + resourceManagerId +
-					" does not match with the leader id of the connected resource manager " +
-					resourceManagerConnection.getTargetLeaderId() + '.';
-
-				log.debug(message);
-				throw new SlotAllocationException(message);
+				throw new TaskManagerException(message);
 			}
 
 			if (taskSlotTable.isSlotFree(slotId.getSlotNumber())) {
@@ -788,8 +791,8 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway {
 					throw new SlotAllocationException("Could not add job to job leader service.", e);
 				}
 			}
-		} catch (SlotAllocationException slotAllocationException) {
-			return FutureUtils.completedExceptionally(slotAllocationException);
+		} catch (TaskManagerException taskManagerException) {
+			return FutureUtils.completedExceptionally(taskManagerException);
 		}
 
 		return CompletableFuture.completedFuture(Acknowledge.get());
@@ -855,7 +858,7 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway {
 
 	@Override
 	public void disconnectResourceManager(Exception cause) {
-		closeResourceManagerConnection(cause);
+		reconnectToResourceManager(cause);
 	}
 
 	// ======================================================================
@@ -867,56 +870,76 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway {
 	// ------------------------------------------------------------------------
 
 	private void notifyOfNewResourceManagerLeader(String newLeaderAddress, ResourceManagerId newResourceManagerId) {
-		if (resourceManagerConnection != null) {
-			if (newLeaderAddress != null) {
-				// the resource manager switched to a new leader
-				log.info("ResourceManager leader changed from {} to {}. Registering at new leader.",
-					resourceManagerConnection.getTargetAddress(), newLeaderAddress);
-			}
-			else {
-				// address null means that the current leader is lost without a new leader being there, yet
-				log.info("Current ResourceManager {} lost leader status. Waiting for new ResourceManager leader.",
-					resourceManagerConnection.getTargetAddress());
-			}
+		resourceManagerAddress = createResourceManagerAddress(newLeaderAddress, newResourceManagerId);
+		reconnectToResourceManager(new FlinkException(String.format("ResourceManager leader changed to new address %s", resourceManagerAddress)));
+	}
 
-			// drop the current connection or connection attempt
-			closeResourceManagerConnection(
-				new FlinkException("New ResourceManager leader found under: " + newLeaderAddress +
-					'(' + newResourceManagerId + ')'));
+	@Nullable
+	private ResourceManagerAddress createResourceManagerAddress(@Nullable String newLeaderAddress, @Nullable ResourceManagerId newResourceManagerId) {
+		if (newLeaderAddress == null) {
+			return null;
+		} else {
+			assert(newResourceManagerId != null);
+			return new ResourceManagerAddress(newLeaderAddress, newResourceManagerId);
 		}
+	}
 
-		// establish a connection to the new leader
-		if (newLeaderAddress != null) {
-			createResourceManagerConnection(newLeaderAddress, newResourceManagerId);
+	private void reconnectToResourceManager(Exception cause) {
+		closeResourceManagerConnection(cause);
+		tryConnectToResourceManager();
+	}
+
+	private void tryConnectToResourceManager() {
+		if (resourceManagerAddress != null) {
+			connectToResourceManager();
 		}
 	}
 
-	private void createResourceManagerConnection(String newLeaderAddress, ResourceManagerId newResourceManagerId) {
-		log.info("Attempting to register at ResourceManager {} ({})", newLeaderAddress, newResourceManagerId);
+	private void connectToResourceManager() {
+		assert(resourceManagerAddress != null);
+		assert(establishedResourceManagerConnection == null);
+		assert(resourceManagerConnection == null);
+
+		log.info("Connecting to ResourceManager {}.", resourceManagerAddress);
+
 		resourceManagerConnection =
 			new TaskExecutorToResourceManagerConnection(
 				log,
 				getRpcService(),
 				getAddress(),
 				getResourceID(),
-				taskSlotTable.createSlotReport(getResourceID()),
 				taskManagerLocation.dataPort(),
 				hardwareDescription,
-				newLeaderAddress,
-				newResourceManagerId,
+				resourceManagerAddress.getAddress(),
+				resourceManagerAddress.getResourceManagerId(),
 				getMainThreadExecutor(),
 				new ResourceManagerRegistrationListener());
 		resourceManagerConnection.start();
 	}
 
 	private void establishResourceManagerConnection(
+			ResourceManagerGateway resourceManagerGateway,
 			ResourceID resourceManagerResourceId,
+			InstanceID taskExecutorRegistrationId,
 			ClusterInformation clusterInformation) {
+
+		final CompletableFuture<Acknowledge> slotReportResponseFuture = resourceManagerGateway.sendSlotReport(
+			getResourceID(),
+			taskExecutorRegistrationId,
+			taskSlotTable.createSlotReport(getResourceID()),
+			taskManagerConfiguration.getTimeout());
+
+		slotReportResponseFuture.whenCompleteAsync(
+			(acknowledge, throwable) -> {
+				if (throwable != null) {
+					reconnectToResourceManager(new TaskManagerException("Failed to send initial slot report to ResourceManager.", throwable));
+				}
+			}, getMainThreadExecutor());
+
 		// monitor the resource manager as heartbeat target
 		resourceManagerHeartbeatManager.monitorTarget(resourceManagerResourceId, new HeartbeatTarget<SlotReport>() {
 			@Override
 			public void receiveHeartbeat(ResourceID resourceID, SlotReport slotReport) {
-				ResourceManagerGateway resourceManagerGateway = resourceManagerConnection.getTargetGateway();
 				resourceManagerGateway.heartbeatFromTaskManager(resourceID, slotReport);
 			}
 
@@ -933,25 +956,35 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway {
 
 		blobCacheService.setBlobServerAddress(blobServerAddress);
 
+		establishedResourceManagerConnection = new EstablishedResourceManagerConnection(
+			resourceManagerGateway,
+			resourceManagerResourceId,
+			taskExecutorRegistrationId);
+
 		stopRegistrationTimeout();
 	}
 
 	private void closeResourceManagerConnection(Exception cause) {
-		if (resourceManagerConnection != null) {
-
-			if (resourceManagerConnection.isConnected()) {
-				if (log.isDebugEnabled()) {
-					log.debug("Close ResourceManager connection {}.",
-						resourceManagerConnection.getResourceManagerId(), cause);
-				} else {
-					log.info("Close ResourceManager connection {}.",
-						resourceManagerConnection.getResourceManagerId());
-				}
-				resourceManagerHeartbeatManager.unmonitorTarget(resourceManagerConnection.getResourceManagerId());
+		if (establishedResourceManagerConnection != null) {
+			final ResourceID resourceManagerResourceId = establishedResourceManagerConnection.getResourceManagerResourceId();
 
-				ResourceManagerGateway resourceManagerGateway = resourceManagerConnection.getTargetGateway();
-				resourceManagerGateway.disconnectTaskManager(getResourceID(), cause);
+			if (log.isDebugEnabled()) {
+				log.debug("Close ResourceManager connection {}.",
+					resourceManagerResourceId, cause);
 			} else {
+				log.info("Close ResourceManager connection {}.",
+					resourceManagerResourceId);
+			}
+			resourceManagerHeartbeatManager.unmonitorTarget(resourceManagerResourceId);
+
+			ResourceManagerGateway resourceManagerGateway = establishedResourceManagerConnection.getResourceManagerGateway();
+			resourceManagerGateway.disconnectTaskManager(getResourceID(), cause);
+
+			establishedResourceManagerConnection = null;
+		}
+
+		if (resourceManagerConnection != null) {
+			if (!resourceManagerConnection.isConnected()) {
 				if (log.isDebugEnabled()) {
 					log.debug("Terminating registration attempts towards ResourceManager {}.",
 						resourceManagerConnection.getTargetAddress(), cause);
@@ -1327,10 +1360,10 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway {
 
 				if (isConnectedToResourceManager()) {
 					// the slot was freed. Tell the RM about it
-					ResourceManagerGateway resourceManagerGateway = resourceManagerConnection.getTargetGateway();
+					ResourceManagerGateway resourceManagerGateway = establishedResourceManagerConnection.getResourceManagerGateway();
 
 					resourceManagerGateway.notifySlotAvailable(
-						resourceManagerConnection.getRegistrationId(),
+						establishedResourceManagerConnection.getTaskExecutorRegistrationId(),
 						new SlotID(getResourceID(), slotIndex),
 						allocationId);
 				}
@@ -1375,7 +1408,11 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway {
 	// ------------------------------------------------------------------------
 
 	private boolean isConnectedToResourceManager() {
-		return (resourceManagerConnection != null && resourceManagerConnection.isConnected());
+		return establishedResourceManagerConnection != null;
+	}
+
+	private boolean isConnectedToResourceManager(ResourceManagerId resourceManagerId) {
+		return establishedResourceManagerConnection != null && resourceManagerAddress != null && resourceManagerAddress.getResourceManagerId().equals(resourceManagerId);
 	}
 
 	private boolean isJobManagerConnectionValid(JobID jobId, JobMasterId jobMasterId) {
@@ -1478,16 +1515,27 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway {
 		}
 	}
 
-	private final class ResourceManagerRegistrationListener implements RegistrationConnectionListener<TaskExecutorRegistrationSuccess> {
+	private final class ResourceManagerRegistrationListener implements RegistrationConnectionListener<TaskExecutorToResourceManagerConnection, TaskExecutorRegistrationSuccess> {
 
 		@Override
-		public void onRegistrationSuccess(TaskExecutorRegistrationSuccess success) {
+		public void onRegistrationSuccess(TaskExecutorToResourceManagerConnection connection, TaskExecutorRegistrationSuccess success) {
 			final ResourceID resourceManagerId = success.getResourceManagerId();
+			final InstanceID taskExecutorRegistrationId = success.getRegistrationId();
 			final ClusterInformation clusterInformation = success.getClusterInformation();
+			final ResourceManagerGateway resourceManagerGateway = connection.getTargetGateway();
 
 			runAsync(
-				() -> establishResourceManagerConnection(resourceManagerId, clusterInformation)
-			);
+				() -> {
+					// filter out outdated connections
+					//noinspection ObjectEquality
+					if (resourceManagerConnection == connection) {
+						establishResourceManagerConnection(
+							resourceManagerGateway,
+							resourceManagerId,
+							taskExecutorRegistrationId,
+							clusterInformation);
+					}
+				});
 		}
 
 		@Override
@@ -1598,18 +1646,11 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway {
 		public void notifyHeartbeatTimeout(final ResourceID resourceId) {
 			runAsync(() -> {
 				// first check whether the timeout is still valid
-				if (resourceManagerConnection != null && resourceManagerConnection.getResourceManagerId().equals(resourceId)) {
+				if (establishedResourceManagerConnection != null && establishedResourceManagerConnection.getResourceManagerResourceId().equals(resourceId)) {
 					log.info("The heartbeat of ResourceManager with id {} timed out.", resourceId);
 
-					final String resourceManagerAddress = resourceManagerConnection.getTargetAddress();
-					final ResourceManagerId resourceManagerId = resourceManagerConnection.getTargetLeaderId();
-
-					closeResourceManagerConnection(
-						new TimeoutException(
-							"The heartbeat of ResourceManager with id " + resourceId + " timed out."));
-
-					// try to reconnect to old ResourceManager
-					createResourceManagerConnection(resourceManagerAddress, resourceManagerId);
+					reconnectToResourceManager(new TaskManagerException(
+						String.format("The heartbeat of ResourceManager with id %s timed out.", resourceId)));
 				} else {
 					log.debug("Received heartbeat timeout for outdated ResourceManager id {}. Ignoring the timeout.", resourceId);
 				}

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutorToResourceManagerConnection.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutorToResourceManagerConnection.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutorToResourceManagerConnection.java
index 1288009..9e1de52 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutorToResourceManagerConnection.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutorToResourceManagerConnection.java
@@ -21,16 +21,14 @@ package org.apache.flink.runtime.taskexecutor;
 import org.apache.flink.api.common.time.Time;
 import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.runtime.instance.HardwareDescription;
-import org.apache.flink.runtime.instance.InstanceID;
 import org.apache.flink.runtime.registration.RegisteredRpcConnection;
 import org.apache.flink.runtime.registration.RegistrationConnectionListener;
-import org.apache.flink.runtime.resourcemanager.ResourceManagerId;
-import org.apache.flink.runtime.rpc.RpcService;
 import org.apache.flink.runtime.registration.RegistrationResponse;
 import org.apache.flink.runtime.registration.RetryingRegistration;
 import org.apache.flink.runtime.resourcemanager.ResourceManagerGateway;
+import org.apache.flink.runtime.resourcemanager.ResourceManagerId;
+import org.apache.flink.runtime.rpc.RpcService;
 
-import org.apache.flink.util.Preconditions;
 import org.slf4j.Logger;
 
 import java.util.concurrent.CompletableFuture;
@@ -50,43 +48,34 @@ public class TaskExecutorToResourceManagerConnection
 
 	private final ResourceID taskManagerResourceId;
 
-	private final SlotReport slotReport;
-
 	private final int dataPort;
 
 	private final HardwareDescription hardwareDescription;
 
-	private final RegistrationConnectionListener<TaskExecutorRegistrationSuccess> registrationListener;
-
-	private InstanceID registrationId;
-
-	private ResourceID resourceManagerResourceId;
+	private final RegistrationConnectionListener<TaskExecutorToResourceManagerConnection, TaskExecutorRegistrationSuccess> registrationListener;
 
 	public TaskExecutorToResourceManagerConnection(
 			Logger log,
 			RpcService rpcService,
 			String taskManagerAddress,
 			ResourceID taskManagerResourceId,
-			SlotReport slotReport,
 			int dataPort,
 			HardwareDescription hardwareDescription,
 			String resourceManagerAddress,
 			ResourceManagerId resourceManagerId,
 			Executor executor,
-			RegistrationConnectionListener<TaskExecutorRegistrationSuccess> registrationListener) {
+			RegistrationConnectionListener<TaskExecutorToResourceManagerConnection, TaskExecutorRegistrationSuccess> registrationListener) {
 
 		super(log, resourceManagerAddress, resourceManagerId, executor);
 
-		this.rpcService = Preconditions.checkNotNull(rpcService);
-		this.taskManagerAddress = Preconditions.checkNotNull(taskManagerAddress);
-		this.taskManagerResourceId = Preconditions.checkNotNull(taskManagerResourceId);
-		this.slotReport = Preconditions.checkNotNull(slotReport);
+		this.rpcService = checkNotNull(rpcService);
+		this.taskManagerAddress = checkNotNull(taskManagerAddress);
+		this.taskManagerResourceId = checkNotNull(taskManagerResourceId);
 		this.dataPort = dataPort;
-		this.hardwareDescription = Preconditions.checkNotNull(hardwareDescription);
-		this.registrationListener = Preconditions.checkNotNull(registrationListener);
+		this.hardwareDescription = checkNotNull(hardwareDescription);
+		this.registrationListener = checkNotNull(registrationListener);
 	}
 
-
 	@Override
 	protected RetryingRegistration<ResourceManagerId, ResourceManagerGateway, TaskExecutorRegistrationSuccess> generateRegistration() {
 		return new TaskExecutorToResourceManagerConnection.ResourceManagerRegistration(
@@ -96,7 +85,6 @@ public class TaskExecutorToResourceManagerConnection
 			getTargetLeaderId(),
 			taskManagerAddress,
 			taskManagerResourceId,
-			slotReport,
 			dataPort,
 			hardwareDescription);
 	}
@@ -106,9 +94,7 @@ public class TaskExecutorToResourceManagerConnection
 		log.info("Successful registration at resource manager {} under registration id {}.",
 			getTargetAddress(), success.getRegistrationId());
 
-		registrationId = success.getRegistrationId();
-		resourceManagerResourceId = success.getResourceManagerId();
-		registrationListener.onRegistrationSuccess(success);
+		registrationListener.onRegistrationSuccess(this, success);
 	}
 
 	@Override
@@ -118,21 +104,6 @@ public class TaskExecutorToResourceManagerConnection
 		registrationListener.onRegistrationFailure(failure);
 	}
 
-	/**
-	 * Gets the ID under which the TaskExecutor is registered at the ResourceManager.
-	 * This returns null until the registration is completed.
-	 */
-	public InstanceID getRegistrationId() {
-		return registrationId;
-	}
-
-	/**
-	 * Gets the unique id of ResourceManager, that is returned when registration success.
-	 */
-	public ResourceID getResourceManagerId() {
-		return resourceManagerResourceId;
-	}
-
 	// ------------------------------------------------------------------------
 	//  Utilities
 	// ------------------------------------------------------------------------
@@ -141,10 +112,8 @@ public class TaskExecutorToResourceManagerConnection
 			extends RetryingRegistration<ResourceManagerId, ResourceManagerGateway, TaskExecutorRegistrationSuccess> {
 
 		private final String taskExecutorAddress;
-		
-		private final ResourceID resourceID;
 
-		private final SlotReport slotReport;
+		private final ResourceID resourceID;
 
 		private final int dataPort;
 
@@ -157,14 +126,12 @@ public class TaskExecutorToResourceManagerConnection
 				ResourceManagerId resourceManagerId,
 				String taskExecutorAddress,
 				ResourceID resourceID,
-				SlotReport slotReport,
 				int dataPort,
 				HardwareDescription hardwareDescription) {
 
 			super(log, rpcService, "ResourceManager", ResourceManagerGateway.class, targetAddress, resourceManagerId);
 			this.taskExecutorAddress = checkNotNull(taskExecutorAddress);
 			this.resourceID = checkNotNull(resourceID);
-			this.slotReport = checkNotNull(slotReport);
 			this.dataPort = dataPort;
 			this.hardwareDescription = checkNotNull(hardwareDescription);
 		}
@@ -177,7 +144,6 @@ public class TaskExecutorToResourceManagerConnection
 			return resourceManager.registerTaskExecutor(
 				taskExecutorAddress,
 				resourceID,
-				slotReport,
 				dataPort,
 				hardwareDescription,
 				timeout);

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/test/java/org/apache/flink/runtime/clusterframework/ResourceManagerTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/clusterframework/ResourceManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/clusterframework/ResourceManagerTest.java
index a1b3fb6..8898ef0 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/clusterframework/ResourceManagerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/clusterframework/ResourceManagerTest.java
@@ -55,7 +55,6 @@ import org.apache.flink.runtime.resourcemanager.StandaloneResourceManager;
 import org.apache.flink.runtime.resourcemanager.slotmanager.SlotManager;
 import org.apache.flink.runtime.rpc.RpcUtils;
 import org.apache.flink.runtime.rpc.TestingRpcService;
-import org.apache.flink.runtime.taskexecutor.SlotReport;
 import org.apache.flink.runtime.taskexecutor.TaskExecutorGateway;
 import org.apache.flink.runtime.taskexecutor.TaskExecutorRegistrationSuccess;
 import org.apache.flink.runtime.testingUtils.TestingUtils;
@@ -540,12 +539,10 @@ public class ResourceManagerTest extends TestLogger {
 			final UUID rmLeaderSessionId = UUID.randomUUID();
 			rmLeaderElectionService.isLeader(rmLeaderSessionId).get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);
 
-			final SlotReport slotReport = new SlotReport();
 			// test registration response successful and it will trigger monitor heartbeat target, schedule heartbeat request at interval time
 			CompletableFuture<RegistrationResponse> successfulFuture = rmGateway.registerTaskExecutor(
 				taskManagerAddress,
 				taskManagerResourceID,
-				slotReport,
 				dataPort,
 				hardwareDescription,
 				timeout);

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/ResourceManagerTaskExecutorTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/ResourceManagerTaskExecutorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/ResourceManagerTaskExecutorTest.java
index 58dc6ca..540db69 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/ResourceManagerTaskExecutorTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/ResourceManagerTaskExecutorTest.java
@@ -28,20 +28,20 @@ import org.apache.flink.runtime.instance.HardwareDescription;
 import org.apache.flink.runtime.leaderelection.LeaderElectionService;
 import org.apache.flink.runtime.leaderelection.TestingLeaderElectionService;
 import org.apache.flink.runtime.metrics.MetricRegistryImpl;
+import org.apache.flink.runtime.registration.RegistrationResponse;
 import org.apache.flink.runtime.resourcemanager.slotmanager.SlotManager;
 import org.apache.flink.runtime.rest.messages.taskmanager.TaskManagerInfo;
 import org.apache.flink.runtime.rpc.FatalErrorHandler;
 import org.apache.flink.runtime.rpc.RpcUtils;
 import org.apache.flink.runtime.rpc.TestingRpcService;
-import org.apache.flink.runtime.registration.RegistrationResponse;
 import org.apache.flink.runtime.rpc.exceptions.FencingTokenException;
-import org.apache.flink.runtime.taskexecutor.SlotReport;
 import org.apache.flink.runtime.taskexecutor.TaskExecutorGateway;
 import org.apache.flink.runtime.taskexecutor.TaskExecutorRegistrationSuccess;
 import org.apache.flink.runtime.testingUtils.TestingUtils;
 import org.apache.flink.runtime.util.TestingFatalErrorHandler;
 import org.apache.flink.util.ExceptionUtils;
 import org.apache.flink.util.TestLogger;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -65,8 +65,6 @@ public class ResourceManagerTaskExecutorTest extends TestLogger {
 
 	private TestingRpcService rpcService;
 
-	private SlotReport slotReport = new SlotReport();
-
 	private int dataPort = 1234;
 
 	private HardwareDescription hardwareDescription = new HardwareDescription(1, 2L, 3L, 4L);
@@ -115,7 +113,7 @@ public class ResourceManagerTaskExecutorTest extends TestLogger {
 		try {
 			// test response successful
 			CompletableFuture<RegistrationResponse> successfulFuture =
-				rmGateway.registerTaskExecutor(taskExecutorAddress, taskExecutorResourceID, slotReport, dataPort, hardwareDescription, timeout);
+				rmGateway.registerTaskExecutor(taskExecutorAddress, taskExecutorResourceID, dataPort, hardwareDescription, timeout);
 			RegistrationResponse response = successfulFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);
 			assertTrue(response instanceof TaskExecutorRegistrationSuccess);
 			final TaskManagerInfo taskManagerInfo = rmGateway.requestTaskManagerInfo(
@@ -125,7 +123,7 @@ public class ResourceManagerTaskExecutorTest extends TestLogger {
 
 			// test response successful with instanceID not equal to previous when receive duplicate registration from taskExecutor
 			CompletableFuture<RegistrationResponse> duplicateFuture =
-				rmGateway.registerTaskExecutor(taskExecutorAddress, taskExecutorResourceID, slotReport, dataPort, hardwareDescription, timeout);
+				rmGateway.registerTaskExecutor(taskExecutorAddress, taskExecutorResourceID, dataPort, hardwareDescription, timeout);
 			RegistrationResponse duplicateResponse = duplicateFuture.get();
 			assertTrue(duplicateResponse instanceof TaskExecutorRegistrationSuccess);
 			assertNotEquals(((TaskExecutorRegistrationSuccess) response).getRegistrationId(), ((TaskExecutorRegistrationSuccess) duplicateResponse).getRegistrationId());
@@ -144,7 +142,7 @@ public class ResourceManagerTaskExecutorTest extends TestLogger {
 		try {
 			// test throw exception when receive a registration from taskExecutor which takes unmatched leaderSessionId
 			CompletableFuture<RegistrationResponse> unMatchedLeaderFuture =
-				wronglyFencedGateway.registerTaskExecutor(taskExecutorAddress, taskExecutorResourceID, slotReport, dataPort, hardwareDescription, timeout);
+				wronglyFencedGateway.registerTaskExecutor(taskExecutorAddress, taskExecutorResourceID, dataPort, hardwareDescription, timeout);
 
 			try {
 				unMatchedLeaderFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);
@@ -168,7 +166,7 @@ public class ResourceManagerTaskExecutorTest extends TestLogger {
 			// test throw exception when receive a registration from taskExecutor which takes invalid address
 			String invalidAddress = "/taskExecutor2";
 			CompletableFuture<RegistrationResponse> invalidAddressFuture =
-				rmGateway.registerTaskExecutor(invalidAddress, taskExecutorResourceID, slotReport, dataPort, hardwareDescription, timeout);
+				rmGateway.registerTaskExecutor(invalidAddress, taskExecutorResourceID, dataPort, hardwareDescription, timeout);
 			assertTrue(invalidAddressFuture.get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS) instanceof RegistrationResponse.Decline);
 		} finally {
 			if (testingFatalErrorHandler.hasExceptionOccurred()) {

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/ResourceManagerTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/ResourceManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/ResourceManagerTest.java
index 08f5605..f5b6d3a 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/ResourceManagerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/ResourceManagerTest.java
@@ -30,7 +30,6 @@ import org.apache.flink.runtime.resourcemanager.slotmanager.SlotManager;
 import org.apache.flink.runtime.rest.messages.taskmanager.TaskManagerInfo;
 import org.apache.flink.runtime.rpc.RpcUtils;
 import org.apache.flink.runtime.rpc.TestingRpcService;
-import org.apache.flink.runtime.taskexecutor.SlotReport;
 import org.apache.flink.runtime.taskexecutor.TaskExecutorGateway;
 import org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGatewayBuilder;
 import org.apache.flink.runtime.testingUtils.TestingUtils;
@@ -120,7 +119,6 @@ public class ResourceManagerTest extends TestLogger {
 			CompletableFuture<RegistrationResponse> registrationResponseFuture = resourceManagerGateway.registerTaskExecutor(
 				taskExecutorGateway.getAddress(),
 				taskManagerId,
-				new SlotReport(),
 				dataPort,
 				hardwareDescription,
 				TestingUtils.TIMEOUT());

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/slotmanager/SlotManagerTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/slotmanager/SlotManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/slotmanager/SlotManagerTest.java
index 56ec12b..af6f3e4 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/slotmanager/SlotManagerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/slotmanager/SlotManagerTest.java
@@ -20,12 +20,14 @@ package org.apache.flink.runtime.resourcemanager.slotmanager;
 
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.api.common.time.Time;
+import org.apache.flink.api.java.tuple.Tuple5;
 import org.apache.flink.runtime.clusterframework.types.AllocationID;
 import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.runtime.clusterframework.types.ResourceProfile;
 import org.apache.flink.runtime.clusterframework.types.SlotID;
 import org.apache.flink.runtime.clusterframework.types.TaskManagerSlot;
 import org.apache.flink.runtime.concurrent.Executors;
+import org.apache.flink.runtime.concurrent.FutureUtils;
 import org.apache.flink.runtime.concurrent.ScheduledExecutor;
 import org.apache.flink.runtime.instance.InstanceID;
 import org.apache.flink.runtime.messages.Acknowledge;
@@ -40,12 +42,15 @@ import org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGateway;
 import org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGatewayBuilder;
 import org.apache.flink.runtime.taskexecutor.exceptions.SlotAllocationException;
 import org.apache.flink.runtime.testingUtils.TestingUtils;
+import org.apache.flink.util.FlinkException;
 import org.apache.flink.util.TestLogger;
 
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
 
 import java.util.Arrays;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.Executor;
@@ -73,6 +78,9 @@ import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+/**
+ * Tests for the {@link SlotManager}.
+ */
 public class SlotManagerTest extends TestLogger {
 
 	/**
@@ -1139,6 +1147,61 @@ public class SlotManagerTest extends TestLogger {
 		}
 	}
 
+	/**
+	 * Testst that the SlotManager retries allocating a slot if the TaskExecutor#requestSlot call
+	 * fails.
+	 */
+	@Test
+	public void testSlotRequestFailure() throws Exception {
+		try (final SlotManager slotManager = createSlotManager(ResourceManagerId.generate(), new TestingResourceActions())) {
+
+			final SlotRequest slotRequest = new SlotRequest(new JobID(), new AllocationID(), ResourceProfile.UNKNOWN, "foobar");
+			slotManager.registerSlotRequest(slotRequest);
+
+			final BlockingQueue<Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId>> requestSlotQueue = new ArrayBlockingQueue<>(1);
+			final BlockingQueue<CompletableFuture<Acknowledge>> responseQueue = new ArrayBlockingQueue<>(1);
+
+			final TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder()
+				.setRequestSlotFunction(slotIDJobIDAllocationIDStringResourceManagerIdTuple5 -> {
+					requestSlotQueue.offer(slotIDJobIDAllocationIDStringResourceManagerIdTuple5);
+					try {
+						return responseQueue.take();
+					} catch (InterruptedException ignored) {
+						return FutureUtils.completedExceptionally(new FlinkException("Response queue was interrupted."));
+					}
+				})
+				.createTestingTaskExecutorGateway();
+
+			final ResourceID taskExecutorResourceId = ResourceID.generate();
+			final TaskExecutorConnection taskExecutionConnection = new TaskExecutorConnection(taskExecutorResourceId, testingTaskExecutorGateway);
+			final SlotReport slotReport = new SlotReport(new SlotStatus(new SlotID(taskExecutorResourceId, 0), ResourceProfile.UNKNOWN));
+
+			final CompletableFuture<Acknowledge> firstManualSlotRequestResponse = new CompletableFuture<>();
+			responseQueue.offer(firstManualSlotRequestResponse);
+
+			slotManager.registerTaskManager(taskExecutionConnection, slotReport);
+
+			final Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId> firstRequest = requestSlotQueue.take();
+
+			final CompletableFuture<Acknowledge> secondManualSlotRequestResponse = new CompletableFuture<>();
+			responseQueue.offer(secondManualSlotRequestResponse);
+
+			// fail first request
+			firstManualSlotRequestResponse.completeExceptionally(new SlotAllocationException("Test exception"));
+
+			final Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId> secondRequest = requestSlotQueue.take();
+
+			assertThat(secondRequest.f2, equalTo(firstRequest.f2));
+			assertThat(secondRequest.f0, equalTo(firstRequest.f0));
+
+			secondManualSlotRequestResponse.complete(Acknowledge.get());
+
+			final TaskManagerSlot slot = slotManager.getSlot(secondRequest.f0);
+			assertThat(slot.getState(), equalTo(TaskManagerSlot.State.ALLOCATED));
+			assertThat(slot.getAllocationId(), equalTo(secondRequest.f2));
+		}
+	}
+
 	private SlotManager createSlotManager(ResourceManagerId resourceManagerId, ResourceActions resourceManagerActions) {
 		SlotManager slotManager = new SlotManager(
 			TestingUtils.defaultScheduledExecutor(),

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/slotmanager/TestingResourceActions.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/slotmanager/TestingResourceActions.java b/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/slotmanager/TestingResourceActions.java
new file mode 100644
index 0000000..915f142
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/slotmanager/TestingResourceActions.java
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.resourcemanager.slotmanager;
+
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.runtime.clusterframework.types.AllocationID;
+import org.apache.flink.runtime.clusterframework.types.ResourceProfile;
+import org.apache.flink.runtime.instance.InstanceID;
+
+/**
+ * Testing implementation of the {@link ResourceActions}.
+ */
+public class TestingResourceActions implements ResourceActions {
+	@Override
+	public void releaseResource(InstanceID instanceId, Exception cause) {
+
+	}
+
+	@Override
+	public void allocateResource(ResourceProfile resourceProfile) {
+
+	}
+
+	@Override
+	public void notifyAllocationFailure(JobID jobId, AllocationID allocationId, Exception cause) {
+
+	}
+}

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/utils/TestingResourceManagerGateway.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/utils/TestingResourceManagerGateway.java b/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/utils/TestingResourceManagerGateway.java
index 106fd73..daa2383 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/utils/TestingResourceManagerGateway.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/utils/TestingResourceManagerGateway.java
@@ -21,8 +21,8 @@ package org.apache.flink.runtime.resourcemanager.utils;
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.api.common.time.Time;
 import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.api.java.tuple.Tuple3;
 import org.apache.flink.api.java.tuple.Tuple4;
-import org.apache.flink.api.java.tuple.Tuple5;
 import org.apache.flink.runtime.blob.TransientBlobKey;
 import org.apache.flink.runtime.clusterframework.ApplicationStatus;
 import org.apache.flink.runtime.clusterframework.types.AllocationID;
@@ -50,6 +50,7 @@ import java.util.Collection;
 import java.util.Collections;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Function;
 
@@ -78,12 +79,18 @@ public class TestingResourceManagerGateway implements ResourceManagerGateway {
 
 	private volatile Consumer<Tuple2<JobID, Throwable>> disconnectJobManagerConsumer;
 
-	private volatile Function<Tuple5<String, ResourceID, SlotReport, Integer, HardwareDescription>, CompletableFuture<RegistrationResponse>> registerTaskExecutorFunction;
+	private volatile Function<Tuple4<String, ResourceID, Integer, HardwareDescription>, CompletableFuture<RegistrationResponse>> registerTaskExecutorFunction;
 
 	private volatile Function<Tuple2<ResourceID, FileType>, CompletableFuture<TransientBlobKey>> requestTaskManagerFileUploadFunction;
 
 	private volatile Consumer<Tuple2<ResourceID, Throwable>> disconnectTaskExecutorConsumer;
 
+	private volatile Function<Tuple3<ResourceID, InstanceID, SlotReport>, CompletableFuture<Acknowledge>> sendSlotReportFunction;
+
+	private volatile BiConsumer<ResourceID, SlotReport> taskExecutorHeartbeatConsumer;
+
+	private volatile Consumer<Tuple3<InstanceID, SlotID, AllocationID>> notifySlotAvailableConsumer;
+
 	public TestingResourceManagerGateway() {
 		this(
 			ResourceManagerId.generate(),
@@ -133,7 +140,7 @@ public class TestingResourceManagerGateway implements ResourceManagerGateway {
 		this.disconnectJobManagerConsumer = disconnectJobManagerConsumer;
 	}
 
-	public void setRegisterTaskExecutorFunction(Function<Tuple5<String, ResourceID, SlotReport, Integer, HardwareDescription>, CompletableFuture<RegistrationResponse>> registerTaskExecutorFunction) {
+	public void setRegisterTaskExecutorFunction(Function<Tuple4<String, ResourceID, Integer, HardwareDescription>, CompletableFuture<RegistrationResponse>> registerTaskExecutorFunction) {
 		this.registerTaskExecutorFunction = registerTaskExecutorFunction;
 	}
 
@@ -145,6 +152,18 @@ public class TestingResourceManagerGateway implements ResourceManagerGateway {
 		this.disconnectTaskExecutorConsumer = disconnectTaskExecutorConsumer;
 	}
 
+	public void setSendSlotReportFunction(Function<Tuple3<ResourceID, InstanceID, SlotReport>, CompletableFuture<Acknowledge>> sendSlotReportFunction) {
+		this.sendSlotReportFunction = sendSlotReportFunction;
+	}
+
+	public void setTaskExecutorHeartbeatConsumer(BiConsumer<ResourceID, SlotReport> taskExecutorHeartbeatConsumer) {
+		this.taskExecutorHeartbeatConsumer = taskExecutorHeartbeatConsumer;
+	}
+
+	public void setNotifySlotAvailableConsumer(Consumer<Tuple3<InstanceID, SlotID, AllocationID>> notifySlotAvailableConsumer) {
+		this.notifySlotAvailableConsumer = notifySlotAvailableConsumer;
+	}
+
 	@Override
 	public CompletableFuture<RegistrationResponse> registerJobManager(JobMasterId jobMasterId, ResourceID jobMasterResourceId, String jobMasterAddress, JobID jobId, Time timeout) {
 		final Consumer<Tuple4<JobMasterId, ResourceID, String, JobID>> currentConsumer = registerJobManagerConsumer;
@@ -187,11 +206,22 @@ public class TestingResourceManagerGateway implements ResourceManagerGateway {
 	}
 
 	@Override
-	public CompletableFuture<RegistrationResponse> registerTaskExecutor(String taskExecutorAddress, ResourceID resourceId, SlotReport slotReport, int dataPort, HardwareDescription hardwareDescription, Time timeout) {
-		final Function<Tuple5<String, ResourceID, SlotReport, Integer, HardwareDescription>, CompletableFuture<RegistrationResponse>> currentFunction = registerTaskExecutorFunction;
+	public CompletableFuture<Acknowledge> sendSlotReport(ResourceID taskManagerResourceId, InstanceID taskManagerRegistrationId, SlotReport slotReport, Time timeout) {
+		final Function<Tuple3<ResourceID, InstanceID, SlotReport>, CompletableFuture<Acknowledge>> currentSendSlotReportFunction = sendSlotReportFunction;
+
+		if (currentSendSlotReportFunction != null) {
+			return currentSendSlotReportFunction.apply(Tuple3.of(taskManagerResourceId, taskManagerRegistrationId, slotReport));
+		} else {
+			return CompletableFuture.completedFuture(Acknowledge.get());
+		}
+	}
+
+	@Override
+	public CompletableFuture<RegistrationResponse> registerTaskExecutor(String taskExecutorAddress, ResourceID resourceId, int dataPort, HardwareDescription hardwareDescription, Time timeout) {
+		final Function<Tuple4<String, ResourceID, Integer, HardwareDescription>, CompletableFuture<RegistrationResponse>> currentFunction = registerTaskExecutorFunction;
 
 		if (currentFunction != null) {
-			return currentFunction.apply(Tuple5.of(taskExecutorAddress, resourceId, slotReport, dataPort, hardwareDescription));
+			return currentFunction.apply(Tuple4.of(taskExecutorAddress, resourceId, dataPort, hardwareDescription));
 		} else {
 			return CompletableFuture.completedFuture(
 				new TaskExecutorRegistrationSuccess(
@@ -204,7 +234,11 @@ public class TestingResourceManagerGateway implements ResourceManagerGateway {
 
 	@Override
 	public void notifySlotAvailable(InstanceID instanceId, SlotID slotID, AllocationID oldAllocationId) {
+		final Consumer<Tuple3<InstanceID, SlotID, AllocationID>> currentNotifySlotAvailableConsumer = notifySlotAvailableConsumer;
 
+		if (currentNotifySlotAvailableConsumer != null) {
+			currentNotifySlotAvailableConsumer.accept(Tuple3.of(instanceId, slotID, oldAllocationId));
+		}
 	}
 
 	@Override
@@ -229,7 +263,11 @@ public class TestingResourceManagerGateway implements ResourceManagerGateway {
 
 	@Override
 	public void heartbeatFromTaskManager(ResourceID heartbeatOrigin, SlotReport slotReport) {
+		final BiConsumer<ResourceID, SlotReport> currentTaskExecutorHeartbeatConsumer = taskExecutorHeartbeatConsumer;
 
+		if (currentTaskExecutorHeartbeatConsumer != null) {
+			currentTaskExecutorHeartbeatConsumer.accept(heartbeatOrigin, slotReport);
+		}
 	}
 
 	@Override


[3/4] flink git commit: [FLINK-9427] Fix registration and request slot race condition in TaskExecutor

Posted by tr...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorTest.java
index ca31f76..9c1f2f2 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorTest.java
@@ -21,6 +21,8 @@ package org.apache.flink.runtime.taskexecutor;
 import org.apache.flink.api.common.ExecutionConfig;
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.api.common.time.Time;
+import org.apache.flink.api.java.tuple.Tuple3;
+import org.apache.flink.api.java.tuple.Tuple4;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.TaskManagerOptions;
 import org.apache.flink.core.testutils.OneShotLatch;
@@ -31,6 +33,7 @@ import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.runtime.clusterframework.types.ResourceProfile;
 import org.apache.flink.runtime.clusterframework.types.SlotID;
 import org.apache.flink.runtime.concurrent.Executors;
+import org.apache.flink.runtime.concurrent.FutureUtils;
 import org.apache.flink.runtime.concurrent.ScheduledExecutor;
 import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
 import org.apache.flink.runtime.deployment.ResultPartitionDeploymentDescriptor;
@@ -79,6 +82,7 @@ import org.apache.flink.runtime.rpc.RpcUtils;
 import org.apache.flink.runtime.rpc.TestingRpcService;
 import org.apache.flink.runtime.state.TaskExecutorLocalStateStoresManager;
 import org.apache.flink.runtime.taskexecutor.exceptions.RegistrationTimeoutException;
+import org.apache.flink.runtime.taskexecutor.exceptions.TaskManagerException;
 import org.apache.flink.runtime.taskexecutor.slot.SlotOffer;
 import org.apache.flink.runtime.taskexecutor.slot.TaskSlotTable;
 import org.apache.flink.runtime.taskexecutor.slot.TimerService;
@@ -108,6 +112,8 @@ import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
 import org.slf4j.Logger;
 
+import javax.annotation.Nonnull;
+
 import java.io.File;
 import java.io.IOException;
 import java.util.Arrays;
@@ -118,10 +124,13 @@ import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 
 import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.is;
@@ -132,6 +141,7 @@ import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 import static org.mockito.Mockito.RETURNS_MOCKS;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyInt;
@@ -143,6 +153,9 @@ import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+/**
+ * Tests for the {@link TaskExecutor}.
+ */
 public class TaskExecutorTest extends TestLogger {
 
 	@Rule
@@ -413,20 +426,32 @@ public class TaskExecutorTest extends TestLogger {
 		final long verificationTimeout = 1000L;
 		final long heartbeatTimeout = 10000L;
 		final String rmAddress = "rm";
-		final ResourceID rmResourceId = new ResourceID(rmAddress);
 		final UUID rmLeaderId = UUID.randomUUID();
 
 		// register the mock resource manager gateway
-		ResourceManagerGateway rmGateway = mock(ResourceManagerGateway.class);
-		when(rmGateway.registerTaskExecutor(
-			anyString(), any(ResourceID.class), any(SlotReport.class), anyInt(), any(HardwareDescription.class), any(Time.class)))
-			.thenReturn(
-				CompletableFuture.completedFuture(
-					new TaskExecutorRegistrationSuccess(
-						new InstanceID(),
-						rmResourceId,
-						10L,
-						new ClusterInformation("localhost", 1234))));
+		final TestingResourceManagerGateway rmGateway = new TestingResourceManagerGateway();
+		final CompletableFuture<ResourceID> taskExecutorRegistrationFuture = new CompletableFuture<>();
+		final ResourceID rmResourceId = rmGateway.getOwnResourceId();
+		final CompletableFuture<RegistrationResponse> registrationResponse = CompletableFuture.completedFuture(
+			new TaskExecutorRegistrationSuccess(
+				new InstanceID(),
+				rmResourceId,
+				10L,
+				new ClusterInformation("localhost", 1234)));
+
+		rmGateway.setRegisterTaskExecutorFunction(stringResourceIDIntegerHardwareDescriptionTuple4 -> {
+            taskExecutorRegistrationFuture.complete(stringResourceIDIntegerHardwareDescriptionTuple4.f1);
+            return registrationResponse;
+        });
+
+		final CompletableFuture<SlotReport> initialSlotReportFuture = new CompletableFuture<>();
+		rmGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
+            initialSlotReportFuture.complete(resourceIDInstanceIDSlotReportTuple3.f2);
+            return CompletableFuture.completedFuture(Acknowledge.get());
+        });
+
+		final CompletableFuture<SlotReport> heartbeatSlotReportFuture = new CompletableFuture<>();
+		rmGateway.setTaskExecutorHeartbeatConsumer((resourceID, slotReport) -> heartbeatSlotReportFuture.complete(slotReport));
 
 		rpc.registerGateway(rmAddress, rmGateway);
 
@@ -498,8 +523,8 @@ public class TaskExecutorTest extends TestLogger {
 			resourceManagerLeaderRetriever.notifyListener(rmAddress, rmLeaderId);
 
 			// register resource manager success will trigger monitoring heartbeat target between tm and rm
-			verify(rmGateway, timeout(verificationTimeout).atLeast(1)).registerTaskExecutor(
-				eq(taskManager.getAddress()), eq(taskManagerLocation.getResourceID()), eq(slotReport1), anyInt(), any(HardwareDescription.class), any(Time.class));
+			assertThat(taskExecutorRegistrationFuture.get(), equalTo(taskManagerLocation.getResourceID()));
+			assertThat(initialSlotReportFuture.get(), equalTo(slotReport1));
 
 			verify(heartbeatManager, timeout(verificationTimeout)).monitorTarget(any(ResourceID.class), any(HeartbeatTarget.class));
 
@@ -508,20 +533,13 @@ public class TaskExecutorTest extends TestLogger {
 			// trigger the heartbeat asynchronously
 			taskExecutorGateway.heartbeatFromResourceManager(rmResourceId);
 
-			ArgumentCaptor<SlotReport> slotReportArgumentCaptor = ArgumentCaptor.forClass(SlotReport.class);
-
 			// wait for heartbeat response
-			verify(rmGateway, timeout(verificationTimeout)).heartbeatFromTaskManager(
-				eq(taskManagerLocation.getResourceID()),
-				slotReportArgumentCaptor.capture());
-
-			SlotReport actualSlotReport = slotReportArgumentCaptor.getValue();
+			SlotReport actualSlotReport = heartbeatSlotReportFuture.get();
 
 			// the new slot report should be reported
 			assertEquals(slotReport2, actualSlotReport);
 		} finally {
-			taskManager.shutDown();
-			taskManager.getTerminationFuture().get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);
+			RpcUtils.terminateRpcEndpoint(taskManager, timeout);
 		}
 	}
 
@@ -536,7 +554,7 @@ public class TaskExecutorTest extends TestLogger {
 		// register a mock resource manager gateway
 		ResourceManagerGateway rmGateway = mock(ResourceManagerGateway.class);
 		when(rmGateway.registerTaskExecutor(
-					anyString(), any(ResourceID.class), any(SlotReport.class), anyInt(), any(HardwareDescription.class), any(Time.class)))
+					anyString(), any(ResourceID.class), anyInt(), any(HardwareDescription.class), any(Time.class)))
 			.thenReturn(CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(
 				new InstanceID(), resourceManagerResourceId, 10L, new ClusterInformation("localhost", 1234))));
 
@@ -578,7 +596,7 @@ public class TaskExecutorTest extends TestLogger {
 			String taskManagerAddress = taskManager.getAddress();
 
 			verify(rmGateway, Mockito.timeout(timeout.toMilliseconds())).registerTaskExecutor(
-					eq(taskManagerAddress), eq(taskManagerLocation.getResourceID()), eq(slotReport), anyInt(), any(HardwareDescription.class), any(Time.class));
+					eq(taskManagerAddress), eq(taskManagerLocation.getResourceID()), anyInt(), any(HardwareDescription.class), any(Time.class));
 		}
 		finally {
 			taskManager.shutDown();
@@ -600,11 +618,11 @@ public class TaskExecutorTest extends TestLogger {
 		ResourceManagerGateway rmGateway2 = mock(ResourceManagerGateway.class);
 
 		when(rmGateway1.registerTaskExecutor(
-					anyString(), any(ResourceID.class), any(SlotReport.class), anyInt(), any(HardwareDescription.class), any(Time.class)))
+					anyString(), any(ResourceID.class), anyInt(), any(HardwareDescription.class), any(Time.class)))
 			.thenReturn(CompletableFuture.completedFuture(
 				new TaskExecutorRegistrationSuccess(new InstanceID(), rmResourceId1, 10L, new ClusterInformation("localhost", 1234))));
 		when(rmGateway2.registerTaskExecutor(
-					anyString(), any(ResourceID.class), any(SlotReport.class), anyInt(), any(HardwareDescription.class), any(Time.class)))
+					anyString(), any(ResourceID.class), anyInt(), any(HardwareDescription.class), any(Time.class)))
 			.thenReturn(CompletableFuture.completedFuture(
 				new TaskExecutorRegistrationSuccess(new InstanceID(), rmResourceId2, 10L, new ClusterInformation("localhost", 1234))));
 
@@ -647,7 +665,7 @@ public class TaskExecutorTest extends TestLogger {
 			resourceManagerLeaderRetriever.notifyListener(address1, leaderId1);
 
 			verify(rmGateway1, Mockito.timeout(timeout.toMilliseconds())).registerTaskExecutor(
-					eq(taskManagerAddress), eq(taskManagerLocation.getResourceID()), any(SlotReport.class), anyInt(), any(HardwareDescription.class), any(Time.class));
+					eq(taskManagerAddress), eq(taskManagerLocation.getResourceID()), anyInt(), any(HardwareDescription.class), any(Time.class));
 			assertNotNull(taskManager.getResourceManagerConnection());
 
 			// cancel the leader 
@@ -657,7 +675,7 @@ public class TaskExecutorTest extends TestLogger {
 			resourceManagerLeaderRetriever.notifyListener(address2, leaderId2);
 
 			verify(rmGateway2, Mockito.timeout(timeout.toMilliseconds())).registerTaskExecutor(
-					eq(taskManagerAddress), eq(taskManagerLocation.getResourceID()), eq(slotReport), anyInt(), any(HardwareDescription.class), any(Time.class));
+					eq(taskManagerAddress), eq(taskManagerLocation.getResourceID()), anyInt(), any(HardwareDescription.class), any(Time.class));
 			assertNotNull(taskManager.getResourceManagerConnection());
 		}
 		finally {
@@ -792,7 +810,7 @@ public class TaskExecutorTest extends TestLogger {
 	}
 
 	/**
-	 * Tests that a TaskManager detects a job leader for which has reserved slots. Upon detecting
+	 * Tests that a TaskManager detects a job leader for which it has reserved slots. Upon detecting
 	 * the job leader, it will offer all reserved slots to the JobManager.
 	 */
 	@Test
@@ -803,18 +821,14 @@ public class TaskExecutorTest extends TestLogger {
 
 		final String resourceManagerAddress = "rm";
 		final ResourceManagerId resourceManagerLeaderId = ResourceManagerId.generate();
-		final ResourceID resourceManagerResourceId = new ResourceID(resourceManagerAddress);
 
-		final ResourceManagerGateway resourceManagerGateway = mock(ResourceManagerGateway.class);
-		final InstanceID registrationId = new InstanceID();
+		final TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
+		CompletableFuture<Void> initialSlotReportFuture = new CompletableFuture<>();
+		resourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
+            initialSlotReportFuture.complete(null);
+            return CompletableFuture.completedFuture(Acknowledge.get());
 
-		when(resourceManagerGateway.registerTaskExecutor(
-			any(String.class),
-			eq(taskManagerLocation.getResourceID()),
-			any(SlotReport.class),
-			anyInt(),
-			any(HardwareDescription.class),
-			any(Time.class))).thenReturn(CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(registrationId, resourceManagerResourceId, 1000L, new ClusterInformation("localhost", 1234))));
+        });
 
 		final String jobManagerAddress = "jm";
 		final UUID jobManagerLeaderId = UUID.randomUUID();
@@ -870,6 +884,9 @@ public class TaskExecutorTest extends TestLogger {
 			// tell the task manager about the rm leader
 			resourceManagerLeaderRetriever.notifyListener(resourceManagerAddress, resourceManagerLeaderId.toUUID());
 
+			// wait for the initial slot report
+			initialSlotReportFuture.get();
+
 			// request slots from the task manager under the given allocation id
 			CompletableFuture<Acknowledge> slotRequestAck = tmGateway.requestSlot(
 				slotId,
@@ -890,8 +907,7 @@ public class TaskExecutorTest extends TestLogger {
 					(Collection<SlotOffer>)Matchers.argThat(contains(slotOffer)),
 					any(Time.class));
 		} finally {
-			taskManager.shutDown();
-			taskManager.getTerminationFuture().get(timeout.toMilliseconds(), TimeUnit.MILLISECONDS);
+			RpcUtils.terminateRpcEndpoint(taskManager, timeout);
 		}
 	}
 
@@ -907,7 +923,6 @@ public class TaskExecutorTest extends TestLogger {
 
 		final String resourceManagerAddress = "rm";
 		final UUID resourceManagerLeaderId = UUID.randomUUID();
-		final ResourceID resourceManagerResourceId = new ResourceID(resourceManagerAddress);
 
 		final String jobManagerAddress = "jm";
 		final UUID jobManagerLeaderId = UUID.randomUUID();
@@ -915,16 +930,20 @@ public class TaskExecutorTest extends TestLogger {
 		resourceManagerLeaderRetriever.notifyListener(resourceManagerAddress, resourceManagerLeaderId);
 		jobManagerLeaderRetriever.notifyListener(jobManagerAddress, jobManagerLeaderId);
 
-		final ResourceManagerGateway resourceManagerGateway = mock(ResourceManagerGateway.class);
+		final TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
+		final ResourceID resourceManagerResourceId = resourceManagerGateway.getOwnResourceId();
 		final InstanceID registrationId = new InstanceID();
 
-		when(resourceManagerGateway.registerTaskExecutor(
-			any(String.class),
-			eq(taskManagerLocation.getResourceID()),
-			any(SlotReport.class),
-			anyInt(),
-			any(HardwareDescription.class),
-			any(Time.class))).thenReturn(CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(registrationId, resourceManagerResourceId, 1000L, new ClusterInformation("localhost", 1234))));
+		final CompletableFuture<ResourceID> registrationFuture = new CompletableFuture<>();
+		resourceManagerGateway.setRegisterTaskExecutorFunction(
+			stringResourceIDIntegerHardwareDescriptionTuple4 -> {
+                registrationFuture.complete(stringResourceIDIntegerHardwareDescriptionTuple4.f1);
+                return CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(registrationId, resourceManagerResourceId, 1000L, new ClusterInformation("localhost", 1234)));
+            }
+		);
+
+		final CompletableFuture<Tuple3<InstanceID, SlotID, AllocationID>> availableSlotFuture = new CompletableFuture<>();
+		resourceManagerGateway.setNotifySlotAvailableConsumer(availableSlotFuture::complete);
 
 		final ResourceID jmResourceId = new ResourceID(jobManagerAddress);
 
@@ -975,14 +994,7 @@ public class TaskExecutorTest extends TestLogger {
 		try {
 			taskManager.start();
 
-			// wait for the registration at the ResourceManager
-			verify(resourceManagerGateway, Mockito.timeout(timeout.toMilliseconds())).registerTaskExecutor(
-				eq(taskManager.getAddress()),
-				eq(taskManagerLocation.getResourceID()),
-				any(SlotReport.class),
-				anyInt(),
-				any(HardwareDescription.class),
-				any(Time.class));
+			assertThat(registrationFuture.get(), equalTo(taskManagerLocation.getResourceID()));
 
 			taskSlotTable.allocateSlot(0, jobId, allocationId1, Time.milliseconds(10000L));
 			taskSlotTable.allocateSlot(1, jobId, allocationId2, Time.milliseconds(10000L));
@@ -991,10 +1003,11 @@ public class TaskExecutorTest extends TestLogger {
 			// been properly started.
 			jobLeaderService.addJob(jobId, jobManagerAddress);
 
-			verify(resourceManagerGateway, Mockito.timeout(timeout.toMilliseconds())).notifySlotAvailable(
-				eq(registrationId),
-				eq(new SlotID(taskManagerLocation.getResourceID(), 1)),
-				eq(allocationId2));
+			final Tuple3<InstanceID, SlotID, AllocationID> instanceIDSlotIDAllocationIDTuple3 = availableSlotFuture.get();
+
+			final Tuple3<InstanceID, SlotID, AllocationID> expectedResult = Tuple3.of(registrationId, new SlotID(taskManagerLocation.getResourceID(), 1), allocationId2);
+
+			assertThat(instanceIDSlotIDAllocationIDTuple3, equalTo(expectedResult));
 
 			assertTrue(taskSlotTable.existsActiveSlot(jobId, allocationId1));
 			assertFalse(taskSlotTable.existsActiveSlot(jobId, allocationId2));
@@ -1016,7 +1029,6 @@ public class TaskExecutorTest extends TestLogger {
 
 		final String resourceManagerAddress = "rm";
 		final UUID resourceManagerLeaderId = UUID.randomUUID();
-		final ResourceID resourceManagerResourceId = new ResourceID(resourceManagerAddress);
 
 		final String jobManagerAddress = "jm";
 		final JobMasterId jobMasterId = JobMasterId.generate();
@@ -1024,17 +1036,10 @@ public class TaskExecutorTest extends TestLogger {
 		resourceManagerLeaderRetriever.notifyListener(resourceManagerAddress, resourceManagerLeaderId);
 		jobManagerLeaderRetriever.notifyListener(jobManagerAddress, jobMasterId.toUUID());
 
-		final ResourceManagerGateway resourceManagerGateway = mock(ResourceManagerGateway.class);
-		final InstanceID registrationId = new InstanceID();
+		final TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
 
-		when(resourceManagerGateway.registerTaskExecutor(
-			any(String.class),
-			eq(taskManagerLocation.getResourceID()),
-			any(SlotReport.class),
-			anyInt(),
-			any(HardwareDescription.class),
-			any(Time.class))).thenReturn(
-				CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(registrationId, resourceManagerResourceId, 1000L, new ClusterInformation("localhost", 1234))));
+		final CompletableFuture<Tuple3<InstanceID, SlotID, AllocationID>> availableSlotFuture = new CompletableFuture<>();
+		resourceManagerGateway.setNotifySlotAvailableConsumer(availableSlotFuture::complete);
 
 		final ResourceID jmResourceId = new ResourceID(jobManagerAddress);
 
@@ -1169,10 +1174,8 @@ public class TaskExecutorTest extends TestLogger {
 			// acknowledge the offered slots
 			offerResultFuture.complete(Collections.singleton(offer1));
 
-			verify(resourceManagerGateway, Mockito.timeout(timeout.toMilliseconds())).notifySlotAvailable(
-				eq(registrationId),
-				eq(new SlotID(taskManagerLocation.getResourceID(), 1)),
-				any(AllocationID.class));
+			final Tuple3<InstanceID, SlotID, AllocationID> instanceIDSlotIDAllocationIDTuple3 = availableSlotFuture.get();
+			assertThat(instanceIDSlotIDAllocationIDTuple3.f1, equalTo(new SlotID(taskManagerLocation.getResourceID(), 1)));
 
 			assertTrue(taskSlotTable.existsActiveSlot(jobId, allocationId1));
 			assertFalse(taskSlotTable.existsActiveSlot(jobId, allocationId2));
@@ -1353,6 +1356,11 @@ public class TaskExecutorTest extends TestLogger {
 
 		try {
 			final TestingResourceManagerGateway resourceManagerGateway = new TestingResourceManagerGateway();
+			final CompletableFuture<Void> initialSlotReport = new CompletableFuture<>();
+			resourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
+                initialSlotReport.complete(null);
+                return CompletableFuture.completedFuture(Acknowledge.get());
+            });
 			final ResourceManagerId resourceManagerId = resourceManagerGateway.getFencingToken();
 
 			rpc.registerGateway(resourceManagerGateway.getAddress(), resourceManagerGateway);
@@ -1377,6 +1385,9 @@ public class TaskExecutorTest extends TestLogger {
 			final JobLeaderService jobLeaderService = taskManagerServices.getJobLeaderService();
 			assertThat(jobLeaderService.containsJob(jobId), is(false));
 
+			// wait for the initial slot report
+			initialSlotReport.get();
+
 			taskExecutorGateway.requestSlot(
 				slotId,
 				jobId,
@@ -1403,15 +1414,7 @@ public class TaskExecutorTest extends TestLogger {
 	public void testMaximumRegistrationDuration() throws Exception {
 		configuration.setString(TaskManagerOptions.REGISTRATION_TIMEOUT, "10 ms");
 
-		final TaskExecutor taskExecutor = new TaskExecutor(
-			rpc,
-			TaskManagerConfiguration.fromConfiguration(configuration),
-			haServices,
-			new TaskManagerServicesBuilder().build(),
-			new HeartbeatServices(1000L, 1000L),
-			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
-			dummyBlobCacheService,
-			testingFatalErrorHandler);
+		final TaskExecutor taskExecutor = createTaskExecutor(new TaskManagerServicesBuilder().build());
 
 		taskExecutor.start();
 
@@ -1483,6 +1486,216 @@ public class TaskExecutorTest extends TestLogger {
 		}
 	}
 
+	/**
+	 * Tests that we ignore slot requests if the TaskExecutor is not
+	 * registered at a ResourceManager.
+	 */
+	@Test
+	public void testIgnoringSlotRequestsIfNotRegistered() throws Exception {
+		final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
+		final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).build();
+
+		final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);
+
+		taskExecutor.start();
+
+		try {
+			final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
+
+			final CompletableFuture<RegistrationResponse> registrationFuture = new CompletableFuture<>();
+			final CompletableFuture<ResourceID> taskExecutorResourceIdFuture = new CompletableFuture<>();
+
+			testingResourceManagerGateway.setRegisterTaskExecutorFunction(stringResourceIDSlotReportIntegerHardwareDescriptionTuple5 -> {
+                taskExecutorResourceIdFuture.complete(stringResourceIDSlotReportIntegerHardwareDescriptionTuple5.f1);
+                return registrationFuture;
+            });
+
+			rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
+			resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());
+
+			final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);
+
+			final ResourceID resourceId = taskExecutorResourceIdFuture.get();
+
+			final SlotID slotId = new SlotID(resourceId, 0);
+			final CompletableFuture<Acknowledge> slotRequestResponse = taskExecutorGateway.requestSlot(slotId, jobId, new AllocationID(), "foobar", testingResourceManagerGateway.getFencingToken(), timeout);
+
+			try {
+				slotRequestResponse.get();
+				fail("We should not be able to request slots before the TaskExecutor is registered at the ResourceManager.");
+			} catch (ExecutionException ee) {
+				assertThat(ExceptionUtils.stripExecutionException(ee), instanceOf(TaskManagerException.class));
+			}
+		} finally {
+			RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
+		}
+	}
+
+	/**
+	 * Tests that the TaskExecutor tries to reconnect to a ResourceManager from which it
+	 * was explicitly disconnected.
+	 */
+	@Test
+	public void testReconnectionAttemptIfExplicitlyDisconnected() throws Exception {
+		final long heartbeatInterval = 1000L;
+		final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
+		final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
+		final TaskExecutor taskExecutor = new TaskExecutor(
+			rpc,
+			TaskManagerConfiguration.fromConfiguration(configuration),
+			haServices,
+			new TaskManagerServicesBuilder()
+				.setTaskSlotTable(taskSlotTable)
+				.setTaskManagerLocation(taskManagerLocation)
+				.build(),
+			new HeartbeatServices(heartbeatInterval, 1000L),
+			UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
+			dummyBlobCacheService,
+			testingFatalErrorHandler);
+
+		taskExecutor.start();
+
+		try {
+			final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
+			final ClusterInformation clusterInformation = new ClusterInformation("foobar", 1234);
+			final CompletableFuture<RegistrationResponse> registrationResponseFuture = CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(new InstanceID(), ResourceID.generate(), heartbeatInterval, clusterInformation));
+			final BlockingQueue<ResourceID> registrationQueue = new ArrayBlockingQueue<>(1);
+
+			testingResourceManagerGateway.setRegisterTaskExecutorFunction(stringResourceIDSlotReportIntegerHardwareDescriptionTuple5 -> {
+				registrationQueue.offer(stringResourceIDSlotReportIntegerHardwareDescriptionTuple5.f1);
+				return registrationResponseFuture;
+			});
+			rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
+
+			resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());
+
+			final ResourceID firstRegistrationAttempt = registrationQueue.take();
+
+			assertThat(firstRegistrationAttempt, equalTo(taskManagerLocation.getResourceID()));
+
+			final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);
+
+			assertThat(registrationQueue, is(empty()));
+
+			taskExecutorGateway.disconnectResourceManager(new FlinkException("Test exception"));
+
+			final ResourceID secondRegistrationAttempt = registrationQueue.take();
+
+			assertThat(secondRegistrationAttempt, equalTo(taskManagerLocation.getResourceID()));
+
+		} finally {
+			RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
+		}
+	}
+
+	/**
+	 * Tests that the {@link TaskExecutor} sends the initial slot report after it
+	 * registered at the ResourceManager.
+	 */
+	@Test
+	public void testInitialSlotReport() throws Exception {
+		final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
+		final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
+		final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
+			.setTaskSlotTable(taskSlotTable)
+			.setTaskManagerLocation(taskManagerLocation)
+			.build();
+		final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);
+
+		taskExecutor.start();
+
+		try {
+			final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
+			final CompletableFuture<ResourceID> initialSlotReportFuture = new CompletableFuture<>();
+
+			testingResourceManagerGateway.setSendSlotReportFunction(
+				resourceIDInstanceIDSlotReportTuple3 -> {
+                    initialSlotReportFuture.complete(resourceIDInstanceIDSlotReportTuple3.f0);
+                    return CompletableFuture.completedFuture(Acknowledge.get());
+                }
+			);
+
+			rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
+			resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());
+
+			assertThat(initialSlotReportFuture.get(), equalTo(taskManagerLocation.getResourceID()));
+		} finally {
+			RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
+		}
+	}
+
+	/**
+	 * Tests that the {@link TaskExecutor} tries to reconnect if the initial slot report
+	 * fails.
+	 */
+	@Test
+	public void testInitialSlotReportFailure() throws Exception {
+		final TaskSlotTable taskSlotTable = new TaskSlotTable(Collections.singleton(ResourceProfile.UNKNOWN), timerService);
+		final TaskManagerLocation taskManagerLocation = new LocalTaskManagerLocation();
+		final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder()
+			.setTaskSlotTable(taskSlotTable)
+			.setTaskManagerLocation(taskManagerLocation)
+			.build();
+		final TaskExecutor taskExecutor = createTaskExecutor(taskManagerServices);
+
+		taskExecutor.start();
+
+		try {
+			final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
+
+			final BlockingQueue<CompletableFuture<Acknowledge>> responseQueue = new ArrayBlockingQueue<>(2);
+			testingResourceManagerGateway.setSendSlotReportFunction(
+				resourceIDInstanceIDSlotReportTuple3 -> {
+					try {
+						return responseQueue.take();
+					} catch (InterruptedException e) {
+						return FutureUtils.completedExceptionally(e);
+					}
+				});
+
+			final CompletableFuture<RegistrationResponse> registrationResponse = CompletableFuture.completedFuture(
+				new TaskExecutorRegistrationSuccess(
+					new InstanceID(),
+					testingResourceManagerGateway.getOwnResourceId(),
+					1000L,
+					new ClusterInformation("foobar", 1234)	));
+
+			final CountDownLatch numberRegistrations = new CountDownLatch(2);
+
+			testingResourceManagerGateway.setRegisterTaskExecutorFunction(new Function<Tuple4<String, ResourceID, Integer, HardwareDescription>, CompletableFuture<RegistrationResponse>>() {
+				@Override
+				public CompletableFuture<RegistrationResponse> apply(Tuple4<String, ResourceID, Integer, HardwareDescription> stringResourceIDIntegerHardwareDescriptionTuple4) {
+					numberRegistrations.countDown();
+					return registrationResponse;
+				}
+			});
+
+			responseQueue.offer(FutureUtils.completedExceptionally(new FlinkException("Test exception")));
+			responseQueue.offer(CompletableFuture.completedFuture(Acknowledge.get()));
+
+			rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
+			resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());
+
+			//wait for the second registration attempt
+			numberRegistrations.await();
+		} finally {
+			RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
+		}
+	}
+
+	@Nonnull
+	private TaskExecutor createTaskExecutor(TaskManagerServices taskManagerServices) {
+		return new TaskExecutor(
+                rpc,
+                TaskManagerConfiguration.fromConfiguration(configuration),
+                haServices,
+                taskManagerServices,
+                new HeartbeatServices(1000L, 1000L),
+                UnregisteredMetricGroups.createUnregisteredTaskManagerMetricGroup(),
+                dummyBlobCacheService,
+                testingFatalErrorHandler);
+	}
+
 	private static final class StartStopNotifyingLeaderRetrievalService implements LeaderRetrievalService {
 		private final CompletableFuture<LeaderRetrievalListener> startFuture;
 

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TestingTaskExecutorGateway.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TestingTaskExecutorGateway.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TestingTaskExecutorGateway.java
index fde64a0..a9e9949 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TestingTaskExecutorGateway.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TestingTaskExecutorGateway.java
@@ -20,6 +20,7 @@ package org.apache.flink.runtime.taskexecutor;
 
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.api.common.time.Time;
+import org.apache.flink.api.java.tuple.Tuple5;
 import org.apache.flink.runtime.blob.TransientBlobKey;
 import org.apache.flink.runtime.checkpoint.CheckpointOptions;
 import org.apache.flink.runtime.clusterframework.types.AllocationID;
@@ -35,11 +36,11 @@ import org.apache.flink.runtime.messages.StackTraceSampleResponse;
 import org.apache.flink.runtime.resourcemanager.ResourceManagerId;
 import org.apache.flink.util.Preconditions;
 
-
 import java.util.concurrent.CompletableFuture;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 /**
  * Simple {@link TaskExecutorGateway} implementation for testing purposes.
@@ -56,17 +57,20 @@ public class TestingTaskExecutorGateway implements TaskExecutorGateway {
 
 	private final BiFunction<TaskDeploymentDescriptor, JobMasterId, CompletableFuture<Acknowledge>> submitTaskConsumer;
 
-	public TestingTaskExecutorGateway(String address, String hostname, Consumer<ResourceID> heartbeatJobManagerConsumer, BiConsumer<JobID, Throwable> disconnectJobManagerConsumer, BiFunction<TaskDeploymentDescriptor, JobMasterId, CompletableFuture<Acknowledge>> submitTaskConsumer) {
+	private final Function<Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId>, CompletableFuture<Acknowledge>> requestSlotFunction;
+
+	TestingTaskExecutorGateway(String address, String hostname, Consumer<ResourceID> heartbeatJobManagerConsumer, BiConsumer<JobID, Throwable> disconnectJobManagerConsumer, BiFunction<TaskDeploymentDescriptor, JobMasterId, CompletableFuture<Acknowledge>> submitTaskConsumer, Function<Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId>, CompletableFuture<Acknowledge>> requestSlotFunction) {
 		this.address = Preconditions.checkNotNull(address);
 		this.hostname = Preconditions.checkNotNull(hostname);
 		this.heartbeatJobManagerConsumer = Preconditions.checkNotNull(heartbeatJobManagerConsumer);
 		this.disconnectJobManagerConsumer = Preconditions.checkNotNull(disconnectJobManagerConsumer);
 		this.submitTaskConsumer = Preconditions.checkNotNull(submitTaskConsumer);
+		this.requestSlotFunction = Preconditions.checkNotNull(requestSlotFunction);
 	}
 
 	@Override
 	public CompletableFuture<Acknowledge> requestSlot(SlotID slotId, JobID jobId, AllocationID allocationId, String targetAddress, ResourceManagerId resourceManagerId, Time timeout) {
-		return CompletableFuture.completedFuture(Acknowledge.get());
+		return requestSlotFunction.apply(Tuple5.of(slotId, jobId, allocationId, targetAddress, resourceManagerId));
 	}
 
 	@Override

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TestingTaskExecutorGatewayBuilder.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TestingTaskExecutorGatewayBuilder.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TestingTaskExecutorGatewayBuilder.java
index 2550abf..1c2f132 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TestingTaskExecutorGatewayBuilder.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TestingTaskExecutorGatewayBuilder.java
@@ -19,15 +19,20 @@
 package org.apache.flink.runtime.taskexecutor;
 
 import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.java.tuple.Tuple5;
+import org.apache.flink.runtime.clusterframework.types.AllocationID;
 import org.apache.flink.runtime.clusterframework.types.ResourceID;
+import org.apache.flink.runtime.clusterframework.types.SlotID;
 import org.apache.flink.runtime.deployment.TaskDeploymentDescriptor;
 import org.apache.flink.runtime.jobmaster.JobMasterId;
 import org.apache.flink.runtime.messages.Acknowledge;
+import org.apache.flink.runtime.resourcemanager.ResourceManagerId;
 
 import java.util.concurrent.CompletableFuture;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 /**
  * Builder for a {@link TestingTaskExecutorGateway}.
@@ -37,12 +42,14 @@ public class TestingTaskExecutorGatewayBuilder {
 	private static final Consumer<ResourceID> NOOP_HEARTBEAT_JOBMANAGER_CONSUMER = ignored -> {};
 	private static final BiConsumer<JobID, Throwable> NOOP_DISCONNECT_JOBMANAGER_CONSUMER = (ignoredA, ignoredB) -> {};
 	private static final BiFunction<TaskDeploymentDescriptor, JobMasterId, CompletableFuture<Acknowledge>> NOOP_SUBMIT_TASK_CONSUMER = (ignoredA, ignoredB) -> CompletableFuture.completedFuture(Acknowledge.get());
+	private static final Function<Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId>, CompletableFuture<Acknowledge>> NOOP_REQUEST_SLOT_FUNCTION = ignored -> CompletableFuture.completedFuture(Acknowledge.get());
 
 	private String address = "foobar:1234";
 	private String hostname = "foobar";
 	private Consumer<ResourceID> heartbeatJobManagerConsumer = NOOP_HEARTBEAT_JOBMANAGER_CONSUMER;
 	private BiConsumer<JobID, Throwable> disconnectJobManagerConsumer = NOOP_DISCONNECT_JOBMANAGER_CONSUMER;
 	private BiFunction<TaskDeploymentDescriptor, JobMasterId, CompletableFuture<Acknowledge>> submitTaskConsumer = NOOP_SUBMIT_TASK_CONSUMER;
+	private Function<Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId>, CompletableFuture<Acknowledge>> requestSlotFunction = NOOP_REQUEST_SLOT_FUNCTION;
 
 	public TestingTaskExecutorGatewayBuilder setAddress(String address) {
 		this.address = address;
@@ -69,7 +76,12 @@ public class TestingTaskExecutorGatewayBuilder {
 		return this;
 	}
 
+	public TestingTaskExecutorGatewayBuilder setRequestSlotFunction(Function<Tuple5<SlotID, JobID, AllocationID, String, ResourceManagerId>, CompletableFuture<Acknowledge>> requestSlotFunction) {
+		this.requestSlotFunction = requestSlotFunction;
+		return this;
+	}
+
 	public TestingTaskExecutorGateway createTestingTaskExecutorGateway() {
-		return new TestingTaskExecutorGateway(address, hostname, heartbeatJobManagerConsumer, disconnectJobManagerConsumer, submitTaskConsumer);
+		return new TestingTaskExecutorGateway(address, hostname, heartbeatJobManagerConsumer, disconnectJobManagerConsumer, submitTaskConsumer, requestSlotFunction);
 	}
 }

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-yarn-tests/src/test/java/org/apache/flink/yarn/YarnConfigurationITCase.java
----------------------------------------------------------------------
diff --git a/flink-yarn-tests/src/test/java/org/apache/flink/yarn/YarnConfigurationITCase.java b/flink-yarn-tests/src/test/java/org/apache/flink/yarn/YarnConfigurationITCase.java
index a5f7c10..f770135 100644
--- a/flink-yarn-tests/src/test/java/org/apache/flink/yarn/YarnConfigurationITCase.java
+++ b/flink-yarn-tests/src/test/java/org/apache/flink/yarn/YarnConfigurationITCase.java
@@ -159,10 +159,11 @@ public class YarnConfigurationITCase extends YarnTestBase {
 
 					taskManagerInfos = taskManagersInfo.getTaskManagerInfos();
 
-					if (taskManagerInfos.isEmpty()) {
-						Thread.sleep(100L);
-					} else {
+					// wait until the task manager has registered and reported its slots
+					if (hasTaskManagerConnectedAndReportedSlots(taskManagerInfos)) {
 						break;
+					} else {
+						Thread.sleep(100L);
 					}
 				}
 
@@ -190,4 +191,13 @@ public class YarnConfigurationITCase extends YarnTestBase {
 			clusterDescriptor.close();
 		}
 	}
+
+	private boolean hasTaskManagerConnectedAndReportedSlots(Collection<TaskManagerInfo> taskManagerInfos) {
+		if (taskManagerInfos.isEmpty()) {
+			return false;
+		} else {
+			final TaskManagerInfo taskManagerInfo = taskManagerInfos.iterator().next();
+			return taskManagerInfo.getNumberSlots() > 0;
+		}
+	}
 }

http://git-wip-us.apache.org/repos/asf/flink/blob/47dc6997/flink-yarn/src/test/java/org/apache/flink/yarn/YarnResourceManagerTest.java
----------------------------------------------------------------------
diff --git a/flink-yarn/src/test/java/org/apache/flink/yarn/YarnResourceManagerTest.java b/flink-yarn/src/test/java/org/apache/flink/yarn/YarnResourceManagerTest.java
index 03f3ef2..4fffc2b 100644
--- a/flink-yarn/src/test/java/org/apache/flink/yarn/YarnResourceManagerTest.java
+++ b/flink-yarn/src/test/java/org/apache/flink/yarn/YarnResourceManagerTest.java
@@ -36,6 +36,7 @@ import org.apache.flink.runtime.highavailability.HighAvailabilityServices;
 import org.apache.flink.runtime.highavailability.TestingHighAvailabilityServices;
 import org.apache.flink.runtime.instance.HardwareDescription;
 import org.apache.flink.runtime.leaderelection.TestingLeaderElectionService;
+import org.apache.flink.runtime.messages.Acknowledge;
 import org.apache.flink.runtime.metrics.MetricRegistry;
 import org.apache.flink.runtime.metrics.NoOpMetricRegistry;
 import org.apache.flink.runtime.registration.RegistrationResponse;
@@ -50,6 +51,7 @@ import org.apache.flink.runtime.rpc.TestingRpcService;
 import org.apache.flink.runtime.taskexecutor.SlotReport;
 import org.apache.flink.runtime.taskexecutor.SlotStatus;
 import org.apache.flink.runtime.taskexecutor.TaskExecutorGateway;
+import org.apache.flink.runtime.taskexecutor.TaskExecutorRegistrationSuccess;
 import org.apache.flink.runtime.testutils.DirectScheduledExecutorService;
 import org.apache.flink.runtime.util.TestingFatalErrorHandler;
 import org.apache.flink.util.TestLogger;
@@ -93,8 +95,10 @@ import static org.apache.flink.yarn.YarnConfigKeys.ENV_FLINK_CLASSPATH;
 import static org.apache.flink.yarn.YarnConfigKeys.ENV_HADOOP_USER_NAME;
 import static org.apache.flink.yarn.YarnConfigKeys.FLINK_JAR_PATH;
 import static org.apache.flink.yarn.YarnConfigKeys.FLINK_YARN_FILES;
+import static org.hamcrest.Matchers.instanceOf;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.eq;
@@ -362,12 +366,21 @@ public class YarnResourceManagerTest extends TestLogger {
 				.registerTaskExecutor(
 					taskHost,
 					taskManagerResourceId,
-					slotReport,
 					dataPort,
 					hardwareDescription,
 					Time.seconds(10L))
+				.thenCompose(
+					(RegistrationResponse response) -> {
+						assertThat(response, instanceOf(TaskExecutorRegistrationSuccess.class));
+						final TaskExecutorRegistrationSuccess success = (TaskExecutorRegistrationSuccess) response;
+						return rmGateway.sendSlotReport(
+							taskManagerResourceId,
+							success.getRegistrationId(),
+							slotReport,
+							Time.seconds(10L));
+					})
 				.handleAsync(
-					(RegistrationResponse response, Throwable throwable) -> rmServices.slotManager.getNumberRegisteredSlots(),
+					(Acknowledge ignored, Throwable throwable) -> rmServices.slotManager.getNumberRegisteredSlots(),
 					resourceManager.getMainThreadExecutorForTesting());
 
 			final int numberRegisteredSlots = numberRegisteredSlotsFuture.get();


[2/4] flink git commit: [FLINK-9421] Remove job from RunningJobsRegistry when it reaches a terminal state

Posted by tr...@apache.org.
[FLINK-9421] Remove job from RunningJobsRegistry when it reaches a terminal state

This commit lets the Dispatcher remove the RunningJobsRegistry entry for a completed job
when it is removed from the Dispatcher.

This closes #6068.


Project: http://git-wip-us.apache.org/repos/asf/flink/repo
Commit: http://git-wip-us.apache.org/repos/asf/flink/commit/5ad0ca23
Tree: http://git-wip-us.apache.org/repos/asf/flink/tree/5ad0ca23
Diff: http://git-wip-us.apache.org/repos/asf/flink/diff/5ad0ca23

Branch: refs/heads/master
Commit: 5ad0ca2392a9672a92756337243834d1b466a24d
Parents: 47dc699
Author: Till Rohrmann <tr...@apache.org>
Authored: Wed May 23 23:48:38 2018 +0200
Committer: Till Rohrmann <tr...@apache.org>
Committed: Thu May 24 16:45:15 2018 +0200

----------------------------------------------------------------------
 .../flink/runtime/dispatcher/Dispatcher.java    |  6 ++
 .../highavailability/RunningJobsRegistry.java   | 24 +++---
 .../StandaloneRunningJobsRegistry.java          |  4 +-
 .../zookeeper/ZooKeeperRunningJobsRegistry.java |  7 +-
 .../DispatcherResourceCleanupTest.java          | 88 +++++++++++++++++++-
 .../TestingHighAvailabilityServices.java        |  5 +-
 6 files changed, 115 insertions(+), 19 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/flink/blob/5ad0ca23/flink-runtime/src/main/java/org/apache/flink/runtime/dispatcher/Dispatcher.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/dispatcher/Dispatcher.java b/flink-runtime/src/main/java/org/apache/flink/runtime/dispatcher/Dispatcher.java
index 5022388..f5fdd27 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/dispatcher/Dispatcher.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/dispatcher/Dispatcher.java
@@ -563,6 +563,12 @@ public abstract class Dispatcher extends FencedRpcEndpoint<DispatcherId> impleme
 					} catch (Exception e) {
 						log.warn("Could not properly remove job {} from submitted job graph store.", jobId);
 					}
+
+					try {
+						runningJobsRegistry.clearJob(jobId);
+					} catch (IOException e) {
+						log.warn("Could not properly remove job {} from the running jobs registry.", jobId);
+					}
 				}
 			},
 			getRpcService().getExecutor());

http://git-wip-us.apache.org/repos/asf/flink/blob/5ad0ca23/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/RunningJobsRegistry.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/RunningJobsRegistry.java b/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/RunningJobsRegistry.java
index 43e5ac5..2e24f27 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/RunningJobsRegistry.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/RunningJobsRegistry.java
@@ -24,16 +24,16 @@ import java.io.IOException;
 
 /**
  * A simple registry that tracks if a certain job is pending execution, running, or completed.
- * 
+ *
  * <p>This registry is used in highly-available setups with multiple master nodes,
- * to determine whether a new leader should attempt to recover a certain job (because the 
+ * to determine whether a new leader should attempt to recover a certain job (because the
  * job is still running), or whether the job has already finished successfully (in case of a
  * finite job) and the leader has only been granted leadership because the previous leader
  * quit cleanly after the job was finished.
- * 
+ *
  * <p>In addition, the registry can help to determine whether a newly assigned leader JobManager
  * should attempt reconciliation with running TaskManagers, or immediately schedule the job from
- * the latest checkpoint/savepoint. 
+ * the latest checkpoint/savepoint.
  */
 public interface RunningJobsRegistry {
 
@@ -42,13 +42,13 @@ public interface RunningJobsRegistry {
 	 */
 	enum JobSchedulingStatus {
 
-		/** Job has not been scheduled, yet */
+		/** Job has not been scheduled, yet. */
 		PENDING,
 
-		/** Job has been scheduled and is not yet finished */
+		/** Job has been scheduled and is not yet finished. */
 		RUNNING,
 
-		/** Job has been finished, successfully or unsuccessfully */
+		/** Job has been finished, successfully or unsuccessfully. */
 		DONE;
 	}
 
@@ -57,7 +57,7 @@ public interface RunningJobsRegistry {
 	/**
 	 * Marks a job as running. Requesting the job's status via the {@link #getJobSchedulingStatus(JobID)}
 	 * method will return {@link JobSchedulingStatus#RUNNING}.
-	 * 
+	 *
 	 * @param jobID The id of the job.
 	 *
 	 * @throws IOException Thrown when the communication with the highly-available storage or registry
@@ -70,7 +70,7 @@ public interface RunningJobsRegistry {
 	 * method will return {@link JobSchedulingStatus#DONE}.
 	 *
 	 * @param jobID The id of the job.
-	 * 
+	 *
 	 * @throws IOException Thrown when the communication with the highly-available storage or registry
 	 *                     failed and could not be retried.
 	 */
@@ -81,17 +81,17 @@ public interface RunningJobsRegistry {
 	 *
 	 * @param jobID The id of the job to check.
 	 * @return The job scheduling status.
-	 * 
+	 *
 	 * @throws IOException Thrown when the communication with the highly-available storage or registry
 	 *                     failed and could not be retried.
 	 */
 	JobSchedulingStatus getJobSchedulingStatus(JobID jobID) throws IOException;
 
 	/**
-	 * Clear job state form the registry, usually called after job finish
+	 * Clear job state form the registry, usually called after job finish.
 	 *
 	 * @param jobID The id of the job to check.
-	 * 
+	 *
 	 * @throws IOException Thrown when the communication with the highly-available storage or registry
 	 *                     failed and could not be retried.
 	 */

http://git-wip-us.apache.org/repos/asf/flink/blob/5ad0ca23/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/nonha/standalone/StandaloneRunningJobsRegistry.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/nonha/standalone/StandaloneRunningJobsRegistry.java b/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/nonha/standalone/StandaloneRunningJobsRegistry.java
index 585ef34..d289cc6 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/nonha/standalone/StandaloneRunningJobsRegistry.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/nonha/standalone/StandaloneRunningJobsRegistry.java
@@ -30,7 +30,7 @@ import static org.apache.flink.util.Preconditions.checkNotNull;
  */
 public class StandaloneRunningJobsRegistry implements RunningJobsRegistry {
 
-	/** The currently running jobs */
+	/** The currently running jobs. */
 	private final HashMap<JobID, JobSchedulingStatus> jobStatus = new HashMap<>();
 
 	@Override
@@ -54,7 +54,7 @@ public class StandaloneRunningJobsRegistry implements RunningJobsRegistry {
 	@Override
 	public JobSchedulingStatus getJobSchedulingStatus(JobID jobID) {
 		checkNotNull(jobID);
-		
+
 		synchronized (jobStatus) {
 			JobSchedulingStatus status = jobStatus.get(jobID);
 			return status == null ? JobSchedulingStatus.PENDING : status;

http://git-wip-us.apache.org/repos/asf/flink/blob/5ad0ca23/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/zookeeper/ZooKeeperRunningJobsRegistry.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/zookeeper/ZooKeeperRunningJobsRegistry.java b/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/zookeeper/ZooKeeperRunningJobsRegistry.java
index 8a083d1..1a84b37 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/zookeeper/ZooKeeperRunningJobsRegistry.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/highavailability/zookeeper/ZooKeeperRunningJobsRegistry.java
@@ -18,11 +18,12 @@
 
 package org.apache.flink.runtime.highavailability.zookeeper;
 
-import org.apache.curator.framework.CuratorFramework;
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.HighAvailabilityOptions;
 import org.apache.flink.runtime.highavailability.RunningJobsRegistry;
+
+import org.apache.curator.framework.CuratorFramework;
 import org.apache.zookeeper.data.Stat;
 
 import java.io.IOException;
@@ -38,7 +39,7 @@ public class ZooKeeperRunningJobsRegistry implements RunningJobsRegistry {
 
 	private static final Charset ENCODING = Charset.forName("utf-8");
 
-	/** The ZooKeeper client to use */
+	/** The ZooKeeper client to use. */
 	private final CuratorFramework client;
 
 	private final String runningJobPath;
@@ -88,7 +89,7 @@ public class ZooKeeperRunningJobsRegistry implements RunningJobsRegistry {
 						return JobSchedulingStatus.valueOf(name);
 					}
 					catch (IllegalArgumentException e) {
-						throw new IOException("Found corrupt data in ZooKeeper: " + 
+						throw new IOException("Found corrupt data in ZooKeeper: " +
 								Arrays.toString(data) + " is no valid job status");
 					}
 				}

http://git-wip-us.apache.org/repos/asf/flink/blob/5ad0ca23/flink-runtime/src/test/java/org/apache/flink/runtime/dispatcher/DispatcherResourceCleanupTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/dispatcher/DispatcherResourceCleanupTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/dispatcher/DispatcherResourceCleanupTest.java
index 77b3fc4..f768251 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/dispatcher/DispatcherResourceCleanupTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/dispatcher/DispatcherResourceCleanupTest.java
@@ -22,6 +22,7 @@ import org.apache.flink.api.common.JobID;
 import org.apache.flink.api.common.time.Time;
 import org.apache.flink.configuration.BlobServerOptions;
 import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.testutils.OneShotLatch;
 import org.apache.flink.runtime.blob.BlobKey;
 import org.apache.flink.runtime.blob.BlobServer;
 import org.apache.flink.runtime.blob.BlobStore;
@@ -32,6 +33,7 @@ import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.runtime.executiongraph.ArchivedExecutionGraph;
 import org.apache.flink.runtime.heartbeat.HeartbeatServices;
 import org.apache.flink.runtime.highavailability.HighAvailabilityServices;
+import org.apache.flink.runtime.highavailability.RunningJobsRegistry;
 import org.apache.flink.runtime.highavailability.TestingHighAvailabilityServices;
 import org.apache.flink.runtime.jobgraph.JobGraph;
 import org.apache.flink.runtime.jobgraph.JobStatus;
@@ -54,6 +56,7 @@ import org.apache.flink.runtime.rpc.TestingRpcService;
 import org.apache.flink.runtime.testtasks.NoOpInvokable;
 import org.apache.flink.runtime.testutils.InMemorySubmittedJobGraphStore;
 import org.apache.flink.runtime.util.TestingFatalErrorHandler;
+import org.apache.flink.util.Preconditions;
 import org.apache.flink.util.TestLogger;
 
 import org.junit.After;
@@ -64,6 +67,7 @@ import org.junit.ClassRule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
 
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
 import java.io.File;
@@ -101,6 +105,12 @@ public class DispatcherResourceCleanupTest extends TestLogger {
 
 	private TestingLeaderElectionService dispatcherLeaderElectionService;
 
+	private SingleRunningJobsRegistry runningJobsRegistry;
+
+	private TestingHighAvailabilityServices highAvailabilityServices;
+
+	private OneShotLatch clearedJobLatch;
+
 	private TestingDispatcher dispatcher;
 
 	private DispatcherGateway dispatcherGateway;
@@ -134,9 +144,12 @@ public class DispatcherResourceCleanupTest extends TestLogger {
 		configuration = new Configuration();
 		configuration.setString(BlobServerOptions.STORAGE_DIRECTORY, temporaryFolder.newFolder().getAbsolutePath());
 
-		final TestingHighAvailabilityServices highAvailabilityServices = new TestingHighAvailabilityServices();
+		highAvailabilityServices = new TestingHighAvailabilityServices();
 		dispatcherLeaderElectionService = new TestingLeaderElectionService();
 		highAvailabilityServices.setDispatcherLeaderElectionService(dispatcherLeaderElectionService);
+		clearedJobLatch = new OneShotLatch();
+		runningJobsRegistry = new SingleRunningJobsRegistry(jobId, clearedJobLatch);
+		highAvailabilityServices.setRunningJobsRegistry(runningJobsRegistry);
 
 		storedBlobFuture = new CompletableFuture<>();
 		deleteAllFuture = new CompletableFuture<>();
@@ -271,6 +284,79 @@ public class DispatcherResourceCleanupTest extends TestLogger {
 		assertThat(deleteAllFuture.isDone(), is(false));
 	}
 
+	/**
+	 * Tests that the {@link RunningJobsRegistry} entries are cleared after the
+	 * job reached a terminal state.
+	 */
+	@Test
+	public void testRunningJobsRegistryCleanup() throws Exception {
+		submitJob();
+
+		runningJobsRegistry.setJobRunning(jobId);
+		assertThat(runningJobsRegistry.contains(jobId), is(true));
+
+		resultFuture.complete(new ArchivedExecutionGraphBuilder().setState(JobStatus.FINISHED).setJobID(jobId).build());
+
+		// wait for the clearing
+		clearedJobLatch.await();
+
+		assertThat(runningJobsRegistry.contains(jobId), is(false));
+	}
+
+	private static final class SingleRunningJobsRegistry implements RunningJobsRegistry {
+
+		@Nonnull
+		private final JobID expectedJobId;
+
+		@Nonnull
+		private final OneShotLatch clearedJobLatch;
+
+		private JobSchedulingStatus jobSchedulingStatus = JobSchedulingStatus.PENDING;
+
+		private boolean containsJob = false;
+
+		private SingleRunningJobsRegistry(@Nonnull JobID expectedJobId, @Nonnull OneShotLatch clearedJobLatch) {
+			this.expectedJobId = expectedJobId;
+			this.clearedJobLatch = clearedJobLatch;
+		}
+
+		@Override
+		public void setJobRunning(JobID jobID) {
+			checkJobId(jobID);
+			containsJob = true;
+			jobSchedulingStatus = JobSchedulingStatus.RUNNING;
+		}
+
+		private void checkJobId(JobID jobID) {
+			Preconditions.checkArgument(expectedJobId.equals(jobID));
+		}
+
+		@Override
+		public void setJobFinished(JobID jobID) {
+			checkJobId(jobID);
+			containsJob = true;
+			jobSchedulingStatus = JobSchedulingStatus.DONE;
+		}
+
+		@Override
+		public JobSchedulingStatus getJobSchedulingStatus(JobID jobID) {
+			checkJobId(jobID);
+			return jobSchedulingStatus;
+		}
+
+		public boolean contains(JobID jobId) {
+			checkJobId(jobId);
+			return containsJob;
+		}
+
+		@Override
+		public void clearJob(JobID jobID) {
+			checkJobId(jobID);
+			containsJob = false;
+			clearedJobLatch.trigger();
+		}
+	}
+
 	private static final class TestingDispatcher extends Dispatcher {
 		public TestingDispatcher(RpcService rpcService, String endpointId, Configuration configuration, HighAvailabilityServices highAvailabilityServices, SubmittedJobGraphStore submittedJobGraphStore, ResourceManagerGateway resourceManagerGateway, BlobServer blobServer, HeartbeatServices heartbeatServices, JobManagerMetricGroup jobManagerMetricGroup, @Nullable String metricServiceQueryPath, ArchivedExecutionGraphStore archivedExecutionGraphStore, JobManagerRunnerFactory jobManagerRunnerFactory, FatalErrorHandler fatalErrorHandler, @Nullable String restAddress, HistoryServerArchivist historyServerArchivist) throws Exception {
 			super(rpcService, endpointId, configuration, highAvailabilityServices, submittedJobGraphStore, resourceManagerGateway, blobServer, heartbeatServices, jobManagerMetricGroup, metricServiceQueryPath, archivedExecutionGraphStore, jobManagerRunnerFactory, fatalErrorHandler, restAddress, historyServerArchivist);

http://git-wip-us.apache.org/repos/asf/flink/blob/5ad0ca23/flink-runtime/src/test/java/org/apache/flink/runtime/highavailability/TestingHighAvailabilityServices.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/highavailability/TestingHighAvailabilityServices.java b/flink-runtime/src/test/java/org/apache/flink/runtime/highavailability/TestingHighAvailabilityServices.java
index 98d5d74..2489f16 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/highavailability/TestingHighAvailabilityServices.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/highavailability/TestingHighAvailabilityServices.java
@@ -56,7 +56,7 @@ public class TestingHighAvailabilityServices implements HighAvailabilityServices
 
 	private volatile SubmittedJobGraphStore submittedJobGraphStore;
 
-	private final RunningJobsRegistry runningJobsRegistry = new StandaloneRunningJobsRegistry();
+	private volatile RunningJobsRegistry runningJobsRegistry = new StandaloneRunningJobsRegistry();
 
 	// ------------------------------------------------------------------------
 	//  Setters for mock / testing implementations
@@ -102,6 +102,9 @@ public class TestingHighAvailabilityServices implements HighAvailabilityServices
 		this.submittedJobGraphStore = submittedJobGraphStore;
 	}
 
+	public void setRunningJobsRegistry(RunningJobsRegistry runningJobsRegistry) {
+		this.runningJobsRegistry = runningJobsRegistry;
+	}
 	// ------------------------------------------------------------------------
 	//  HA Services Methods
 	// ------------------------------------------------------------------------